backend/core/agent_loop.py

"""
The agent loop: model proposes tool calls, we execute them, results go back in.

This replaces a fixed pipeline (detect intent -> pick tables -> write SQL ->
execute -> explain) with an adaptive loop. The difference that matters is that
the agent can *observe before it acts*: check a column's real values, look at a
schema, run a cheap probe query, then commit. The old pipeline had to guess in
one shot and had exactly one retry.

The loop is written explicitly rather than using the SDK's automatic function
calling, because every step needs to be streamed to the UI as a reasoning entry
and because tool failures must be fed back to the model rather than raised.
"""

from __future__ import annotations

import asyncio
import json
import logging
import re
from typing import Any, AsyncIterator, Dict, List, Optional

from google.genai import types

from backend.core.agent_tools import (
    AgentContext, build_tools, call_tool, serialize_result,
)

logger = logging.getLogger(__name__)

# Enough turns for discover -> inspect -> probe -> build -> summarize, with room
# to recover from a couple of mistakes. Beyond this the agent is looping.
MAX_ITERATIONS = 14

# Words that mean the user expects something to look at. Used only to catch the
# case where the agent answers in prose without producing the artifact.
_VISUAL_REQUEST = re.compile(
    r"\b(plot|map|chart|graph|show|display|visuali[sz]e|animate|draw|"
    r"compare|rank|distribution|where)\b",
    re.IGNORECASE,
)

AGENT_SYSTEM_PROMPT = """You are Perch, a geospatial analyst for avian biodiversity data.

You answer questions by calling tools. You can see the data before you commit to
an answer — use that.

## How to work

1. **Find the data.** If you do not already know which table holds the answer,
   call `search_datasets`. Then `describe_table` for the exact columns.
2. **Check before you filter.** Before writing a WHERE clause on a value whose
   format you are not certain of — a region code, a category, a date — call
   `sample_values`. A wrong guess returns zero rows silently, which is worse
   than an error. This is the single most common way to get a confidently wrong
   answer, so spend the one extra call.
3. **Build the output — this is not optional.** If the user asked to see, show,
   map, plot, chart, visualize, display, compare or rank anything, you MUST call
   `add_map_layer` (for anything with geometry) and/or `make_chart` (for
   rankings and comparisons) BEFORE you answer. Running `run_sql` only computes
   numbers for you; it puts nothing in front of the user. An answer that
   describes results the user cannot see has failed, however accurate the prose.
   Many questions deserve both a map and a chart.

   **One layer per subject; one layer for all the categories of a subject.**

   A *subject* is a thing being compared against another thing — a species, a
   region, a year. Comparing three species means three `add_map_layer` calls,
   named so they can be told apart. Layers accumulate, so a later call never
   replaces an earlier one, and mapping only one of the things being compared
   makes the map contradict your own answer.

   But the *categories within* one subject belong on a single layer, returned by
   one query with the category column included — the four seasonal ranges of one
   species are one `add_map_layer` call over a result with a `season` column, not
   four calls. The map colours them by category and gives the user a checkbox per
   category, which is the only way to read ranges that overlap. Splitting them
   into separate layers throws that away and buries the real comparison in a list
   of near-identical entries.

   So: three species' breeding ranges → three layers. One species' four seasons →
   one layer. Three species across four seasons → three layers, each carrying its
   own season column.
4. **Then answer.** When you have what you need, reply with plain text. That
   final message is what the user reads.

## Working efficiently

You have a limited number of steps, so spend them on progress rather than
reassurance:

- Do not re-verify something you already established. One `sample_values` per
  uncertain column is enough.
- Do not run a query just to preview what a later query will return. Go
  straight to `add_map_layer` or `make_chart` once you know the shape.
- Batch independent lookups into a single turn — several tool calls at once run
  in parallel.
- Aim to be producing output by roughly your fifth step.

## Judgement

- Prefer acting over asking. Make a sensible choice and say what you chose.
  Use `ask_user` only when the request is genuinely ambiguous and guessing would
  waste their time.
- If a query returns zero rows, do not report "no data" until you have checked
  the actual values with `sample_values`. The data is usually there.
- Non-spatial tables have no geometry. To map their values, join to a boundary
  table to borrow its geometry.

## Writing efficient SQL

Results are not truncated, so what you ask for is what gets built and rendered.
Ask for the right thing:

- **Return what answers the question, not the whole table.** "Where is this
  species most abundant" wants the abundance grid; it does not want every week
  of every species. Select the columns you need, not `*`.
- **Aggregate in SQL, not by returning rows.** For per-region answers, GROUP BY
  the region and return one row per region — not every hexagon inside it. For a
  yearly summary use the year-round table rather than averaging 52 weekly rows
  per hexagon yourself.
- **Filter early.** Put the season, week, species or country filter in the query
  rather than returning everything and describing a subset.
- **Pick the right grain.** A weekly table has ~52x the rows of its year-round
  equivalent. Use `_weekly` only when the question is about change over time,
  `_seasonal` for season comparisons, `_abundance` for a single summary.
- If a query is refused as too large to render, do not retry it unchanged —
  aggregate it or narrow it, then retry.
- To summarise a fine grid (hexagons) into regions, spatially join the grid to
  boundary polygons, GROUP BY the region, and keep the region geometry with
  ANY_VALUE(geometry) so the result can still be mapped.
- Relative abundance is the mean count expected on a standard eBird checklist.
  It is an index, not a population census — describe it as relative abundance.

## Answering

Write for someone who did not watch you work. Lead with the finding, give the
numbers that support it, name the tables you used. Be concise and concrete; do
not narrate your tool calls. If a result was capped or a caveat applies, say so
plainly rather than implying the answer is complete.
"""


