# Gemini Support Implementation Plan

> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.

**Goal:** Add Google Gemini as a third supported LLM provider alongside OpenAI and Anthropic.

**Architecture:** New `GeminiProvider` class reads `response.usage_metadata.prompt_token_count` / `candidates_token_count`. Detection added to `detect_provider()` via isinstance + fingerprint + model prefix. `session.track()` gains an optional `model=` override for cases where the response doesn't carry the model name (older `google-generativeai` SDK).

**Tech Stack:** Python 3.9+, pytest, SimpleNamespace mocks (no real Google SDK needed for tests)

---

### Task 1: Add Gemini models to pricing.py

**Files:**
- Modify: `agentbudget/pricing.py`
- Test: `tests/test_providers.py` (will add Gemini tests there in Task 3)

**Step 1: Add Gemini section to MODEL_PRICING dict**

Open `agentbudget/pricing.py`. After the last Anthropic entry, add a new `# Google Gemini` section:

```python
    # Google Gemini
    "gemini-2.0-flash": {"input": 0.10, "output": 0.40},
    "gemini-2.0-flash-lite": {"input": 0.075, "output": 0.30},
    "gemini-1.5-pro": {"input": 1.25, "output": 5.00},
    "gemini-1.5-flash": {"input": 0.075, "output": 0.30},
    "gemini-1.5-flash-8b": {"input": 0.0375, "output": 0.15},
```

**Step 2: Verify pricing calculation**

```bash
cd "/Users/davidezaninello/CODING/Personal Project/SOFTWARE/AgentBudget"
uv run python -c "
from agentbudget.pricing import calculate_cost
# gemini-2.0-flash: 1000 input ($0.0001) + 500 output ($0.0002) = $0.0003
print(calculate_cost('gemini-2.0-flash', 1000, 500))
assert abs(calculate_cost('gemini-2.0-flash', 1000, 500) - 0.0003) < 1e-10
print('OK')
"
```
Expected: `0.0003` then `OK`

**Step 3: Run existing tests to confirm no regression**

```bash
uv run pytest tests/ -q
```
Expected: `108 passed`

**Step 4: Commit**

```bash
git add agentbudget/pricing.py
git commit -m "feat(pricing): add Google Gemini models (gemini-2.0-flash, 1.5-pro, 1.5-flash)"
```

---

### Task 2: Create GeminiProvider

**Files:**
- Create: `agentbudget/providers/gemini.py`
- Test: `tests/test_providers.py` (add tests in Task 3)

**Step 1: Write the file**

Create `agentbudget/providers/gemini.py` with this exact content:

```python
"""Gemini provider adapter for agentbudget."""

from typing import Any

from agentbudget.providers.base import BaseProvider, UsageData


class GeminiProvider(BaseProvider):
    """Provider adapter for Google Gemini API responses.

    Handles GenerateContentResponse objects from both google-generativeai
    and google-genai SDKs. Extracts prompt_token_count (input) and
    candidates_token_count (output) from response.usage_metadata.

    Model name is read from response.model or response.model_version
    (available in google-genai >= 0.5). Falls back to "gemini" if absent,
    which will trigger a $0.00 pricing warning — pass model= to
    session.track() to override.
    """

    def extract_usage(self, response: Any) -> UsageData:
        """Extract token usage from a Gemini GenerateContentResponse.

        Args:
            response: A GenerateContentResponse from google-generativeai
                or google-genai SDK.

        Returns:
            UsageData with model, input_tokens, output_tokens, provider="gemini".

        Raises:
            ValueError: If the response has no usage_metadata.
        """
        if not hasattr(response, "usage_metadata") or response.usage_metadata is None:
            raise ValueError("Gemini response has no usage_metadata.")

        # Try multiple attributes; newer google-genai SDK exposes model_version
        model: str = (
            getattr(response, "model", None)
            or getattr(response, "model_version", None)
            or "gemini"
        )

        usage = response.usage_metadata
        input_tokens: int = getattr(usage, "prompt_token_count", 0) or 0
        output_tokens: int = getattr(usage, "candidates_token_count", 0) or 0

        return UsageData(
            model=model,
            input_tokens=input_tokens,
            output_tokens=output_tokens,
            provider="gemini",
        )
```

**Step 2: Verify syntax**

```bash
uv run python -m py_compile agentbudget/providers/gemini.py && echo "OK"
```
Expected: `OK`

**Step 3: Commit**

```bash
git add agentbudget/providers/gemini.py
git commit -m "feat(providers): add GeminiProvider extracting usage_metadata"
```

