diff --git a/backend/chat/agent.py b/backend/chat/agent.py index 70690790..2ca42fe1 100644 --- a/backend/chat/agent.py +++ b/backend/chat/agent.py @@ -2,6 +2,7 @@ from __future__ import annotations +import json import logging from collections.abc import AsyncIterator from typing import Any @@ -45,14 +46,26 @@ # export_data returns inline CSV, which the browser gets by reference instead. AGENT_DATA_TOOLS = tuple(name for name in TOOL_MODELS if name != "export_data") -INSTRUCTIONS = """\ +# How the agent learns which datasets exist. Normally they are listed in the +# instructions (dataset_catalog); FIND_WITH_TOOL is the fallback when the +# warehouse couldn't be read at startup. +FIND_IN_PROMPT = """\ +- The datasets are listed above. Pick from that list, then use + describe_dataset and search_variables or search_locations before querying. + Use returned identifiers exactly.""" +FIND_WITH_TOOL = """\ +- Start with list_datasets, then describe_dataset and search_variables or + search_locations before querying. Use returned identifiers exactly.""" + +INTRO = """\ You answer questions about Vermont public data (housing, demographics, economy, health, zoning, flood hazard, wastewater) for planners, select boards, developers, and reporters. Write plainly for non-technical readers. +""" +RULES = """\ Finding data: -- Start with list_datasets, then describe_dataset and search_variables or - search_locations before querying. Use returned identifiers exactly. +{find} - Keep geography, year, measure, and units consistent when comparing places. Never average medians, interpolate missing years, or invent values. If the data can't answer the question, say so and suggest what it can answer. @@ -127,6 +140,98 @@ def preview(result: dict[str, Any], rows_shown: int) -> dict[str, Any]: } +# Short gap lists are worth showing; long ones (historic census years) aren't. +MAX_GAPS_LISTED = 10 + + +def compact_datasets(result: dict[str, Any]) -> dict[str, Any]: + """Trim list_datasets to what the model needs to pick a dataset. + + The full listing repeats the same caveats on most datasets and carries + lineage and column plumbing that describe_dataset also returns. The model + reads this result on every later request of the turn, so its size is paid + for many times over. Caveats shared by several datasets are stated once + and referenced by key; nothing a reader needs is dropped. + """ + datasets = result.get("datasets") + if not isinstance(datasets, list): + return result + counts: dict[str, int] = {} + for dataset in datasets: + for caveat in dataset.get("caveats", []): + counts[caveat] = counts.get(caveat, 0) + 1 + shared = { + caveat: f"C{i}" + for i, caveat in enumerate((c for c, n in counts.items() if n > 1), start=1) + } + compact = [] + for dataset in datasets: + caveats = dataset.get("caveats", []) + gaps = dataset.get("gaps") or [] + entry = { + "dataset_id": dataset.get("dataset_id"), + "label": dataset.get("label"), + "description": dataset.get("description"), + "source": (dataset.get("source") or {}).get("name"), + "kind": dataset.get("kind"), + # Record datasets (zoning, wastewater) have no year. + "years": [dataset["year_min"], dataset.get("year_max")] + if dataset.get("year_min") is not None + else None, + "gaps": gaps + if len(gaps) <= MAX_GAPS_LISTED + else f"{len(gaps)} missing years; describe_dataset lists years_available", + "geo_type": dataset.get("fixed_geo_type"), + "variable_columns": dataset.get("variable_columns"), + # Margin columns are covered by the shared caveat that explains them. + "value_columns": { + column: unit + for column, unit in (dataset.get("value_columns") or {}).items() + if not column.endswith("_moe") + }, + "shared_caveats": [shared[c] for c in caveats if c in shared], + "caveats": [c for c in caveats if c not in shared], + "available": None if dataset.get("available", True) else False, + } + compact.append({k: v for k, v in entry.items() if v not in (None, [], {})}) + return { + "datasets": compact, + "shared_caveats": {key: caveat for caveat, key in shared.items()}, + "warehouse_version": result.get("warehouse_version"), + } + + +def dataset_catalog(service: DataToolService) -> str | None: + """The compact dataset list as text for the instructions. + + Every question needs this list, so it goes in the prompt instead of + behind a tool call: that saves a model request per question (models + often made the call two or three times), and as part of the fixed prefix + it is read from cache by every later request and question. None if the + warehouse can't be read right now; the agent then keeps list_datasets. + """ + try: + listing = compact_datasets(service.call("list_datasets", {})) + except Exception as exc: # noqa: BLE001 -- chat still starts, with the tool + logger.error("Chat couldn't list datasets (%s)", type(exc).__name__) + return None + return json.dumps(listing, separators=(",", ":"), default=str) + + +def instructions(catalog: str | None) -> str: + """The dataset list goes before the rules, not after them. With the list + last, the rules sat ahead of 20,000 characters of JSON, and the eval + scored 86% against 89% with the rules last + (docs/reports/chat-eval-history.md, 2026-11-13).""" + if catalog is None: + return f"{INTRO}\n{RULES.format(find=FIND_WITH_TOOL)}" + return ( + f"{INTRO}\nDatasets (shared_caveats keys refer to the map of that name; " + "describe_dataset returns columns, filter values, and year coverage):\n" + f"{catalog}\n\n{RULES.format(find=FIND_IN_PROMPT)}" + ) + + def agent_schema(name: str) -> dict[str, Any]: """A data tool's schema with the filter format spelled out for models.""" schema = TOOL_MODELS[name].model_json_schema() @@ -140,6 +245,8 @@ def data_tool(service: DataToolService, name: str, rows_shown: int) -> Tool: def call(**arguments: Any) -> dict[str, Any]: try: result = service.call(name, normalize_filters(arguments)) + if name == "list_datasets": + return compact_datasets(result) return preview(result, rows_shown) except Exception as exc: # noqa: BLE001 -- the model gets a safe message return {"error": tool_error(exc, name)} @@ -245,14 +352,27 @@ def build_agent(settings: ChatSettings, service: DataToolService) -> Agent: # Ask OpenRouter to report each response's cost (see turn_cost). model_settings["openrouter_usage"] = {"include": True} model_settings["openrouter_provider"] = PROVIDER_ROUTING + # One question is several model requests, each re-sending the tool + # definitions, the instructions, and every earlier tool result. Mark + # all three as cacheable so later requests read that prefix at the + # cached rate. Providers without explicit cache control ignore these. + model_settings["openrouter_cache_tool_definitions"] = True + model_settings["openrouter_cache_instructions"] = True + model_settings["openrouter_cache_messages"] = True + # Built once: computing it per request would put it after the cached part + # of the instructions. A replaced warehouse shows here after a restart + # (describe_dataset always reads the live one). + catalog = dataset_catalog(service) tools = [ - data_tool(service, name, settings.preview_rows) for name in AGENT_DATA_TOOLS + data_tool(service, name, settings.preview_rows) + for name in AGENT_DATA_TOOLS + if catalog is None or name != "list_datasets" ] tools.append(rank_tool(service, settings)) tools.append(chart_tool(service, settings)) return Agent( model, - instructions=INSTRUCTIONS, + instructions=instructions(catalog), tools=tools, # Let the API start without a provider key; requests report the problem. defer_model_check=True, @@ -354,6 +474,9 @@ async def stream_answer( "tool_calls": usage.tool_calls, "input_tokens": usage.input_tokens, "output_tokens": usage.output_tokens, + # Part of input_tokens; shows whether caching works. + "cache_read_tokens": usage.cache_read_tokens, + "cache_write_tokens": usage.cache_write_tokens, "cost_usd": turn_cost(event.result.new_messages()), }, } diff --git a/backend/tests/test_chat.py b/backend/tests/test_chat.py index 19765daf..57ab9c2e 100644 --- a/backend/tests/test_chat.py +++ b/backend/tests/test_chat.py @@ -4,6 +4,7 @@ import json import duckdb +import httpx2 import pytest from fastapi import FastAPI from fastapi.testclient import TestClient @@ -11,9 +12,18 @@ from pydantic_ai import ToolReturn from pydantic_ai.exceptions import ModelHTTPError from pydantic_ai.models.function import FunctionModel +from pydantic_ai.models.openrouter import OpenRouterModel +from pydantic_ai.providers.openrouter import OpenRouterProvider from chat import access -from chat.agent import build_agent, chart_tool, preview, rank_tool, stream_answer +from chat.agent import ( + build_agent, + chart_tool, + compact_datasets, + preview, + rank_tool, + stream_answer, +) from chat.charts import ChartSpec, DataRef, overlapping_rows, resolve_rows from chat.config import ChatSettings from chat.demo import CHART @@ -88,6 +98,70 @@ def test_preview_trims_rows_for_the_model(): assert preview(small, 5) is small +def test_dataset_list_is_compact_and_states_shared_caveats_once(service): + full = service.call("list_datasets", {}) + compact = compact_datasets(full) + assert [d["dataset_id"] for d in compact["datasets"]] == [ + d["dataset_id"] for d in full["datasets"] + ] + assert len(json.dumps(compact)) < len(json.dumps(full, default=str)) / 2 + income = next(d for d in compact["datasets"] if d["dataset_id"] == INCOME) + # Enough to query a wide series without describing it first. + assert "median_household_income" in income["value_columns"] + assert not any(c.endswith("_moe") for c in income["value_columns"]) + assert "warehouse_table" not in income + # Every caveat survives, either on the dataset or once in the shared map. + for before, after in zip(full["datasets"], compact["datasets"], strict=True): + kept = after.get("caveats", []) + [ + compact["shared_caveats"][key] for key in after.get("shared_caveats", []) + ] + assert sorted(kept) == sorted(before["caveats"]) + assert all( + sum(c in d["caveats"] for d in full["datasets"]) > 1 + for c in compact["shared_caveats"].values() + ) + + +def seen_by_model(agent, settings): + """The instructions and tool names one request hands the model.""" + seen = {} + + async def answers(_messages, info): + seen["instructions"] = info.instructions + seen["tools"] = [tool.name for tool in info.function_tools] + yield "ok" + + async def run(): + with agent.override(model=FunctionModel(stream_function=answers)): + return [e async for e in stream_answer(agent, settings, "hi", [])] + + assert asyncio.run(run())[-1]["type"] == "done" + return seen + + +def test_datasets_are_listed_in_the_instructions_not_behind_a_tool(warehouse): + settings = ChatSettings(model="demo", warehouse_path=warehouse) + seen = seen_by_model(build_agent(settings, DataToolService(warehouse)), settings) + assert INCOME in seen["instructions"] + assert "shared_caveats" in seen["instructions"] + assert "list_datasets" not in seen["tools"] + assert "list_datasets" not in seen["instructions"] + assert "describe_dataset" in seen["tools"] + # The rules come after the list, where the model weighs them most. + assert seen["instructions"].index(INCOME) < seen["instructions"].index( + "Charts and tables:" + ) + + +def test_agent_keeps_the_dataset_tool_when_the_warehouse_cant_be_read(tmp_path): + missing = tmp_path / "absent.duckdb" + settings = ChatSettings(model="demo", warehouse_path=missing) + seen = seen_by_model(build_agent(settings, DataToolService(missing)), settings) + assert "list_datasets" in seen["tools"] + assert "Start with list_datasets" in seen["instructions"] + assert INCOME not in seen["instructions"] + + def test_data_ref_drops_paging_and_validates_tool_args(): ref = DataRef( tool="get_timeseries", @@ -424,3 +498,95 @@ def test_openrouter_models_report_cost_and_skip_data_collecting_providers(wareho agent = build_agent(settings, DataToolService(warehouse)) assert agent.model_settings["openrouter_provider"] == {"data_collection": "deny"} assert agent.model_settings["openrouter_usage"] == {"include": True} + + +def openrouter_requests(warehouse, slug): + """Play one tool-using turn against a fake OpenRouter; return the request + bodies it received and the turn's events.""" + bodies = [] + + def respond(request): + body = json.loads(request.content) + bodies.append(body) + delta = ( + { + "role": "assistant", + "tool_calls": [ + { + "index": 0, + "id": "call_1", + "type": "function", + "function": { + "name": "describe_dataset", + "arguments": json.dumps({"dataset_id": INCOME}), + }, + } + ], + } + if len(bodies) == 1 + else {"role": "assistant", "content": "Done."} + ) + chunk = { + "id": "x", + "object": "chat.completion.chunk", + "created": 1, + "model": body["model"], + "choices": [ + { + "index": 0, + "delta": delta, + "finish_reason": "tool_calls" if len(bodies) == 1 else "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 1, + "total_tokens": 11, + "prompt_tokens_details": {"cached_tokens": 7, "cache_write_tokens": 2}, + }, + } + return httpx2.Response( + 200, + content=f"data: {json.dumps(chunk)}\n\ndata: [DONE]\n\n", + headers={"content-type": "text/event-stream"}, + ) + + settings = ChatSettings(model=slug, warehouse_path=warehouse) + agent = build_agent(settings, DataToolService(warehouse)) + + async def run(): + client = httpx2.AsyncClient(transport=httpx2.MockTransport(respond)) + model = OpenRouterModel( + slug, provider=OpenRouterProvider(api_key="test", http_client=client) + ) + with agent.override(model=model): + return [e async for e in stream_answer(agent, settings, "hi", [])] + + return bodies, asyncio.run(run()) + + +def cache_marked(part) -> bool: + return "cache_control" in json.dumps(part) + + +def test_claude_requests_mark_the_repeated_prefix_as_cacheable(warehouse): + bodies, events = openrouter_requests(warehouse, "anthropic/claude-haiku-5.5") + assert len(bodies) == 2 + for body in bodies: + # The last tool definition, the instructions, and the newest message: + # everything before each marker is read from cache on the next request. + assert [cache_marked(tool) for tool in body["tools"]][-1] + assert cache_marked(body["messages"][0]) + assert cache_marked(body["messages"][-1]) + # The dataset list rides in the cached instructions. + assert INCOME in json.dumps(body["messages"][0]) + assert bodies[1]["messages"][-1]["role"] == "tool" + usage = events[-1]["usage"] + assert usage["cache_read_tokens"] == 14 + assert usage["cache_write_tokens"] == 4 + + +def test_other_models_get_no_cache_markers(warehouse): + bodies, events = openrouter_requests(warehouse, "openai/gpt-6.1-sol") + assert events[-1]["type"] == "done" + assert not any(cache_marked(body) for body in bodies) diff --git a/docs/architecture/data-chat.md b/docs/architecture/data-chat.md index 680a77f7..4cff4b2f 100644 --- a/docs/architecture/data-chat.md +++ b/docs/architecture/data-chat.md @@ -141,6 +141,15 @@ That's a few dozen tokens however large the result is. plus counts) before they reach the model. The model reads enough to write the narrative; full rows go to the browser by reference. This is also the biggest lever on token cost, since tool results are re-sent on every loop iteration. +- **The dataset list is in the instructions,** not behind a tool + (`dataset_catalog` in `chat/agent.py`). Every question needs it, so listing + it up front saves a model request, and as part of the fixed prefix it is + cached (see Decision 4). It goes ahead of the rules, which the model follows + better when they come last. It is a compact listing: caveats shared by several + datasets are stated once, and lineage and column plumbing are left to + `describe_dataset`. It is built when the API starts, so restart the API after + replacing the warehouse; if the warehouse can't be read then, the agent keeps + a `list_datasets` tool instead. The MCP's `list_datasets` is unchanged. - **Fits the existing types:** `ItemDataRef` in `frontend/src/types/cachedCharts.ts` describes exactly this, though nothing uses it yet. Agent charts become ordinary `ChartItem`s with a `dataRef`, so "Add to working report" works with the current @@ -223,6 +232,17 @@ against live warehouse data, with `data_collection: "deny"` routing. hosts, many running reduced-precision (fp4) builds, so quality can vary by host. +These runs predate prompt caching and the dataset list in the instructions. +Claude requests now mark the tool definitions, the instructions (with the +dataset list), and the newest message as cacheable (`build_agent` in +`chat/agent.py`). The fixed prefix is about 42,000 characters and is shared by +every question while the cache is warm; within a turn, each later request also +reads the earlier tool results at the cached rate. The turn's `usage` reports +`cache_read_tokens` and `cache_write_tokens`. Other providers ignore the +markers. Measured on 2026-11-13 against `main` the same day: $0.0039 per chat +instead of $0.0146, scoring 89% against 91% +([reports/chat-eval-history.md](../reports/chat-eval-history.md)). + Costs are per eval chat. Real questions may be longer or chattier, so treat the last column as an upper bound on volume until we log real usage. The earlier planning estimate (about $0.08 per question for Sol) was roughly double Sol's diff --git a/docs/reports/chat-eval-history.md b/docs/reports/chat-eval-history.md index 2da34504..bcd35976 100644 --- a/docs/reports/chat-eval-history.md +++ b/docs/reports/chat-eval-history.md @@ -17,6 +17,47 @@ Case columns: **Trend** = income-trend, **Rank** = renter-burden-ranking, **Compare** = home-values-follow-up, **Mix** = jobs-mix, **Limits** = forecast-limits. +## 2026-11-13: dataset list in the prompt with caching, Haiku 5.5, 3 repeats + +- **Setup:** Claude Haiku 5.5 only, 3 repeats, 15 chats per run, all on the + warehouse rebuilt with margins of error on 2026-11-09. Three runs: `main` + as a same-day baseline (`e58f7ab8`, run `20261113-135232`); the compact + dataset list at the end of the instructions with prompt caching + (`88413638`, run `20261112-230106`); and the same with the list ahead of + the rules (`a5b7a259`, run `20261112-234945`). +- **Decision:** ship the list-first version. Quality is level with `main` + within run-to-run variation, at about a quarter of the cost. + +| Version | Score | Quality | Cost / chat | Time / chat | Tokens / chat | Errors | +| --------------------- | ----- | ------- | ----------- | ----------- | ------------- | ------ | +| `main` (list by tool) | 91% | 7.5/10 | $0.0146 | 16s | 135,528 | 0 | +| List in prompt, last | 86% | 7.1/10 | $0.0041 | 14s | 125,697 | 0 | +| List in prompt, first | 89% | 7.4/10 | $0.0039 | 13s | 122,037 | 0 | + +| Version | Trend | Rank | Compare | Mix | Limits | +| --------------------- | ----- | ---- | ------- | --- | ------ | +| `main` (list by tool) | 93% | 97% | 82% | 87% | 96% | +| List in prompt, last | 96% | 90% | 78% | 87% | 82% | +| List in prompt, first | 100% | 83% | 89% | 93% | 82% | + +- **Cost of the runs:** $0.34 in chat calls ($0.22 of it the baseline) plus + $1.77 for the judge. +- **Notes:** + - With caching, 72% to 94% of each chat's input tokens were read from + cache. Tokens per chat fell less than cost because cached tokens still + count as input. + - Charts were drawn in 6 of the 9 chats that expect one on `main` and with + the list first (5 of 9 with the list last). The October run drew 9 of 9, + so this changed with the warehouse, not with the prompt. + - Limits is lower with the list in the prompt in both runs: the agent + answers Burlington's population from ACS estimates, where on `main` it + uses the decennial census series the reference is built from. A possible + cause, not tested: the compact list describes the historic population + series as having "206 missing years". + - Rank drew no chart in the list-first run (3 of 3 on `main`): the agent + answered from the query preview without `rank_rows` or `render_chart`. + The rankings themselves were right. Cause not established. + ## 2026-10-16: Haiku 5.5 vs Sol vs DeepSeek, 3 repeats (decision run) - **Setup:** 3 repeats per model, 45 chats. OpenRouter routing with