def _fn_args(call: Any) -> Dict[str, Any]:
    """Normalize a function call's arguments to a plain dict."""
    args = getattr(call, "args", None) or {}
    if isinstance(args, dict):
        return dict(args)
    try:
        return json.loads(args)
    except (TypeError, ValueError):
        return {}


def _describe_call(name: str, args: Dict[str, Any]) -> str:
    """One-line human summary of a tool call, for the reasoning timeline."""
    if name == "search_datasets":
        return f"Searching datasets for “{args.get('query', '')}”"
    if name == "describe_table":
        return f"Inspecting schema of {args.get('table', '')}"
    if name == "sample_values":
        return f"Checking real values of {args.get('table', '')}.{args.get('column', '')}"
    if name == "run_sql":
        return args.get("purpose") or "Running a query"
    if name == "add_map_layer":
        return f"Mapping “{args.get('name', 'result')}”"
    if name == "make_chart":
        return f"Charting “{args.get('title', '')}”"
    if name == "compute_stats":
        return f"Computing statistics for {args.get('column', '')}"
    if name == "ask_user":
        return "Asking a clarifying question"
    if name == "spawn_subagents":
        tasks = args.get("tasks") or []
        return f"Delegating {len(tasks)} parallel investigations"
    return f"Calling {name}"


def _summarize_outcome(name: str, payload: Dict[str, Any]) -> Optional[str]:
    """Short result line for the timeline, or None to stay quiet."""
    if not payload.get("ok"):
        return f"↳ {payload.get('error', 'failed')}"
    r = payload.get("result") or {}
    if name == "search_datasets":
        n = len(r.get("results") or [])
        return f"↳ {n} candidate dataset(s)"
    if name == "sample_values":
        sample = r.get("sample") or []
        return f"↳ e.g. {', '.join(map(str, sample[:4]))}" if sample else None
    if name == "run_sql":
        return f"↳ {r.get('row_count', 0):,} rows"
    if name == "add_map_layer":
        return f"↳ {r.get('features', 0):,} features on the map"
    if name == "make_chart":
        return f"↳ {r.get('points', 0)} points"
    if name == "describe_table":
        return f"↳ {len(r.get('columns') or [])} columns, {r.get('rows') or 0:,} rows"
    return None