---

### Task 3: Update detect_provider + add Gemini tests

**Files:**
- Modify: `agentbudget/providers/base.py`
- Modify: `tests/test_providers.py`

**Step 1: Write failing tests first**

Open `tests/test_providers.py`. Add at the end (before the last closing line if any):

```python
# ── Gemini ────────────────────────────────────────────────────────────────────

class TestDetectProviderGemini:
    """detect_provider routes Gemini responses correctly."""

    def test_gemini_usage_metadata_fingerprint(self):
        """Response with usage_metadata.prompt_token_count → GeminiProvider."""
        from agentbudget.providers.gemini import GeminiProvider
        response = SimpleNamespace(
            usage_metadata=SimpleNamespace(
                prompt_token_count=100,
                candidates_token_count=50,
            )
        )
        provider = detect_provider(response)
        assert isinstance(provider, GeminiProvider)

    def test_gemini_model_prefix_fallback(self):
        """model='gemini-1.5-flash' prefix → GeminiProvider."""
        from agentbudget.providers.gemini import GeminiProvider
        response = SimpleNamespace(model="gemini-1.5-flash")
        provider = detect_provider(response)
        assert isinstance(provider, GeminiProvider)

    def test_gemini_model_prefix_20_flash(self):
        """model='gemini-2.0-flash' prefix → GeminiProvider."""
        from agentbudget.providers.gemini import GeminiProvider
        response = SimpleNamespace(model="gemini-2.0-flash")
        provider = detect_provider(response)
        assert isinstance(provider, GeminiProvider)


class TestGeminiProviderExtractUsage:
    """GeminiProvider.extract_usage() extracts correct UsageData."""

    def _make_response(
        self,
        model: str = "gemini-2.0-flash",
        prompt_tokens: int = 100,
        candidates_tokens: int = 50,
        use_model_version: bool = False,
    ) -> SimpleNamespace:
        kwargs: dict = {}
        if use_model_version:
            kwargs["model_version"] = model
        else:
            kwargs["model"] = model
        return SimpleNamespace(
            usage_metadata=SimpleNamespace(
                prompt_token_count=prompt_tokens,
                candidates_token_count=candidates_tokens,
            ),
            **kwargs,
        )

    def test_extract_usage_correct_tokens(self):
        from agentbudget.providers.gemini import GeminiProvider
        response = self._make_response(prompt_tokens=500, candidates_tokens=200)
        usage = GeminiProvider().extract_usage(response)
        assert usage.input_tokens == 500
        assert usage.output_tokens == 200

    def test_extract_usage_correct_model(self):
        from agentbudget.providers.gemini import GeminiProvider
        response = self._make_response(model="gemini-1.5-pro")
        usage = GeminiProvider().extract_usage(response)
        assert usage.model == "gemini-1.5-pro"

    def test_extract_usage_model_version_fallback(self):
        """Newer SDK exposes model_version instead of model."""
        from agentbudget.providers.gemini import GeminiProvider
        response = self._make_response(model="gemini-2.0-flash", use_model_version=True)
        usage = GeminiProvider().extract_usage(response)
        assert usage.model == "gemini-2.0-flash"

    def test_extract_usage_no_model_falls_back_to_gemini(self):
        """Older SDK with no model attr → 'gemini' fallback."""
        from agentbudget.providers.gemini import GeminiProvider
        response = SimpleNamespace(
            usage_metadata=SimpleNamespace(
                prompt_token_count=100,
                candidates_token_count=50,
            )
        )
        usage = GeminiProvider().extract_usage(response)
        assert usage.model == "gemini"

    def test_extract_usage_provider_name(self):
        from agentbudget.providers.gemini import GeminiProvider
        response = self._make_response()
        usage = GeminiProvider().extract_usage(response)
        assert usage.provider == "gemini"

    def test_extract_usage_missing_usage_metadata_raises(self):
        from agentbudget.providers.gemini import GeminiProvider
        with pytest.raises(ValueError, match="usage_metadata"):
            GeminiProvider().extract_usage(SimpleNamespace())

    def test_extract_usage_none_usage_metadata_raises(self):
        from agentbudget.providers.gemini import GeminiProvider
        with pytest.raises(ValueError):
            GeminiProvider().extract_usage(SimpleNamespace(usage_metadata=None))
```

**Step 2: Run tests — expect failures**

```bash
uv run pytest tests/test_providers.py -k "Gemini" -v
```
Expected: FAIL (GeminiProvider not yet in detect_provider, so fingerprint/prefix tests fail)

