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 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 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-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 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​

On this page