agent.py

"""Gemini function-calling loop with recorded tool calls.

Manual loop (not SDK auto-calling) so we can:
- emit an event per tool call (for SSE progress),
- keep every tool result for the verification layer.
"""

import json
import logging
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Iterator

from google import genai
from google.genai import errors as genai_errors
from google.genai import types

from config import get_settings
from tools import TOOL_DECLARATIONS, TOOL_FUNCTIONS
from verify import verify_answer

logger = logging.getLogger(__name__)

SYSTEM_PROMPT = (Path(__file__).parent / "prompts" / "system.txt").read_text(encoding="utf-8")

_client: genai.Client | None = None


def get_client() -> genai.Client:
    global _client
    if _client is None:
        settings = get_settings()
        if not settings.gemini_api_key:
            raise RuntimeError("GEMINI_API_KEY missing in backend/.env")
        _client = genai.Client(api_key=settings.gemini_api_key)
    return _client


@dataclass
class AgentEvent:
    type: str                 # tool_call | tool_result | answer | error
    data: dict = field(default_factory=dict)


def _to_contents(messages: list[dict]) -> list[types.Content]:
    contents = []
    for message in messages:
        role = "user" if message["role"] == "user" else "model"
        contents.append(types.Content(role=role, parts=[types.Part(text=message["content"])]))
    return contents


def run_agent(messages: list[dict]) -> Iterator[AgentEvent]:
    """Yield tool events, then a final verified answer event."""
    settings = get_settings()
    client = get_client()

    config = types.GenerateContentConfig(
        system_instruction=SYSTEM_PROMPT,
        temperature=0.2,
        tools=[types.Tool(function_declarations=TOOL_DECLARATIONS)],
        automatic_function_calling=types.AutomaticFunctionCallingConfig(disable=True),
        # Without this, 2.5-flash tends to lay out the numbers in its hidden
        # reasoning and then emit only the qualitative summary as output.
        thinking_config=types.ThinkingConfig(thinking_budget=0),
    )

    contents = _to_contents(messages)
    tool_results: list[dict] = []
    citations: list[dict] = []

    for _round in range(settings.max_tool_rounds):
        response = _generate_with_retry(client, settings.gemini_model, contents, config)
        candidate = response.candidates[0]
        function_calls = [
            part.function_call for part in candidate.content.parts if part.function_call
        ]

        if not function_calls:
            answer = (response.text or "").strip()
            verification = verify_answer(answer, tool_results)
            yield AgentEvent("answer", {
                "answer": answer,
                "citations": citations,
                "verification": verification,
            })
            return

        contents.append(candidate.content)
        response_parts = []
        for call in function_calls:
            name = call.name
            args = dict(call.args or {})
            yield AgentEvent("tool_call", {"name": name, "args": args})

            function = TOOL_FUNCTIONS.get(name)
            if function is None:
                result = {"error": f"unknown tool {name}"}
            else:
                try:
                    result = function(**args)
                except Exception as exc:
                    logger.exception("tool %s failed", name)
                    result = {"error": str(exc)}

            tool_results.append(result)
            _collect_citations(name, result, citations)
            yield AgentEvent("tool_result", {"name": name, "summary": _summarize(result)})
            response_parts.append(types.Part.from_function_response(name=name, response=result))

        contents.append(types.Content(role="tool", parts=response_parts))

    yield AgentEvent("error", {"message": "Tool budget exhausted before an answer was produced."})


def _generate_with_retry(client, model, contents, config, attempts: int = 4):
    """Retry transient Gemini errors (503 overload, 429 rate limit) with backoff."""
    for attempt in range(attempts):
        try:
            return client.models.generate_content(model=model, contents=contents, config=config)
        except genai_errors.APIError as exc:
            if exc.code == 429 and "PerDay" in str(exc):
                # Daily quota exhausted — retrying only burns more of tomorrow's.
                raise RuntimeError(
                    "Gemini free-tier daily quota is exhausted. Enable billing on the "
                    "key or wait for the reset (midnight Pacific)."
                ) from exc
            if exc.code in (429, 503) and attempt < attempts - 1:
                delay = 2 ** attempt * 2  # 2s, 4s, 8s
                logger.warning("Gemini %s, retrying in %ds (%d/%d)", exc.code, delay, attempt + 1, attempts)
                time.sleep(delay)
                continue
            raise


def _summarize(result: dict) -> str:
    if "error" in result:
        return f"error: {result['error']}"
    if "metrics" in result:
        counts = {metric: len(rows) for metric, rows in result["metrics"].items()}
        return f"facts for {result.get('ticker')}: {counts}"
    if "passages" in result:
        return f"{len(result['passages'])} passages"
    if "companies" in result:
        return f"{len(result['companies'])} companies"
    if "price" in result:
        return f"{result.get('symbol')} @ {result['price']}"
    return json.dumps(result)[:120]


def _collect_citations(tool: str, result: dict, citations: list[dict]) -> None:
    if tool == "query_facts" and "metrics" in result:
        for metric, rows in result["metrics"].items():
            for row in rows:
                citations.append({
                    "kind": "xbrl",
                    "ticker": result.get("ticker"),
                    "metric": metric,
                    "fiscal_year": row.get("fiscal_year"),
                    "period_end": row.get("period_end"),
                    "value": row.get("value"),
                    "unit": row.get("unit"),
                    "accession": row.get("source_accession"),
                    "form": row.get("source_form"),
                })
    elif tool == "retrieve_passages" and "passages" in result:
        for passage in result["passages"]:
            citations.append({
                "kind": "passage",
                "ticker": passage.get("ticker"),
                "form": passage.get("form"),
                "filing_date": passage.get("filing_date"),
                "accession": passage.get("accession"),
                "section": passage.get("section"),
                "chunk_id": passage.get("chunk_id"),
                "score": passage.get("score"),
            })