**Step 3: Update detect_provider in base.py**

Open `agentbudget/providers/base.py`. Update the function and docstring:

```python
def detect_provider(response: Any) -> "BaseProvider":
    """Auto-detect the appropriate provider from an API response object.

    Detection order:
    1. isinstance check for openai.types.chat.ChatCompletion → OpenAIProvider
    2. isinstance check for google.generativeai GenerateContentResponse → GeminiProvider
    3. isinstance check for google.genai GenerateContentResponse → GeminiProvider
    4. Check if response.usage has input_tokens attribute → AnthropicProvider
    5. Check if response.usage_metadata has prompt_token_count → GeminiProvider
    6. Fallback: check response.model prefix:
         "gpt-", "o1", "o3" → OpenAI
         "claude-"           → Anthropic
         "gemini-"           → Gemini
    7. Raise ValueError if no match.

    Args:
        response: The raw API response object from any supported provider.

    Returns:
        An instantiated provider adapter for the detected provider.

    Raises:
        ValueError: If the provider cannot be detected from the response.
    """
    # Step 1: isinstance check for OpenAI ChatCompletion
    try:
        import openai.types.chat
        if isinstance(response, openai.types.chat.ChatCompletion):
            from agentbudget.providers.openai import OpenAIProvider
            return OpenAIProvider()
    except ImportError:
        pass

    # Step 2: isinstance check for google-generativeai SDK
    try:
        import google.generativeai.types  # type: ignore[import]
        if isinstance(response, google.generativeai.types.GenerateContentResponse):
            from agentbudget.providers.gemini import GeminiProvider
            return GeminiProvider()
    except ImportError:
        pass

    # Step 3: isinstance check for google-genai SDK
    try:
        import google.genai.types  # type: ignore[import]
        if isinstance(response, google.genai.types.GenerateContentResponse):
            from agentbudget.providers.gemini import GeminiProvider
            return GeminiProvider()
    except ImportError:
        pass

    # Step 4: Check for Anthropic-style usage (input_tokens attribute)
    if hasattr(response, "usage") and hasattr(response.usage, "input_tokens"):
        from agentbudget.providers.anthropic import AnthropicProvider
        return AnthropicProvider()

    # Step 5: Check for Gemini-style usage (usage_metadata.prompt_token_count)
    if (
        hasattr(response, "usage_metadata")
        and hasattr(response.usage_metadata, "prompt_token_count")
    ):
        from agentbudget.providers.gemini import GeminiProvider
        return GeminiProvider()

    # Step 6: Fallback by model name prefix
    model = getattr(response, "model", "")
    if isinstance(model, str):
        if model.startswith(("gpt-", "o1", "o3")):
            from agentbudget.providers.openai import OpenAIProvider
            return OpenAIProvider()
        if model.startswith("claude-"):
            from agentbudget.providers.anthropic import AnthropicProvider
            return AnthropicProvider()
        if model.startswith("gemini-"):
            from agentbudget.providers.gemini import GeminiProvider
            return GeminiProvider()

    raise ValueError("Cannot detect provider from response object.")
```

**Step 4: Run Gemini tests — expect all pass**

```bash
uv run pytest tests/test_providers.py -k "Gemini" -v
```
Expected: all pass

**Step 5: Run full suite — expect no regressions**

```bash
uv run pytest tests/ -q
```
Expected: 108 + new Gemini tests = all pass

**Step 6: Commit**

```bash
git add agentbudget/providers/base.py tests/test_providers.py
git commit -m "feat(providers): add Gemini detection and tests (fingerprint + prefix)"
```

---

### Task 4: Add model= override to session.track()

**Files:**
- Modify: `agentbudget/session.py`
- Modify: `tests/test_session.py`

**Step 1: Write the failing test**

Open `tests/test_session.py`. Add at the end of the file:

