Query with the Google Gemini API

Important

The Google Gemini API is only compatible with Gemini pay per token foundation models and external models. For a unified API that works across all providers, use the Chat Completions API.

The Google Gemini API provides native Google AI SDK compatibility for Gemini models on Azure Databricks. Use this API when you need Gemini-specific features or are migrating existing Google AI SDK code.

Requirements

  • See Requirements.
  • Install the google-genai package on your compute.

Query examples

The following examples show how to query a Foundation Model API pay-per-token endpoint using the Google Gemini API.

Python

from google import genai
from google.genai import types
import os

DATABRICKS_TOKEN = os.environ.get('DATABRICKS_TOKEN')

client = genai.Client(
    api_key="databricks",
    http_options=types.HttpOptions(
        base_url="https://example.staging.cloud.databricks.com/serving-endpoints/gemini",
        headers={
            "Authorization": f"Bearer {DATABRICKS_TOKEN}",
        },
    ),
)

response = client.models.generate_content(
    model="databricks-gemini-2-5-pro",
    contents=[
        types.Content(
            role="user",
            parts=[types.Part(text="What is a mixture of experts model?")],
        ),
    ],
    config=types.GenerateContentConfig(
        max_output_tokens=256,
    ),
)

print(response.text)

REST API

curl \
-u token:$DATABRICKS_TOKEN \
-X POST \
-H "Content-Type: application/json" \
-d '{
  "contents": [
    {
      "role": "user",
      "parts": [{"text": "What is a mixture of experts model?"}]
    }
  ],
  "generationConfig": {
    "maxOutputTokens": 256
  }
}' \
https://<workspace_host>.databricks.com/serving-endpoints/gemini/v1beta/models/databricks-gemini-2-5-pro:generateContent

Supported models

Databricks-hosted foundation models

  • databricks-gemini-3-1-pro
  • databricks-gemini-3-1-flash-lite
  • databricks-gemini-3-pro
  • databricks-gemini-3-flash
  • databricks-gemini-2-5-pro
  • databricks-gemini-2-5-flash

External models

  • Google model provider

Supported input types

Gemini models on Azure Databricks accept text and image inputs. The Pro variants (databricks-gemini-3-1-pro, databricks-gemini-3-pro, databricks-gemini-2-5-pro) and the Flash variants (databricks-gemini-3-1-flash-lite, databricks-gemini-3-flash, databricks-gemini-2-5-flash) also accept video and audio inputs. See Query audio and video models for video and audio request examples, and Query vision models for image format and size requirements. For per-model input types, see Databricks-hosted foundation models available in Foundation Model APIs.

Additional resources

  • Last updated on