class GeoAgent:
    """Runs one question to completion, streaming progress as it goes."""

    def __init__(self, client, model: str, extra_tools: Optional[Dict[str, Any]] = None):
        self.client = client
        self.model = model
        self.extra_tools = extra_tools or {}

    # Transient upstream failures (503 overloaded, 429 rate limit, 500) would
    # otherwise abort a turn that had already done most of its work, losing every
    # layer and chart the agent had built. Retry them; leave real errors alone.
    _TRANSIENT = ("503", "UNAVAILABLE", "429", "RESOURCE_EXHAUSTED", "500", "INTERNAL")

    async def _generate_with_retry(self, contents, config, attempts: int = 3):
        last: Exception | None = None
        for attempt in range(attempts):
            try:
                return await asyncio.wait_for(
                    asyncio.to_thread(
                        self.client.models.generate_content,
                        model=self.model, contents=contents, config=config,
                    ),
                    timeout=120.0,
                )
            except asyncio.TimeoutError:
                raise
            except Exception as e:  # noqa: BLE001
                last = e
                if not any(t in str(e) for t in self._TRANSIENT) or attempt == attempts - 1:
                    raise
                delay = 1.5 * (2 ** attempt)
                logger.warning(f"Transient model error, retrying in {delay:.1f}s: {e}")
                await asyncio.sleep(delay)
        raise last  # pragma: no cover - loop always returns or raises

    async def run(
        self,
        question: str,
        history: List[Dict[str, str]],
        ctx: AgentContext,
        max_iterations: int = MAX_ITERATIONS,
    ) -> AsyncIterator[Dict[str, Any]]:
        """
        Yield progress events, then a final {"type": "final", ...}.

        Events: {"type": "step"|"thought"|"final"|"error", ...}. The caller maps
        these onto SSE; keeping them abstract means the loop does not know about
        the transport.
        """
        tools = build_tools(ctx)
        tools.update(self.extra_tools)

        tool_config = types.Tool(
            function_declarations=[t.declaration() for t in tools.values()]
        )

        contents: List[types.Content] = []
        for msg in history[-8:]:
            role = "model" if msg.get("role") == "assistant" else "user"
            text = (msg.get("content") or "").strip()
            if text:
                contents.append(types.Content(
                    role=role, parts=[types.Part.from_text(text=text)]
                ))
        contents.append(types.Content(
            role="user", parts=[types.Part.from_text(text=question)]
        ))

        config = types.GenerateContentConfig(
            system_instruction=AGENT_SYSTEM_PROMPT,
            tools=[tool_config],
            # Let the model choose freely between calling a tool and answering.
            tool_config=types.ToolConfig(
                function_calling_config=types.FunctionCallingConfig(mode="AUTO")
            ),
            # Low thinking per turn, deliberately. In a tool loop the reasoning is
            # externalized — the agent learns by calling sample_values and reading
            # the result, not by thinking harder about what the value might be. A
            # medium budget on every one of a dozen turns multiplies latency
            # without improving the decisions, which are mostly "which tool next".
            thinking_config=types.ThinkingConfig(thinking_level="low"),
        )

        nudged = False

        for iteration in range(max_iterations):
            try:
                response = await self._generate_with_retry(contents, config)
            except asyncio.TimeoutError:
                yield {"type": "error", "message": "The model timed out. Please try again."}
                return
            except Exception as e:  # noqa: BLE001
                logger.error(f"Agent generate_content failed: {e}", exc_info=True)
                yield {"type": "error", "message": f"Agent error: {e}"}
                return

            candidate = (response.candidates or [None])[0]
            content = getattr(candidate, "content", None)
            parts = list(getattr(content, "parts", None) or [])
            if not parts:
                yield {"type": "final", "text": response.text or "I could not produce an answer."}
                return

            calls = [p.function_call for p in parts if getattr(p, "function_call", None)]
            text_parts = [
                p.text for p in parts
                if getattr(p, "text", None) and not getattr(p, "thought", False)
            ]

            # No tool calls means the model is answering.
            if not calls:
                answer = "\n".join(t for t in text_parts if t).strip()

                # Safety net: the user asked to see something, the agent answered
                # in prose without producing it. Prompting alone is not reliable
                # enough here, because a fluent description reads like success to
                # the model while the user's screen stays empty. Nudge once.
                wants_visual = bool(_VISUAL_REQUEST.search(question))
                produced = bool(ctx.layers) or ctx.chart_data is not None
                if wants_visual and not produced and not nudged and not ctx.pending_question:
                    nudged = True
                    yield {"type": "step", "text": "Producing the visual output"}
                    contents.append(content)
                    contents.append(types.Content(role="user", parts=[types.Part.from_text(
                        text=(
                            "You have not produced anything the user can see. They asked to "
                            "visualize this. Call add_map_layer now if the result has geometry "
                            "(join to a boundary table to borrow geometry if needed), and/or "
                            "make_chart for a ranking or comparison. Then give your answer."
                        )
                    )]))
                    continue

                yield {"type": "final", "text": answer or "I could not produce an answer."}
                return

            # Keep the model's own turn in the transcript so it remembers what it asked for.
            contents.append(content)

            # Execute the requested calls. Several in one turn run concurrently —
            # independent lookups should not be serialized.
            async def _invoke(fc):
                name = fc.name
                args = _fn_args(fc)
                tool = tools.get(name)
                if tool is None:
                    return name, args, {"ok": False, "error": f"Unknown tool '{name}'."}
                return name, args, await call_tool(tool, args)

            for fc in calls:
                yield {"type": "step", "text": _describe_call(fc.name, _fn_args(fc))}

            results = await asyncio.gather(*(_invoke(fc) for fc in calls))

            response_parts = []
            for name, args, payload in results:
                note = _summarize_outcome(name, payload)
                if note:
                    yield {"type": "step", "text": note}
                response_parts.append(types.Part.from_function_response(
                    name=name, response=serialize_result(payload)
                ))

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

            # A clarifying question ends the turn — waiting for the user is the
            # whole point, and continuing would answer a question they did not ask.
            if ctx.pending_question:
                yield {
                    "type": "final",
                    "text": ctx.pending_question["question"],
                    "question": ctx.pending_question,
                }
                return

        # Out of iterations. Ask for a final answer with tools switched off, so the
        # user gets whatever was actually learned rather than an empty failure.
        yield {"type": "step", "text": "Wrapping up"}
        try:
            final = await asyncio.wait_for(
                asyncio.to_thread(
                    self.client.models.generate_content,
                    model=self.model,
                    contents=contents + [types.Content(
                        role="user",
                        parts=[types.Part.from_text(text=(
                            "Stop calling tools and answer now using what you have "
                            "already found. If the answer is incomplete, say so."
                        ))],
                    )],
                    config=types.GenerateContentConfig(
                        system_instruction=AGENT_SYSTEM_PROMPT,
                        thinking_config=types.ThinkingConfig(thinking_level="low"),
                    ),
                ),
                timeout=60.0,
            )
            yield {"type": "final", "text": final.text or "I ran out of steps before finishing."}
        except Exception as e:  # noqa: BLE001
            logger.error(f"Agent wrap-up failed: {e}")
            yield {"type": "final", "text": "I ran out of steps before finishing this question."}