```python
class TestCostSessionTrackGemini:
    """Gemini-specific tracking via session.track()."""

    def _make_gemini_response(self, model: str = "gemini-2.0-flash") -> SimpleNamespace:
        return SimpleNamespace(
            model=model,
            usage_metadata=SimpleNamespace(
                prompt_token_count=500,
                candidates_token_count=200,
            ),
        )

    def test_track_gemini_auto_detect(self, tmp_path, monkeypatch):
        """Gemini response auto-detected and cost calculated correctly."""
        import agentbudget.tracker as tm
        monkeypatch.setattr(tm.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tm.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        with CostSession(name="g", verbose=False) as s:
            s.track(self._make_gemini_response("gemini-2.0-flash"))

        # gemini-2.0-flash: 500 input @ $0.10/M + 200 output @ $0.40/M
        # = $0.00005 + $0.00008 = $0.00013
        assert abs(s.total_cost - 0.00013) < 1e-9
        assert s.total_tokens == 700

    def test_track_gemini_model_override(self, tmp_path, monkeypatch):
        """model= param overrides model from response."""
        import agentbudget.tracker as tm
        monkeypatch.setattr(tm.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tm.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        # Response has no model attr (old SDK) — model passed explicitly
        response = SimpleNamespace(
            usage_metadata=SimpleNamespace(
                prompt_token_count=500,
                candidates_token_count=200,
            )
        )
        with CostSession(name="g", verbose=False) as s:
            s.track(response, model="gemini-1.5-pro")

        # gemini-1.5-pro: 500 input @ $1.25/M + 200 output @ $5.00/M
        # = $0.000625 + $0.001 = $0.001625
        assert abs(s.total_cost - 0.001625) < 1e-9

    def test_track_model_override_also_works_for_openai(self, tmp_path, monkeypatch):
        """model= override works for any provider, not just Gemini."""
        import agentbudget.tracker as tm
        monkeypatch.setattr(tm.CostTracker, "_STORAGE_DIR", tmp_path)
        monkeypatch.setattr(tm.CostTracker, "_STORAGE_FILE", tmp_path / "sessions.json")

        response = SimpleNamespace(
            model="gpt-4o",
            usage=SimpleNamespace(prompt_tokens=100, completion_tokens=50),
        )
        with CostSession(name="s", verbose=False) as s:
            # Override the model reported in the response
            s.track(response, model="gpt-4o-mini")

        # gpt-4o-mini: 100 input @ $0.15/M + 50 output @ $0.60/M
        # = $0.000015 + $0.00003 = $0.000045
        assert abs(s.total_cost - 0.000045) < 1e-9
```

**Step 2: Run tests — expect failures**

