Download backend/core/agent_tools.py from GerardCB/perch: direct link, hf CLI and curl.
- Browser
- Download file 20.5 kB
-
https://huggingface.co/spaces/GerardCB/perch/resolve/main/backend/core/agent_tools.py
- Command line
-
hf download hf://spaces/GerardCB/perch/backend/core/agent_tools.py
-
curl -L -o agent_tools.py https://huggingface.co/spaces/GerardCB/perch/resolve/main/backend/core/agent_tools.py
20.5 kB
| """ | |
| Tools the geospatial agent can call. | |
| Each tool pairs a Gemini `FunctionDeclaration` (what the model sees) with a Python | |
| implementation (what actually runs). The agent decides which to call and reacts to | |
| what comes back, instead of following a fixed detect-intent -> pick-tables -> | |
| write-SQL pipeline. | |
| The point of the toolset is that the agent can **look before it commits**. Every | |
| bug class the old pipeline hit came from writing SQL blind against a schema | |
| summary: guessing `US-%` when the data uses `USA-%`, colouring by `week_index` | |
| because it happened to be the first numeric column, assuming a full-year raster | |
| exists for a resident species. `sample_values` and `describe_table` make those | |
| answerable in one cheap call rather than a wrong answer delivered confidently. | |
| Results are deliberately compact. `run_sql` returns a preview and a row count, | |
| never the full payload β the geometry goes to the map via `add_map_layer`, and | |
| stuffing thousands of rows back into the prompt would blow the context and teach | |
| the model nothing extra. | |
| """ | |
| from __future__ import annotations | |
| import asyncio | |
| import json | |
| import logging | |
| import re | |
| from dataclasses import dataclass, field | |
| from typing import Any, Callable, Dict, List, Optional | |
| from google.genai import types | |
| from backend.core.jsonutil import dumps_safe, json_safe | |
| logger = logging.getLogger(__name__) | |
| # A tool call that runs longer than this is almost certainly a runaway spatial | |
| # join; failing it returns control to the agent, which can narrow the query. | |
| TOOL_TIMEOUT_SECONDS = 90.0 | |
| # Rows echoed back into the conversation from run_sql. Enough to see the shape of | |
| # the result and spot obvious mistakes, small enough not to crowd the context. | |
| PREVIEW_ROWS = 8 | |
| # Statements the agent may execute. Anything that writes is rejected outright: | |
| # the agent is an analyst, not an administrator. | |
| _FORBIDDEN_SQL = re.compile( | |
| r"\b(INSERT|UPDATE|DELETE|DROP|ALTER|CREATE|TRUNCATE|ATTACH|COPY|INSTALL|LOAD|PRAGMA|EXPORT)\b", | |
| re.IGNORECASE, | |
| ) | |
| class ToolError(Exception): | |
| """A tool failed in a way the agent can reason about and retry.""" | |
| class Tool: | |
| """One callable tool: its schema for the model, and its implementation.""" | |
| name: str | |
| description: str | |
| parameters: Dict[str, Any] | |
| run: Callable[..., Any] | |
| # Tools that produce user-visible output rather than just information. | |
| produces_output: bool = False | |
| def declaration(self) -> types.FunctionDeclaration: | |
| return types.FunctionDeclaration( | |
| name=self.name, | |
| description=self.description, | |
| parameters=self.parameters, | |
| ) | |
| class AgentContext: | |
| """ | |
| Mutable state shared across one agent run. | |
| Collects everything the tools produce so the caller can emit a single final | |
| response: the map layer, chart, citations and stats accumulate here rather | |
| than being returned through the model, which would mean serializing large | |
| payloads into the prompt. | |
| """ | |
| allowed_datasets: Optional[List[str]] = None | |
| # A list, not a single slot: the agent legitimately builds several layers for | |
| # one answer (one per species in a comparison). Keeping only the newest threw | |
| # away work it had already done and made the map contradict the text. | |
| layers: List[Dict[str, Any]] = field(default_factory=list) | |
| chart_data: Optional[Dict[str, Any]] = None | |
| raw_data: List[Dict[str, Any]] = field(default_factory=list) | |
| sql_statements: List[str] = field(default_factory=list) | |
| tables_used: List[str] = field(default_factory=list) | |
| stats: Dict[str, Any] = field(default_factory=dict) | |
| pending_question: Optional[Dict[str, Any]] = None | |
| def _validate_sql(sql: str) -> str: | |
| """Reject anything that is not a read-only query.""" | |
| cleaned = sql.strip().rstrip(";").strip() | |
| if not cleaned: | |
| raise ToolError("Empty SQL.") | |
| if _FORBIDDEN_SQL.search(cleaned): | |
| raise ToolError( | |
| "Only read-only SELECT/WITH queries are permitted. " | |
| "Remove any statement that modifies data or schema." | |
| ) | |
| first = cleaned.lstrip("(").split(None, 1)[0].upper() | |
| if first not in ("SELECT", "WITH"): | |
| raise ToolError(f"Query must start with SELECT or WITH, got '{first}'.") | |
| return cleaned | |
| def build_tools(ctx: AgentContext) -> Dict[str, Tool]: | |
| """ | |
| Construct the toolset bound to one agent run. | |
| Imports are local so this module can be imported without spinning up DuckDB | |
| or the embedding index (useful for tests and for the schema-only path). | |
| """ | |
| from backend.core.data_catalog import get_data_catalog, describe_table | |
| from backend.core.geo_engine import get_geo_engine | |
| from backend.core.semantic_search import get_semantic_search | |
| from backend.services.response_formatter import ResponseFormatter | |
| catalog = get_data_catalog() | |
| engine = get_geo_engine() | |
| semantic = get_semantic_search() | |
| def _in_scope(name: str) -> bool: | |
| return ctx.allowed_datasets is None or name in ctx.allowed_datasets | |
| # ---------------------------------------------------------------- discovery | |
| def search_datasets(query: str, limit: int = 8) -> Dict[str, Any]: | |
| hits = semantic.search(query, top_k=max(1, min(limit, 20)), | |
| allowed_datasets=ctx.allowed_datasets) | |
| results = [] | |
| for name, score in hits: | |
| meta = catalog.get_table_metadata(name) | |
| if not meta: | |
| continue | |
| results.append({ | |
| "table": name, | |
| "relevance": round(float(score), 3), | |
| "rows": meta.get("row_count"), | |
| "description": describe_table(meta)[:400], | |
| }) | |
| if not results: | |
| return {"results": [], "hint": "Nothing matched. Try broader wording, " | |
| "or call describe_table on a known table."} | |
| return {"results": results} | |
| def describe_table_tool(table: str) -> Dict[str, Any]: | |
| meta = catalog.get_table_metadata(table) | |
| if not meta: | |
| close = [n for n in catalog.catalog if table.lower() in n.lower()][:5] | |
| raise ToolError( | |
| f"No table named '{table}'." + | |
| (f" Did you mean: {', '.join(close)}?" if close else | |
| " Use search_datasets to find one.") | |
| ) | |
| if not _in_scope(table): | |
| raise ToolError(f"'{table}' is outside the datasets selected for this session.") | |
| if not engine.ensure_table_loaded(table): | |
| raise ToolError(f"'{table}' is in the catalog but its data file could not be loaded.") | |
| columns = engine.describe_columns(table) | |
| has_geometry = any(c[0] in engine.GEOMETRY_COLUMNS for c in columns) | |
| return { | |
| "table": table, | |
| "rows": meta.get("row_count"), | |
| "spatial": has_geometry, | |
| "columns": [{"name": c[0], "type": c[1]} for c in columns], | |
| "description": describe_table(meta), | |
| "attribution": meta.get("attribution"), | |
| } | |
| def sample_values(table: str, column: str, limit: int = 15) -> Dict[str, Any]: | |
| """Distinct values of a column β how the agent learns real formats.""" | |
| if not _in_scope(table): | |
| raise ToolError(f"'{table}' is outside the datasets selected for this session.") | |
| if not engine.ensure_table_loaded(table): | |
| raise ToolError(f"Could not load '{table}'.") | |
| limit = max(1, min(limit, 50)) | |
| try: | |
| distinct = engine.fetch_all( | |
| f'SELECT DISTINCT "{column}" FROM "{table}" ' | |
| f'WHERE "{column}" IS NOT NULL LIMIT {limit}' | |
| ) | |
| total = engine.fetch_one( | |
| f'SELECT COUNT(DISTINCT "{column}") FROM "{table}"' | |
| )[0] | |
| except Exception as e: | |
| raise ToolError(f"Could not sample '{column}' from '{table}': {e}") | |
| values = [r[0] for r in distinct] | |
| out: Dict[str, Any] = { | |
| "table": table, | |
| "column": column, | |
| "distinct_count": total, | |
| "sample": [str(v) for v in values], | |
| } | |
| # Numeric columns: the range matters more than example values. | |
| if values and isinstance(values[0], (int, float)) and not isinstance(values[0], bool): | |
| lo, hi = engine.fetch_one( | |
| f'SELECT MIN("{column}"), MAX("{column}") FROM "{table}"' | |
| ) | |
| out["min"], out["max"] = lo, hi | |
| if total > limit: | |
| out["note"] = f"Showing {len(values)} of {total} distinct values." | |
| return out | |
| # ----------------------------------------------------------------- querying | |
| def _execute(sql: str) -> Dict[str, Any]: | |
| from backend.core.geo_engine import ResultTooLargeError # noqa: F401 | |
| cleaned = _validate_sql(sql) | |
| for name in catalog.catalog: | |
| if re.search(rf'(?<![\w."]){re.escape(name)}(?![\w."])', cleaned, re.IGNORECASE): | |
| if not _in_scope(name): | |
| raise ToolError(f"'{name}' is outside the datasets selected for this session.") | |
| engine.ensure_table_loaded(name) | |
| try: | |
| return engine.execute_spatial_query(cleaned) | |
| except ResultTooLargeError as e: | |
| # Not a failure of the query's logic β it is simply too coarse. The | |
| # message names the fix, so the agent rewrites rather than gives up. | |
| raise ToolError(str(e)) | |
| except Exception as e: | |
| # Hand the database's own message back; it is usually enough for the | |
| # agent to fix the query itself. | |
| raise ToolError(f"SQL failed: {e}") | |
| def run_sql(sql: str, purpose: str = "") -> Dict[str, Any]: | |
| result = _execute(sql) | |
| features = result.get("features", []) | |
| props = result.get("properties", {}) or {} | |
| ctx.sql_statements.append(sql.strip()) | |
| preview = [ | |
| {k: v for k, v in (f.get("properties") or {}).items()} | |
| for f in features[:PREVIEW_ROWS] | |
| ] | |
| out: Dict[str, Any] = { | |
| "row_count": len(features), | |
| "has_geometry": any(f.get("geometry") for f in features), | |
| "columns": list(preview[0].keys()) if preview else [], | |
| "preview": preview, | |
| } | |
| if not features: | |
| out["hint"] = ("Zero rows. Check literal values with sample_values before " | |
| "assuming a format, and verify filters match real data.") | |
| return out | |
| def add_map_layer(sql: str, name: str, color_by: str = "") -> Dict[str, Any]: | |
| result = _execute(sql) | |
| features = result.get("features", []) | |
| if not features: | |
| raise ToolError("Query returned no rows, so there is nothing to map.") | |
| if not any(f.get("geometry") for f in features): | |
| raise ToolError( | |
| "Result has no geometry and cannot be mapped. Either select a geometry " | |
| "column, or join to a boundary table to borrow one." | |
| ) | |
| ctx.sql_statements.append(sql.strip()) | |
| geojson, _layer_id, layer_name = ResponseFormatter.format_geojson_layer( | |
| name, result, features, name, "π", None, color_by=color_by or None, | |
| ) | |
| # Credit the tables THIS query read, not every table the agent touched | |
| # while answering. A three-species comparison runs one query per species, | |
| # and attributing the union to all of them told the reader a layer was | |
| # built from three datasets when it came from one β and made the layer's | |
| # species look ambiguous when it was not. | |
| geojson.setdefault("properties", {})["source_tables"] = ( | |
| ResponseFormatter._tables_referenced_in_sql(sql, list(catalog.catalog.keys())) | |
| ) | |
| ctx.layers.append(geojson) | |
| return { | |
| "layer": layer_name, | |
| "features": len(features), | |
| "layers_on_map": len(ctx.layers), | |
| "note": "Layer added to the map. Call again to add another for a " | |
| "different species or season; existing layers are kept.", | |
| } | |
| def make_chart(sql: str, chart_type: str, title: str) -> Dict[str, Any]: | |
| result = _execute(sql) | |
| features = result.get("features", []) | |
| if not features: | |
| raise ToolError("Query returned no rows, so there is nothing to chart.") | |
| ctx.sql_statements.append(sql.strip()) | |
| chart = ResponseFormatter.generate_chart_data( | |
| sql, features, title, | |
| {"use_chart": True, "type": chart_type, "title": title}, | |
| ) | |
| if not chart: | |
| raise ToolError( | |
| "Could not build a chart from that result β it needs a label column " | |
| "and a numeric column." | |
| ) | |
| ctx.chart_data = chart | |
| ctx.raw_data = ResponseFormatter.prepare_raw_data(features) | |
| return {"chart": chart.get("type"), "title": chart.get("title"), | |
| "points": len(chart.get("data") or [])} | |
| def compute_stats(sql: str, column: str) -> Dict[str, Any]: | |
| result = _execute(sql) | |
| features = result.get("features", []) | |
| values = [ | |
| f["properties"].get(column) for f in features | |
| if isinstance(f.get("properties", {}).get(column), (int, float)) | |
| and not isinstance(f["properties"].get(column), bool) | |
| ] | |
| if not values: | |
| raise ToolError(f"No numeric values found in column '{column}'.") | |
| ctx.sql_statements.append(sql.strip()) | |
| values.sort() | |
| n = len(values) | |
| stats = { | |
| "count": n, | |
| "min": values[0], | |
| "max": values[-1], | |
| "mean": sum(values) / n, | |
| "median": values[n // 2] if n % 2 else (values[n // 2 - 1] + values[n // 2]) / 2, | |
| } | |
| ctx.stats[column] = stats | |
| return stats | |
| def ask_user(question: str, options: Optional[List[str]] = None) -> Dict[str, Any]: | |
| """Record a clarifying question; the loop stops and surfaces it.""" | |
| ctx.pending_question = {"question": question, "options": options or []} | |
| return {"asked": question, "note": "Stop and wait for the user's answer."} | |
| # ------------------------------------------------------------------ schemas | |
| def _obj(props: Dict[str, Any], required: List[str]) -> Dict[str, Any]: | |
| return {"type": "object", "properties": props, "required": required} | |
| _str = {"type": "string"} | |
| _int = {"type": "integer"} | |
| tools = [ | |
| Tool( | |
| name="search_datasets", | |
| description=( | |
| "Find datasets relevant to a question by meaning. Use this first when you " | |
| "do not already know which table holds the answer." | |
| ), | |
| parameters=_obj({ | |
| "query": {**_str, "description": "What you are looking for, in plain language."}, | |
| "limit": {**_int, "description": "Max results (default 8)."}, | |
| }, ["query"]), | |
| run=search_datasets, | |
| ), | |
| Tool( | |
| name="describe_table", | |
| description=( | |
| "Get a table's exact columns, types, row count, whether it has geometry, " | |
| "and its documented meaning. Call before writing SQL against an unfamiliar table." | |
| ), | |
| parameters=_obj({"table": _str}, ["table"]), | |
| run=describe_table_tool, | |
| ), | |
| Tool( | |
| name="sample_values", | |
| description=( | |
| "List real distinct values of a column (and min/max if numeric). Use this " | |
| "BEFORE filtering on a value whose exact format you are unsure of β codes, " | |
| "category names, dates. Guessing a format returns zero rows with no error." | |
| ), | |
| parameters=_obj({ | |
| "table": _str, | |
| "column": _str, | |
| "limit": {**_int, "description": "Max distinct values (default 15)."}, | |
| }, ["table", "column"]), | |
| run=sample_values, | |
| ), | |
| Tool( | |
| name="run_sql", | |
| description=( | |
| "Execute a read-only DuckDB SQL query and get back the row count plus a small " | |
| "preview. Use it to check an approach or compute an answer. It does NOT put " | |
| "anything on the map β use add_map_layer for that." | |
| ), | |
| parameters=_obj({ | |
| "sql": {**_str, "description": "A SELECT or WITH query."}, | |
| "purpose": {**_str, "description": "One short line on what this is for."}, | |
| }, ["sql"]), | |
| run=run_sql, | |
| ), | |
| Tool( | |
| name="add_map_layer", | |
| description=( | |
| "Run a query and add its result to the map as a styled layer. The query MUST " | |
| "select a geometry column. For a temporal animation, return every time step " | |
| "with its date/step column rather than filtering to one." | |
| ), | |
| parameters=_obj({ | |
| "sql": {**_str, "description": "A SELECT that includes a geometry column."}, | |
| "name": {**_str, "description": "Short layer name, 1-4 words."}, | |
| "color_by": {**_str, "description": "Column to colour by (optional)."}, | |
| }, ["sql", "name"]), | |
| run=add_map_layer, | |
| produces_output=True, | |
| ), | |
| Tool( | |
| name="make_chart", | |
| description=( | |
| "Run a query and turn the result into a chart for the Plots tab. Use for " | |
| "rankings, comparisons and distributions." | |
| ), | |
| parameters=_obj({ | |
| "sql": _str, | |
| "chart_type": {**_str, "enum": ["bar", "line", "pie", "histogram"]}, | |
| "title": _str, | |
| }, ["sql", "chart_type", "title"]), | |
| run=make_chart, | |
| produces_output=True, | |
| ), | |
| Tool( | |
| name="compute_stats", | |
| description="Summary statistics (count, min, max, mean, median) for a numeric column.", | |
| parameters=_obj({"sql": _str, "column": _str}, ["sql", "column"]), | |
| run=compute_stats, | |
| ), | |
| Tool( | |
| name="ask_user", | |
| description=( | |
| "Ask ONE clarifying question when the request is genuinely ambiguous and " | |
| "guessing would waste the user's time. Prefer making a reasonable choice and " | |
| "saying what you chose." | |
| ), | |
| parameters=_obj({ | |
| "question": _str, | |
| "options": {"type": "array", "items": _str}, | |
| }, ["question"]), | |
| run=ask_user, | |
| ), | |
| ] | |
| return {t.name: t for t in tools} | |
| async def call_tool(tool: Tool, args: Dict[str, Any]) -> Dict[str, Any]: | |
| """ | |
| Run a tool off the event loop, returning either its result or a structured | |
| error. Errors are returned rather than raised so the agent can read what went | |
| wrong and try something else β an exception here would end the whole turn. | |
| """ | |
| try: | |
| result = await asyncio.wait_for( | |
| asyncio.to_thread(tool.run, **args), timeout=TOOL_TIMEOUT_SECONDS | |
| ) | |
| return {"ok": True, "result": result} | |
| except asyncio.TimeoutError: | |
| return {"ok": False, "error": | |
| f"{tool.name} timed out after {TOOL_TIMEOUT_SECONDS:.0f}s. " | |
| "Narrow the query (filter by region, season or species) and retry."} | |
| except ToolError as e: | |
| return {"ok": False, "error": str(e)} | |
| except TypeError as e: | |
| return {"ok": False, "error": f"Bad arguments for {tool.name}: {e}"} | |
| except Exception as e: # noqa: BLE001 - surface anything to the agent | |
| logger.warning(f"Tool {tool.name} failed: {e}", exc_info=True) | |
| return {"ok": False, "error": f"{tool.name} failed: {e}"} | |
| def serialize_result(payload: Dict[str, Any]) -> Dict[str, Any]: | |
| """ | |
| Make a tool result safe to hand back to the model. | |
| Note the previous version checked with a bare `json.dumps`, which does NOT | |
| raise on NaN β it emits a bare `NaN` token that the API then rejects with | |
| `400 INVALID_ARGUMENT`. json_safe strips those first. | |
| """ | |
| safe = json_safe(payload) | |
| try: | |
| json.dumps(safe, allow_nan=False) | |
| return safe | |
| except (TypeError, ValueError): | |
| return json.loads(dumps_safe(safe)) | |