Skip to content

Gemini API Integration Guide

Complete guide for using async-batch-llm with Google's Gemini API.

Installation

# Install async-batch-llm with Gemini support
pip install 'async-batch-llm[gemini]'
# or
uv add 'async-batch-llm[gemini]'

This installs:

  • async-batch-llm - Core batch processing framework
  • google-genai - Official Google Gemini SDK
  • pydantic - For response validation

Setup

1. Get API Key

Get a free API key from Google AI Studio: https://aistudio.google.com/apikey

2. Set Environment Variable

export GOOGLE_API_KEY=your_api_key_here

Or in Python:

import os
os.environ["GOOGLE_API_KEY"] = "your_api_key_here"

Tip: google-genai still honors the older GEMINI_API_KEY variable for backward compatibility, but GOOGLE_API_KEY takes precedence when both are set. Set at least one.

3. Verify Setup

from google import genai

client = genai.Client()
response = client.models.generate_content(
    model="gemini-2.5-flash",
    contents="Say hello!"
)
print(response.text)

Usage with async-batch-llm

async-batch-llm provides one built-in Gemini strategy (GeminiStrategy) plus two models — GeminiModel (direct) and GeminiCachedModel (context caching). You pick caching by choosing the model; the strategy is the same either way.

1. GeminiModel (Simple API Calls)

For direct Gemini API calls without caching:

from async_batch_llm import (
    GeminiModel,
    GeminiStrategy,
    LLMWorkItem,
    ParallelBatchProcessor,
    ProcessorConfig,
)
from google import genai
from pydantic import BaseModel

class SummaryOutput(BaseModel):
    """Structured output for summarization."""
    summary: str
    key_points: list[str]

# Create client
client = genai.Client(api_key="your-api-key")

# Create response parser
def parse_response(response) -> SummaryOutput:
    """Parse the normalized LLMResponse into your output model."""
    return SummaryOutput.model_validate_json(response.text)

# Create the model, then wrap it in the strategy. Pass `generation_config` to
# thread Gemini's per-call generation config — here, server-enforced JSON via
# `response_schema` / `response_mime_type` — on every call, no custom strategy
# needed. (For a config that must change per retry attempt, subclass
# `ModelStrategy.execute()` and read `self.generation_config`.)
model = GeminiModel("gemini-2.5-flash", client)
strategy = GeminiStrategy(
    model,
    response_parser=parse_response,
    temperature=0.7,
    generation_config={
        "response_mime_type": "application/json",
        "response_schema": SummaryOutput,
    },
)

# Configure processor
config = ProcessorConfig(max_workers=5, timeout_per_item=30.0)

# Process items
async with ParallelBatchProcessor[str, SummaryOutput, None](config=config) as processor:
    texts = ["Text 1...", "Text 2...", "Text 3..."]

    for i, text in enumerate(texts):
        await processor.add_work(
            LLMWorkItem(
                item_id=f"text_{i}",
                strategy=strategy,
                prompt=f"Summarize: {text}",
            )
        )

    result = await processor.process_all()

# Use results
for item in result.results:
    if item.success:
        print(f"{item.item_id}: {item.output.summary}")
        print(f"  Tokens: {item.token_usage['total_tokens']}")

2. GeminiCachedModel (With Context Caching)

Perfect for RAG applications with large shared context. Since v0.6.0, caching lives on a dedicated model — GeminiCachedModel — which you wrap in the ordinary GeminiStrategy. (The old GeminiCachedStrategy was removed; the model now owns the cache find/create/renew/delete lifecycle.)

from async_batch_llm import GeminiCachedModel, GeminiStrategy
from google import genai

client = genai.Client(api_key="your-api-key")

# Define large context to cache (e.g., retrieved documents)
cached_content = [
    genai.types.Content(
        role="user",
        parts=[
            genai.types.Part(text="Large document or knowledge base to cache...")
        ]
    )
]

# Create ONE cached model and share it across every work item that should
# share this context — constructing a new instance per item defeats caching
# entirely and can cost ~10x more.
cached_model = GeminiCachedModel(
    "gemini-2.5-flash",
    client,
    cached_content=cached_content,
    cache_ttl_seconds=3600,             # Cache for 1 hour
    cache_renewal_buffer_seconds=300,   # Renew 5 min before expiry
    auto_renew=True,
)
strategy = GeminiStrategy(cached_model, response_parser=lambda r: r.text)

