Query model services with the Google Gemini API

Beta

This feature is in Beta. Account admins can control access to this feature from the account console Previews page. See Manage Databricks previews.

important

The Google Gemini API is only compatible with Gemini foundation 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 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 model service using the Google Gemini API.

Python

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://<workspace-url>/ai-gateway/gemini",
        headers={
            "Authorization": f"Bearer {DATABRICKS_TOKEN}",
        },
    ),
)

response = client.models.generate_content(
    model="system.ai.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

Bash

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-url>/ai-gateway/gemini/v1beta/models/system.ai.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

Supported input types

Gemini models on 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 foundation models by type for video and audio request examples, and Query foundation models by type for image format and size requirements. For per-model input types, see Databricks-hosted foundation models available in Foundation Model APIs.

Additional resources

• Query foundation models by type.
• Databricks-hosted foundation models available in Foundation Model APIs.