Carry each answer's tool calls into follow-up questions

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
 
 

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.

backend/chat/router.py

@@ -19,23 +19,39 @@
 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
 from fastapi.responses import StreamingResponse
 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":

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
 

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")

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