# Use in processor
config = ProcessorConfig(max_workers=3, timeout_per_item=30.0)

async with ParallelBatchProcessor[str, str, None](config=config) as processor:
    questions = [
        "What is the main topic?",
        "What are the key findings?",
        "What are the conclusions?",
    ]

    for i, question in enumerate(questions):
        await processor.add_work(
            LLMWorkItem(
                item_id=f"question_{i}",
                strategy=strategy,
                prompt=question,
            )
        )

    result = await processor.process_all()

# Cache lifecycle (driven by the framework via GeminiCachedModel):
# - prepare() finds or creates the cache once per shared instance
# - generate() renews automatically when nearing expiry (auto_renew)
# - cleanup() leaves the cache active for reuse across runs
#   (call cached_model.delete_cache() to force removal)

Advanced Features

Progressive Temperature on Retries

Create a custom strategy that adjusts temperature based on attempt:

from async_batch_llm.llm_strategies import LLMCallStrategy

class ProgressiveTempGeminiStrategy(LLMCallStrategy[SummaryOutput]):
    """Gemini strategy with progressive temperature."""

    def __init__(self, client: genai.Client, temps=[0.0, 0.5, 1.0]):
        self.client = client
        self.temps = temps

    async def execute(
        self, prompt: str, attempt: int, timeout: float, state=None
    ) -> tuple[SummaryOutput, dict[str, int]]:
        # Use higher temperature for retries
        temp = self.temps[min(attempt - 1, len(self.temps) - 1)]

        config = genai.types.GenerateContentConfig(
            temperature=temp,
            response_mime_type="application/json",
            response_schema=SummaryOutput,
        )

        response = await self.client.aio.models.generate_content(
            model="gemini-2.5-flash",
            contents=prompt,
            config=config,
        )

        output = SummaryOutput.model_validate_json(response.text)

        usage = response.usage_metadata
        tokens = {
            "input_tokens": usage.prompt_token_count or 0,
            "output_tokens": usage.candidates_token_count or 0,
            "total_tokens": usage.total_token_count or 0,
        }

        return output, tokens

# Use it
strategy = ProgressiveTempGeminiStrategy(client=client, temps=[0.0, 0.5, 1.0])

Why progressive temperature?

  • Attempt 1 (temp=0.0): Deterministic, most likely to succeed
  • Attempt 2 (temp=0.5): More creative if first attempt had validation errors
  • Attempt 3 (temp=1.0): Maximum creativity as last resort

Model Selection

# Fast, experimental (free tier)
model="gemini-2.0-flash-exp"

# Production-ready, fast
model="gemini-2.5-flash-lite"

# Most capable, slower
model="gemini-2.5-flash"

# With extended thinking
model="gemini-2.5-pro"

See: https://ai.google.dev/gemini-api/docs/models/gemini

Generation Config Options

from google.genai.types import GenerateContentConfig

config = GenerateContentConfig(
    # Temperature: 0.0 (deterministic) to 1.0 (creative)
    temperature=0.7,

    # Nucleus sampling: Consider tokens with cumulative probability top_p
    top_p=0.95,

    # Top-k sampling: Consider only top k tokens
    top_k=40,

    # Maximum tokens in response
    max_output_tokens=2048,

    # Stop sequences
    stop_sequences=["END", "STOP"],

    # Structured output
    response_mime_type="application/json",
    response_schema=YourPydanticModel,

    # System instruction
    system_instruction="You are a helpful assistant...",

    # Safety settings
    safety_settings=[
        {
            "category": "HARM_CATEGORY_HATE_SPEECH",
            "threshold": "BLOCK_MEDIUM_AND_ABOVE"
        }
    ],
)

To apply these on every call without writing a custom strategy, pass the equivalent dict as GeminiStrategy(..., generation_config={...}) — it's merged into the SDK config on each generate(). (safety_settings also has a dedicated GeminiModel(safety_settings=...) argument.)

See: https://ai.google.dev/gemini-api/docs/models/generative-models#model-parameters

Grounding & custom metadata extraction

Request grounding (the google_search tool) via generation_config, and the built-in Gemini models surface the citations by default under metadata['grounding'] — no extractor needed. Read them through the typed views on the result (see the typed auxiliary output section in the API reference; experimental — the shapes may change while they stabilize):

