diff --git a/backend/chat/agent.py b/backend/chat/agent.py index 70690790..b03595fc 100644 --- a/backend/chat/agent.py +++ b/backend/chat/agent.py @@ -20,7 +20,9 @@ TextPart, TextPartDelta, Tool, + ToolCallPart, ToolReturn, + ToolReturnPart, UsageLimitExceeded, UsageLimits, UserPromptPart, @@ -82,6 +84,12 @@ - To chart a ranking, reuse the rank_rows reference (with sort_by and top) as the chart's data. +Follow-up questions: earlier answers keep the tool calls they made but not +the results. Reuse their dataset ids, location ids, filters, and chart +references instead of searching again, and call a tool again when you need +its rows. To change a chart the user has seen, call render_chart again with +the earlier spec adjusted. + Answers: - Lead with a one-sentence answer, then only the key numbers with their year and units. Aim for under 150 words; the reader can ask a follow-up. @@ -285,16 +293,61 @@ def turn_cost(messages: list[ModelMessage]) -> float | None: return round(sum(costs), 6) if costs else None -def to_history(messages: list[dict[str, str]]) -> list[ModelMessage]: - """Rebuild text-only history from the client; tool calls aren't replayed.""" +# Every tool the agent can call, so history can't name anything else. +AGENT_TOOLS = frozenset({*AGENT_DATA_TOOLS, "rank_rows", "render_chart"}) + +# What an earlier tool call "returned" when it is replayed from history. Rows +# are never kept between turns: the client would have to hold and re-send +# them, and the warehouse can answer the same call again. +HISTORY_RESULT = "Result not kept from the earlier turn. Call again if you need it." +HISTORY_CHART = "The chart was shown to the user." + + +def to_history(messages: list[dict[str, Any]]) -> list[ModelMessage]: + """Rebuild history from the client's messages. + + An assistant message may carry the tool calls its answer made (name and + arguments). They are replayed as real tool calls with a placeholder result, + so a follow-up ("how does Rutland compare?", "make that a bar chart") can + reuse the dataset, places, and chart spec without discovering them again. + """ history: list[ModelMessage] = [] - for message in messages: + for index, message in enumerate(messages): if message["role"] == "user": history.append( ModelRequest(parts=[UserPromptPart(content=message["content"])]) ) - else: - history.append(ModelResponse(parts=[TextPart(content=message["content"])])) + continue + tools = message.get("tools") or [] + if tools: + ids = [f"history_{index}_{n}" for n in range(len(tools))] + history.append( + ModelResponse( + parts=[ + ToolCallPart( + tool_name=tool["name"], + args=tool["args"], + tool_call_id=call_id, + ) + for tool, call_id in zip(tools, ids, strict=True) + ] + ) + ) + history.append( + ModelRequest( + parts=[ + ToolReturnPart( + tool_name=tool["name"], + content=HISTORY_CHART + if tool["name"] == "render_chart" + else HISTORY_RESULT, + tool_call_id=call_id, + ) + for tool, call_id in zip(tools, ids, strict=True) + ] + ) + ) + history.append(ModelResponse(parts=[TextPart(content=message["content"])])) return history diff --git a/backend/chat/config.py b/backend/chat/config.py index 04c908cc..9ae6f26b 100644 --- a/backend/chat/config.py +++ b/backend/chat/config.py @@ -44,6 +44,10 @@ class ChatSettings: max_history_messages: int = 20 max_message_chars: int = 4_000 max_history_chars: int = 40_000 + # Tool calls replayed from one earlier answer, and the size of each one's + # arguments (see agent.to_history); larger or extra ones are dropped. + max_history_tool_calls: int = 12 + max_history_tool_chars: int = 2_000 # Questions per client IP per rolling hour, per worker process. rate_limit: int = 30 # Rows /api/chat/data returns for one chart reference. diff --git a/backend/chat/router.py b/backend/chat/router.py index 8374f403..73e39f1f 100644 --- a/backend/chat/router.py +++ b/backend/chat/router.py @@ -19,7 +19,7 @@ import threading import time from collections import defaultdict, deque -from typing import Literal +from typing import Any, Literal import anyio from fastapi import APIRouter, Depends, HTTPException, Request, Response @@ -27,15 +27,31 @@ from pydantic import BaseModel, Field, model_validator from chat import access -from chat.agent import build_agent, stream_answer, to_history, tool_error +from chat.agent import ( + AGENT_TOOLS, + build_agent, + stream_answer, + to_history, + tool_error, +) from chat.charts import DataRef, resolve_rows from chat.config import ChatSettings from data_tools.service import DataToolService +class ToolUse(BaseModel): + """A tool call an earlier answer made, sent back so follow-ups can reuse + its dataset, places, and filters (agent.to_history).""" + + name: str = Field(max_length=50) + args: dict[str, Any] = Field(default_factory=dict) + + class ChatMessage(BaseModel): role: Literal["user", "assistant"] content: str = Field(min_length=1) + # Assistant messages only; ignored on user messages. + tools: list[ToolUse] = Field(default_factory=list, max_length=100) class ChatRequest(BaseModel): @@ -77,23 +93,43 @@ def allow(self, client: str) -> bool: return True +def replayable_tools( + message: ChatMessage, settings: ChatSettings +) -> list[dict[str, Any]]: + """The message's tool calls that are safe and small enough to replay. + + The client sends these, so they are held to the agent's own tool names and + a size bound. Dropping one only costs the model a rediscovery. + """ + if message.role != "assistant": + return [] + tools = [ + {"name": tool.name, "args": tool.args} + for tool in message.tools + if tool.name in AGENT_TOOLS + and len(json.dumps(tool.args, default=str)) <= settings.max_history_tool_chars + ] + return tools[-settings.max_history_tool_calls :] + + def bounded_history( messages: list[ChatMessage], settings: ChatSettings -) -> tuple[str, list[dict[str, str]]]: +) -> tuple[str, list[dict[str, Any]]]: """Split off the question and keep the newest history that fits the bounds.""" question = messages[-1].content if len(question) > settings.max_message_chars: raise HTTPException( 413, f"Questions are limited to {settings.max_message_chars} characters" ) - kept: list[dict[str, str]] = [] + kept: list[dict[str, Any]] = [] total = len(question) for message in reversed(messages[:-1][-settings.max_history_messages :]): content = message.content[: settings.max_message_chars] - total += len(content) + tools = replayable_tools(message, settings) + total += len(content) + len(json.dumps(tools, default=str)) * bool(tools) if total > settings.max_history_chars: break - kept.append({"role": message.role, "content": content}) + kept.append({"role": message.role, "content": content, "tools": tools}) kept.reverse() # History must alternate starting from a user turn. while kept and kept[0]["role"] != "user": diff --git a/backend/evals/chat/runner.py b/backend/evals/chat/runner.py index d60ecf00..c4720a64 100644 --- a/backend/evals/chat/runner.py +++ b/backend/evals/chat/runner.py @@ -40,7 +40,7 @@ async def play_chat( A turn that runs past `turn_timeout` seconds (a stalled provider, say) ends with an error instead of holding up the whole run. """ - history: list[dict[str, str]] = [] + history: list[dict[str, Any]] = [] turns = [] for question in case.turns: started = time.monotonic() @@ -60,6 +60,7 @@ async def play_chat( elif kind == "tool_call": steps[event["id"]] = { "name": event["name"], + "args": event["args"], "ok": True, "error": None, } @@ -87,11 +88,19 @@ async def play_chat( } ) if answer and not error: - # Same text-only history the browser sends (chatStore.toHistory). - notes = [f"[Chart shown: {c['title']}]" for c in charts] + # The same history the browser sends (chatStore.toHistory): the + # answer plus the tool calls that worked. history += [ {"role": "user", "content": question}, - {"role": "assistant", "content": "\n".join([*notes, answer])}, + { + "role": "assistant", + "content": answer, + "tools": [ + {"name": step["name"], "args": step["args"]} + for step in steps.values() + if step["ok"] + ], + }, ] return turns diff --git a/backend/tests/test_chat.py b/backend/tests/test_chat.py index 19765daf..9a91d64e 100644 --- a/backend/tests/test_chat.py +++ b/backend/tests/test_chat.py @@ -8,16 +8,29 @@ from fastapi import FastAPI from fastapi.testclient import TestClient from pydantic import ValidationError -from pydantic_ai import ToolReturn +from pydantic_ai import ModelRequest, ModelResponse, ToolCallPart, ToolReturn from pydantic_ai.exceptions import ModelHTTPError from pydantic_ai.models.function import FunctionModel 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, + preview, + rank_tool, + stream_answer, + to_history, +) from chat.charts import ChartSpec, DataRef, overlapping_rows, resolve_rows from chat.config import ChatSettings from chat.demo import CHART -from chat.router import UNLOCK_ATTEMPTS, ChatMessage, bounded_history, create_router +from chat.router import ( + UNLOCK_ATTEMPTS, + ChatMessage, + ToolUse, + bounded_history, + create_router, +) from data_tools.service import DataToolService INCOME = "acs5_ts_household_income" @@ -186,6 +199,120 @@ def test_bounded_history_keeps_newest_turns_starting_with_user(): assert history[0]["content"] == "user 4" +SERIES_CALL = {"name": "get_timeseries", "args": {"dataset_id": INCOME}} + + +def test_history_replays_earlier_tool_calls_without_their_results(): + history = to_history( + [ + {"role": "user", "content": "Income in Addison?"}, + { + "role": "assistant", + "content": "It rose.", + "tools": [SERIES_CALL, {"name": "render_chart", "args": CHART}], + }, + ] + ) + _question, calls, returns, answer = history + assert [type(m) for m in history] == [ + ModelRequest, + ModelResponse, + ModelRequest, + ModelResponse, + ] + assert [(p.tool_name, p.args) for p in calls.parts] == [ + ("get_timeseries", {"dataset_id": INCOME}), + ("render_chart", CHART), + ] + # Each call is answered by id, with a placeholder instead of rows. + assert [p.tool_call_id for p in returns.parts] == [ + p.tool_call_id for p in calls.parts + ] + assert "not kept" in returns.parts[0].content + assert "shown to the user" in returns.parts[1].content + assert answer.parts[0].content == "It rose." + # An answer with no tool calls stays a plain text message. + assert len(to_history([{"role": "assistant", "content": "Hi."}])) == 1 + + +def test_bounded_history_keeps_only_replayable_tool_calls(): + settings = ChatSettings(max_history_tool_calls=2, max_history_tool_chars=100) + known = ToolUse(**SERIES_CALL) + messages = [ + ChatMessage(role="user", content="q1", tools=[known]), + ChatMessage( + role="assistant", + content="a1", + tools=[ + ToolUse(name="search_locations", args={"query": "Addison"}), + ToolUse(name="run_sql", args={"sql": "DROP TABLE x"}), + ToolUse(name="query_data", args={"dataset_id": "x" * 200}), + known, + known, + ], + ), + ChatMessage(role="user", content="q2"), + ] + _question, history = bounded_history(messages, settings) + # Tools on a user message are ignored; unknown names and oversized + # arguments are dropped; only the newest calls are kept. + assert history[0]["tools"] == [] + assert history[1]["tools"] == [SERIES_CALL, SERIES_CALL] + + +def test_follow_up_question_sees_the_earlier_tool_calls(warehouse): + seen = [] + + async def answers(messages, _info): + seen.extend(messages) + yield "Same dataset." + + settings = ChatSettings(model="demo", warehouse_path=warehouse) + client = make_client(warehouse) + agent_model = FunctionModel(stream_function=answers) + app_agent = build_agent(settings, DataToolService(warehouse)) + + async def run(): + history = to_history( + [ + {"role": "user", "content": "Income in Addison?"}, + {"role": "assistant", "content": "It rose.", "tools": [SERIES_CALL]}, + ] + ) + with app_agent.override(model=agent_model): + return [ + e + async for e in stream_answer( + app_agent, settings, "And Chittenden?", history + ) + ] + + events = asyncio.run(run()) + assert events[-1]["type"] == "done" + replayed = [ + part + for message in seen + for part in message.parts + if isinstance(part, ToolCallPart) + ] + assert [(p.tool_name, p.args) for p in replayed] == [ + ("get_timeseries", {"dataset_id": INCOME}) + ] + # The route accepts the same shape from the browser. + response = client.post( + "/api/chat", + json={ + "messages": [ + {"role": "user", "content": "Income in Addison?"}, + {"role": "assistant", "content": "It rose.", "tools": [SERIES_CALL]}, + {"role": "user", "content": "And Chittenden?"}, + ] + }, + ) + assert response.status_code == 200 + assert '"type": "done"' in response.text + + def test_status_route(warehouse): client = make_client(warehouse) response = client.get("/api/chat/status") diff --git a/docs/architecture/data-chat.md b/docs/architecture/data-chat.md index 680a77f7..1ef80010 100644 --- a/docs/architecture/data-chat.md +++ b/docs/architecture/data-chat.md @@ -145,6 +145,12 @@ That's a few dozen tokens however large the result is. 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 store. +- **Follow-ups reuse references too:** the browser keeps each answer's tool + calls (name and arguments, never results) and sends them back with the + history. The server replays them as tool calls with a placeholder result + (`to_history` in `chat/agent.py`), so "how does Rutland compare?" or "make + that a bar chart" starts from the dataset, places, and chart spec already + used. Only the agent's own tool names are accepted, within a size bound. **Chart kinds:** `line` (change over years), `bar` (compare places or categories; named categories draw as horizontal bars), `stacked_bar` and diff --git a/frontend/src/app/chat/page.tsx b/frontend/src/app/chat/page.tsx index 062f8b12..77ed3a35 100644 --- a/frontend/src/app/chat/page.tsx +++ b/frontend/src/app/chat/page.tsx @@ -266,7 +266,12 @@ export default function ChatPage() { ...t, steps: [ ...(t.steps ?? []), - { id: event.id, name: event.name, status: 'running' }, + { + id: event.id, + name: event.name, + status: 'running', + args: event.args, + }, ], })); break; diff --git a/frontend/src/components/Chat/api.ts b/frontend/src/components/Chat/api.ts index 5b353de4..052f2bd5 100644 --- a/frontend/src/components/Chat/api.ts +++ b/frontend/src/components/Chat/api.ts @@ -72,6 +72,9 @@ export type ChatEvent = export interface HistoryMessage { role: 'user' | 'assistant'; content: string; + // Tool calls an earlier answer made; the server replays them so follow-ups + // can reuse the same dataset, places, and chart. + tools?: { name: string; args: Record }[]; } const OFF: ChatStatus = { diff --git a/frontend/src/components/Chat/chatStore.ts b/frontend/src/components/Chat/chatStore.ts index f8ac9cbe..573695c3 100644 --- a/frontend/src/components/Chat/chatStore.ts +++ b/frontend/src/components/Chat/chatStore.ts @@ -13,6 +13,9 @@ export interface ToolStep { name: string; status: 'running' | 'ok' | 'error'; error?: string | null; + // Kept so follow-up questions can send the call back (see toHistory). + // Chats saved before this was recorded don't have it. + args?: Record; } export interface ChatTurn { @@ -123,19 +126,22 @@ export const useChat = create()( ), ); -/** The text history the API expects; failed or empty answers are skipped. - * Tool calls aren't replayed, so charts are noted by title to keep the - * model aware of what the user has already seen. */ +/** The history the API expects; failed or empty answers are skipped. Each + * answer carries the tool calls that worked, which the server replays so a + * follow-up can reuse their dataset, places, and chart spec. Chats saved + * before calls were recorded fall back to noting their charts by title. */ export function toHistory(turns: ChatTurn[]): HistoryMessage[] { return turns .filter((t) => t.content.trim() && !(t.role === 'assistant' && t.error)) .map((t) => { + if (t.role === 'user') return { role: t.role, content: t.content }; + const tools = (t.steps ?? []).flatMap((s) => + s.status === 'ok' && s.args ? [{ name: s.name, args: s.args }] : [], + ); + if (tools.length > 0) return { role: t.role, content: t.content, tools }; const charts = (t.charts ?? []).map( (c) => `[Chart shown: ${c.spec.title}]`, ); - return { - role: t.role, - content: [...charts, t.content].join('\n'), - }; + return { role: t.role, content: [...charts, t.content].join('\n') }; }); }