Skip to main content
Last updated on

Query with the Google Gemini API

important

The Google Gemini API is only compatible with Gemini pay per token 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 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)

Supported models​

Databricks-hosted foundation models​

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

Additional resources​