# safe-genai

[![PyPI version](https://img.shields.io/pypi/v/safe-genai)](https://pypi.org/project/safe-genai/)
[![Python](https://img.shields.io/pypi/pyversions/safe-genai)](https://pypi.org/project/safe-genai/)
[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
[![Ruff](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/ruff/main/assets/badge/v2.json)](https://github.com/astral-sh/ruff)

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** — runs a semantic pass/fail check on the formatted output

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

## Installation

```bash
pip install safe-genai
```

Requires a Google Gemini API key or Vertex AI project:

```bash
export GEMINI_API_KEY="your-key-here"
```

## Quickstart

```python
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

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

## Configuration

```python
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:

```python
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

```python
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:

```python
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

```bash
uv sync --extra dev
uv run pytest
uv build
```

## License

MIT — see [LICENSE](LICENSE).