```bash
uv run pytest tests/test_session.py -k "Gemini" -v
```
Expected: FAIL (`track()` doesn't accept `model=` yet)

**Step 3: Update session.track() signature**

Open `agentbudget/session.py`. Change the `track` method signature and body:

```python
def track(self, response: Any, label: str = "", model: Optional[str] = None) -> None:
    """Track an LLM API response and accumulate its cost.

    Detects the provider automatically (or uses the provider set in the
    constructor), extracts token usage, calculates cost, and appends the
    result as a step. After each call, checks the budget and warns or
    raises an exception if it is exceeded.

    Args:
        response: The raw API response object from OpenAI, Anthropic, or
            Gemini SDK.
        label: Optional human-readable label for this step (e.g., "generate",
            "summarize"). Defaults to "step-N" if not provided.
        model: Optional model name override. Use this when the response object
            does not include the model name (e.g., with older google-generativeai
            SDK). Example: model="gemini-2.0-flash".

    Raises:
        ValueError: If the provider cannot be detected from the response.
        BudgetExceededError: If budget is set, strict=True, and cost exceeds budget.
    """
    # Detect provider
    if self.provider == "auto":
        provider_adapter = detect_provider(response)
    elif self.provider == "openai":
        from agentbudget.providers.openai import OpenAIProvider
        provider_adapter = OpenAIProvider()
    elif self.provider == "anthropic":
        from agentbudget.providers.anthropic import AnthropicProvider
        provider_adapter = AnthropicProvider()
    elif self.provider == "gemini":
        from agentbudget.providers.gemini import GeminiProvider
        provider_adapter = GeminiProvider()
    else:
        raise ValueError(f"Unknown provider: {self.provider!r}. Use 'auto', 'openai', 'anthropic', or 'gemini'.")

    usage = provider_adapter.extract_usage(response)

    # model= param overrides whatever the response reported
    effective_model = model if model is not None else usage.model

    cost = calculate_cost(effective_model, usage.input_tokens, usage.output_tokens)

    step_label = label or f"step-{len(self._steps) + 1}"

    self._steps.append({
        "label": step_label,
        "model": effective_model,
        "provider": usage.provider,
        "input_tokens": usage.input_tokens,
        "output_tokens": usage.output_tokens,
        "cost": cost,
    })

    # Budget enforcement
    if self.budget is not None and self.total_cost > self.budget:
        msg = (
            f"[agentbudget] Budget of ${self.budget:.2f} exceeded. "
            f"Current cost: ${self.total_cost:.4f}"
        )
        if self.strict:
            raise BudgetExceededError(msg)
        else:
            warnings.warn(msg, BudgetExceededWarning, stacklevel=2)
```

Note: also update the docstring of `__init__` to mention `"gemini"` as a valid `provider` value.

**Step 4: Run Gemini session tests**

```bash
uv run pytest tests/test_session.py -k "Gemini" -v
```
Expected: all pass

**Step 5: Run full suite**

```bash
uv run pytest tests/ -q
```
Expected: all pass (at least 108 tests + the new ones)

**Step 6: Commit**

```bash
git add agentbudget/session.py tests/test_session.py
git commit -m "feat(session): add model= override to track() for Gemini SDK compatibility"
```

---

### Task 5: Update pyproject.toml, README, CHANGELOG

**Files:**
- Modify: `pyproject.toml`
- Modify: `README.md`
- Modify: `CHANGELOG.md`

**Step 1: Update pyproject.toml**

Add `gemini` extra and update `all`:

```toml
[project.optional-dependencies]
openai = ["openai>=1.0.0"]
anthropic = ["anthropic>=0.20.0"]
gemini = ["google-genai>=0.8.0"]
all = ["openai>=1.0.0", "anthropic>=0.20.0", "google-genai>=0.8.0"]
dev = ["pytest>=7.0", "pytest-cov>=4.0"]
```

**Step 2: Update README.md**

In the Features section, update the providers line from:
```
- **Supported providers** — OpenAI and Anthropic (Gemini, Groq, and others coming in v0.3)
```
to:
```
- **Supported providers** — OpenAI, Anthropic, and Google Gemini
```

Add a Gemini example after the Anthropic example in the Examples section:

```markdown
### Gemini

```python
from agentbudget import CostSession

with CostSession(name="gemini-pipeline", budget=0.05) as session:
    response = client.models.generate_content(
        model="gemini-2.0-flash",
        contents="Explain quantum computing"
    )
    session.track(response, label="explain")
```

**With the older `google-generativeai` SDK** (response doesn't include model name):

```python
model = genai.GenerativeModel("gemini-1.5-pro")
response = model.generate_content("Explain quantum computing")
session.track(response, label="explain", model="gemini-1.5-pro")
```
```

Remove Gemini from the Roadmap v0.3 section.

**Step 3: Update CHANGELOG.md**

Add a new entry before `[0.1.0]`:

```markdown
## [0.1.1] - 2026-03-01

### Added
- Google Gemini provider support (`google-genai` and `google-generativeai` SDKs)
- `model=` optional parameter to `session.track()` for cases where the SDK response does not include the model name
- Gemini models in pricing table: gemini-2.0-flash, gemini-2.0-flash-lite, gemini-1.5-pro, gemini-1.5-flash, gemini-1.5-flash-8b
- `gemini` optional extra: `pip install agentbudget[gemini]`
```

**Step 4: Run full test suite**

```bash
uv run pytest -q
```
Expected: all pass

**Step 5: Commit**

```bash
git add pyproject.toml README.md CHANGELOG.md
git commit -m "feat: add Gemini to extras, README examples, and CHANGELOG"
```

---

### Task 6: Final verification

**Step 1: Run full test suite with coverage**

```bash
uv run pytest --cov=agentbudget --cov-report=term-missing -q
```
Expected: all pass, ≥ 80% coverage per module

**Step 2: Smoke test the Gemini flow end-to-end**

```bash
uv run python -c "
from agentbudget import CostSession
from types import SimpleNamespace

# Simulate newer SDK (has model attribute)
response_new = SimpleNamespace(
    model='gemini-2.0-flash',
    usage_metadata=SimpleNamespace(prompt_token_count=500, candidates_token_count=200)
)

# Simulate older SDK (no model attribute)
response_old = SimpleNamespace(
    usage_metadata=SimpleNamespace(prompt_token_count=300, candidates_token_count=100)
)

with CostSession(name='gemini-test', budget=0.10) as s:
    s.track(response_new, label='new-sdk')
    s.track(response_old, label='old-sdk', model='gemini-1.5-flash')
"
```
Expected: formatted table with both steps, correct costs, no errors

**Step 3: Run the anthropic example to verify no regression**

```bash
uv run python examples/anthropic_example.py
uv run python examples/openai_example.py
```
Expected: both print summaries without errors

**Step 4: Tag the release**

```bash
git tag v0.1.1
```
(Push to GitHub when ready: `git push origin main --tags`)
