safe-genai

A Python library that wraps Google Gemini with a mandatory generate → format → validate pipeline, so every response is structured and semantically checked before you receive it.

How it works

Every call to SafeAgent.run() executes three steps in sequence:

1. Generate — calls Gemini (with Google Search grounding) using your system prompt
2. Format — coerces the raw text into your Pydantic schema via a second LLM call
3. Validate — verifies the output fulfills the original prompt requirements, rejecting evasive non-answers (skippable via skip_validation=True)

If validation fails, the pipeline retries from the top (configurable via max_retries). On final failure it raises SafeGenAIValidationError.

Installation

pip install safe-genai

Requires a Google Gemini API key or Vertex AI project:

export GEMINI_API_KEY="your-key-here"

Quickstart

from pydantic import BaseModel
from safe_genai import SafeAgent

class LatestRelease(BaseModel):
    version: str
    release_date: str
    notable_changes: list[str]

agent = SafeAgent(
    output_schema=LatestRelease,
    system_prompt=(
        "You are a software research assistant. "
        "Use Google Search to find the most recent release information. "
        "Always return the latest published version, never guess."
    ),
)

result = agent.run("What is the latest stable release of Python?")
print(result.output.version)
print(result.output.release_date)

Async usage

result = await agent.run_async("What is 2 + 2?")

Configuration

from safe_genai import SafeAgent, AgentConfig

agent = SafeAgent(
    output_schema=Answer,
    system_prompt="You are a helpful assistant.",
    additional_formatter_instructions="Always express numbers as digits, not words.",
    additional_validator_instructions="Reject answers with confidence below 0.8.",
    max_retries=2,
    config=AgentConfig(
        generator_model="gemini-3.1-pro-preview",
        formatter_model="gemini-3-flash-preview",
        validator_model="gemini-3-flash-preview",
        api_key="your-key",  # or set GEMINI_API_KEY env var
    ),
)

Vertex AI (gcloud ADC)

If you're already authenticated with gcloud auth application-default login, you can use Vertex AI instead of an API key:

from safe_genai import SafeAgent, AgentConfig

agent = SafeAgent(
    output_schema=Answer,
    system_prompt="You are a helpful assistant.",
    config=AgentConfig(
        vertexai=True,
        vertex_project="my-gcp-project",   # or set GOOGLE_CLOUD_PROJECT env var
        vertex_location="us-central1",     # default
    ),
)

No GEMINI_API_KEY is needed — the client uses your ambient gcloud credentials.

Error handling

from safe_genai import SafeGenAIValidationError, SafeGenAIFormatterError

try:
    result = agent.run("some prompt")
except SafeGenAIValidationError as e:
    print(f"Validation failed after {e.attempts} attempt(s): {e.feedback}")
    print(e.output)        # last Pydantic model produced
    print(e.raw_response)  # raw text from the generator
except SafeGenAIFormatterError as e:
    print(f"Could not coerce output into {e.schema_name}: {e.cause}")

Token usage

Every result includes a breakdown of token usage by pipeline step, accumulated across all retry attempts:

result = agent.run("What is the latest stable release of Python?")

print(result.usage.generator)  # Usage(input_tokens=..., output_tokens=...)
print(result.usage.formatter)  # Usage(input_tokens=..., output_tokens=...)
print(result.usage.validator)  # Usage(input_tokens=..., output_tokens=...)
print(result.usage.total)      # Usage(input_tokens=..., output_tokens=...)  ← sum of all three

Usage is a plain dataclass with two fields: input_tokens and output_tokens.

Development

uv sync --extra dev
uv run pytest
uv build

License

MIT — see LICENSE.