from async_batch_llm import GeminiModel, GeminiStrategy
from google.genai import types

model = GeminiModel("gemini-2.5-flash", client)
strategy = GeminiStrategy(
    model,
    generation_config={"tools": [types.Tool(google_search=types.GoogleSearch())]},
)

# After processing, grounding is available provider-agnostically on each
# per-item WorkItemResult (the views live on WorkItemResult/LLMResponse,
# not on the BatchResult):
for item in result.successes:
    for source in item.grounding.sources:  # typed view over metadata["grounding"]
        print(source.uri, source.title)
    print(item.grounding.queries)          # the web_search_queries the model ran
    # item.grounding.supports              -> answer-span -> source-index links
    # Or read the plain dicts directly: item.metadata["grounding"]["sources"], ...

Grounding data only exists when you requested the tool, so the default metadata payload is unchanged for non-grounded calls — mixing grounded and non-grounded calls through the same model is fine. (grounding_metadata_extractor is still exported for custom models and explicit opt-in configurations; passing it to the built-in models is redundant but harmless.)

To surface other provider-specific output through the same channel without subclassing the model, pass metadata extractors: hooks that receive the raw provider response and return extra metadata keys. Write your own for anything (it's any Callable[[Any], dict | None]):

def reasoning_extractor(response):
    """Surface a thinking/reasoning trace under metadata['reasoning']."""
    candidate = (response.candidates or [None])[0]
    thought = getattr(candidate, "thought", None) if candidate else None
    return {"reasoning": thought} if thought else None

model = GeminiModel("gemini-2.5-flash", client, metadata_extractors=[reasoning_extractor])

Extractors merge on top of the built-in metadata (your keys win on collision), run independently, and a failing extractor is logged and skipped rather than breaking the call. The same metadata_extractors= constructor argument is available on GeminiCachedModel, OpenAIModel, OpenRouterModel, and DeepSeekModel; the OpenAI-compatible models (OpenAIModel, OpenRouterModel, DeepSeekModel) also accept it through their from_api_key(...) constructor.

Error Handling

async-batch-llm includes GeminiErrorClassifier for Gemini-specific errors:

from async_batch_llm.classifiers import GeminiErrorClassifier

processor = ParallelBatchProcessor(
    config=config,
    error_classifier=GeminiErrorClassifier(),  # Handles 429, 500, etc.
)

The classifier automatically:

  • Detects rate limit errors (429) as retryable and triggers the coordinated cooldown (other 4xx client errors fail fast as non-retryable)
  • Marks server errors (500) as retryable
  • Detects timeout errors
  • Handles validation errors

Rate Limit Handling

Gemini has rate limits. Configure automatic handling:

from async_batch_llm import RateLimitConfig

config = ProcessorConfig(
    max_workers=10,
    timeout_per_item=30.0,
    rate_limit=RateLimitConfig(
        cooldown_seconds=60.0,  # Wait 60s after rate limit
        slow_start_items=50,  # Gradually resume over 50 items
        slow_start_initial_delay=2.0,  # 2s between items initially
        slow_start_final_delay=0.1,  # 0.1s between items finally
    )
)

When rate limit (429) is detected:

  1. All workers pause
  2. Wait for cooldown period
  3. Resume with slow-start (gradual ramp-up)
  4. Automatically retry failed items

Multimodal Inputs

Process images with text:

from google.genai.types import Part, Content
from async_batch_llm.llm_strategies import LLMCallStrategy

class GeminiVisionStrategy(LLMCallStrategy[str]):
    """Strategy for Gemini vision tasks."""

    def __init__(self, client: genai.Client, image_path: str):
        self.client = client
        self.image_path = image_path

    async def execute(
        self, prompt: str, attempt: int, timeout: float, state=None
    ) -> tuple[str, dict[str, int]]:
        # Read image
        with open(self.image_path, "rb") as f:
            image_bytes = f.read()

        # Create multimodal content
        contents = [
            Content(
                parts=[
                    Part.from_bytes(data=image_bytes, mime_type="image/jpeg"),
                    Part.from_text(text=prompt)
                ]
            )
        ]

        response = await self.client.aio.models.generate_content(
            model="gemini-2.5-flash",
            contents=contents,
        )

        usage = response.usage_metadata
        tokens = {
            "input_tokens": usage.prompt_token_count or 0,
            "output_tokens": usage.candidates_token_count or 0,
            "total_tokens": usage.total_token_count or 0,
        }

        return response.text, tokens

# Use it
strategy = GeminiVisionStrategy(client=client, image_path="photo.jpg")

See: https://ai.google.dev/gemini-api/docs/vision

Token Usage Tracking

# After processing
result = await processor.process_all()

print(f"Total input tokens: {result.total_input_tokens}")
print(f"Total output tokens: {result.total_output_tokens}")
print(f"Total tokens: {result.total_input_tokens + result.total_output_tokens}")

# Per-item usage
for item in result.results:
    if item.success:
        tokens = item.token_usage
        cost = tokens["input_tokens"] * 0.00001 + tokens["output_tokens"] * 0.00003
        print(f"{item.item_id}: ${cost:.6f}")

Pricing: https://ai.google.dev/pricing

Complete Example

See examples/example_gemini_direct.py for a complete working example.

Run it:

export GOOGLE_API_KEY=your_key_here
uv run python examples/example_gemini_direct.py

Comparison: PydanticAI vs Direct Strategies

With PydanticAI

from async_batch_llm import PydanticAIStrategy
from pydantic_ai import Agent

agent = Agent('gemini-2.5-flash', output_type=SummaryOutput)
strategy = PydanticAIStrategy(agent=agent)

work_item = LLMWorkItem(
    item_id="item_1",
    strategy=strategy,
    prompt="Summarize this text..."
)

Pros: Simpler API, less code Cons: Less control over API parameters, extra dependency

Direct Gemini Strategy

from async_batch_llm import GeminiModel, GeminiStrategy

model = GeminiModel("gemini-2.5-flash", client)
strategy = GeminiStrategy(model, response_parser=parse_response, temperature=0.7)

work_item = LLMWorkItem(
    item_id="item_1",
    strategy=strategy,
    prompt="Summarize this text..."
)

Pros: Full control, no PydanticAI dependency, custom configurations Cons: More code, manual parsing

Best Practices

  1. Use structured output: Set response_schema for reliable parsing
  2. Implement progressive temperature: Start low (0.0), increase on retries
  3. Set reasonable timeouts: 30s for simple, 120s for complex queries
  4. Handle errors gracefully: Use GeminiErrorClassifier for Gemini errors
  5. Monitor token usage: Track costs using item.token_usage
  6. Respect rate limits: Configure rate_limit settings appropriately
  7. Choose right model: Use flash for speed, pro for quality
  8. Use caching for RAG: GeminiCachedModel saves money on repeated context

Troubleshooting

API Key Not Found

Error: GOOGLE_API_KEY environment variable not set

Fix: Export your API key before running:

export GOOGLE_API_KEY=your_key_here

Rate Limit Errors (429)

google.genai.errors.ClientError: 429 Resource exhausted

Fix: Reduce max_workers or configure rate limiting:

config = ProcessorConfig(
    max_workers=3,  # Lower concurrency
    rate_limit=RateLimitConfig(cooldown_seconds=60.0)
)

Validation Errors

pydantic.ValidationError: response doesn't match schema

Fix:

  1. Check your Pydantic model matches expected output
  2. Use progressive temperature strategy (increases temp on retries)
  3. Add examples in your prompt

Timeout Errors

asyncio.TimeoutError

Fix: Increase timeout:

config = ProcessorConfig(timeout_per_item=60.0)  # 60 seconds

Advanced: Smart Retry with on_error

Use the on_error callback to handle Gemini-specific errors intelligently:

Smart Model Escalation for Gemini

Only escalate to expensive Gemini models on validation errors, not network/rate limit errors:

from async_batch_llm.llm_strategies import LLMCallStrategy
from async_batch_llm import TokenUsage
from pydantic import ValidationError
from google import genai

class SmartGeminiStrategy(LLMCallStrategy[PersonData]):
    """Smart model escalation for Gemini API."""

    MODELS = [
        "gemini-2.5-flash-lite",  # Cheapest, fastest
        "gemini-2.5-flash",       # Production-ready
        "gemini-2.5-pro",         # Most capable
    ]

    def __init__(self, client: genai.Client):
        self.client = client
        self.validation_failures = 0  # Track quality issues only
        self.safety_blocks = 0        # Track Gemini safety blocks

    async def on_error(self, exception: Exception, attempt: int, state=None) -> None:
        """Track Gemini-specific error types."""
        if isinstance(exception, ValidationError):
            self.validation_failures += 1
        elif "SAFETY" in str(exception) or "BLOCKED" in str(exception):
            self.safety_blocks += 1
            # Note: Could adjust safety_settings on retry

    async def execute(
        self, prompt: str, attempt: int, timeout: float, state=None
    ) -> tuple[PersonData, TokenUsage]:
        # Select model based on validation failures (not total attempts)
        model_index = min(self.validation_failures, len(self.MODELS) - 1)
        model = self.MODELS[model_index]

        # Adjust safety settings if we've hit safety blocks
        config = genai.types.GenerateContentConfig(
            temperature=0.7,
            response_mime_type="application/json",
            response_schema=PersonData,
        )

        if self.safety_blocks > 0:
            # Could make safety_settings more permissive
            # config.safety_settings = [...]
            pass

        response = await self.client.aio.models.generate_content(
            model=model,
            contents=prompt,
            config=config,
        )

        output = PersonData.model_validate_json(response.text)
        usage = response.usage_metadata
        tokens: TokenUsage = {
            "input_tokens": usage.prompt_token_count or 0,
            "output_tokens": usage.candidates_token_count or 0,
            "total_tokens": usage.total_token_count or 0,
        }

        return output, tokens

Cost Savings:

  • Validation error → Escalate to gemini-2.5-pro (quality issue)
  • Network error → Retry with gemini-2.5-flash-lite (transient issue)
  • Rate limit error → Retry with gemini-2.5-flash-lite (API quota)
  • Safety block → Retry with same model, adjusted safety settings
  • Result: 60-80% cost reduction vs. always using gemini-2.5-pro

Smart Retry Prompts for Gemini

Build targeted retry prompts based on Gemini validation errors:

class SmartRetryGeminiStrategy(LLMCallStrategy[PersonData]):
    """Tell Gemini exactly what failed in previous attempt."""

    def __init__(self, client: genai.Client):
        self.client = client
        self.last_error = None
        self.last_response = None

    async def on_error(self, exception: Exception, attempt: int, state=None) -> None:
        """Track validation errors for smart retry."""
        if isinstance(exception, ValidationError):
            self.last_error = exception

    async def execute(
        self, prompt: str, attempt: int, timeout: float, state=None
    ) -> tuple[PersonData, TokenUsage]:
        if attempt == 1:
            final_prompt = prompt
        else:
            # Build focused retry prompt
            final_prompt = self._create_retry_prompt(prompt)

        config = genai.types.GenerateContentConfig(
            temperature=0.7,
            response_mime_type="application/json",
            response_schema=PersonData,
        )

        response = await self.client.aio.models.generate_content(
            model="gemini-2.5-flash",
            contents=final_prompt,
            config=config,
        )

        try:
            output = PersonData.model_validate_json(response.text)
            usage = response.usage_metadata
            tokens: TokenUsage = {
                "input_tokens": usage.prompt_token_count or 0,
                "output_tokens": usage.candidates_token_count or 0,
                "total_tokens": usage.total_token_count or 0,
            }
            return output, tokens
        except ValidationError as e:
            self.last_response = response.text
            raise  # Framework calls on_error, then retries

    def _create_retry_prompt(self, original_prompt: str) -> str:
        """Create targeted retry prompt with field-specific feedback."""
        if not self.last_error:
            return original_prompt

        # Parse which fields succeeded vs failed
        failed_fields = []
        for error in self.last_error.errors():
            field = ".".join(str(loc) for loc in error["loc"])
            msg = error["msg"]
            failed_fields.append(f"  - {field}: {msg}")

        retry_prompt = f"""RETRY REQUEST: The previous response had validation errors.

ORIGINAL REQUEST:
{original_prompt}

VALIDATION ERRORS TO FIX:
{chr(10).join(failed_fields)}

Please provide a complete, valid JSON response that fixes these specific validation errors.
Ensure all fields match the required schema exactly."""

        return retry_prompt

Benefits:

  • Gemini knows exactly what went wrong
  • Focused on fixing specific fields
  • Higher success rate on retries
  • Lower token usage (shorter prompts)

Complete Examples:

  • examples/example_smart_model_escalation.py - Full implementation
  • examples/example_gemini_smart_retry.py - Complete smart retry example

Resources

Support

For issues with: