Use recall tool in chat agent

CHANGELOG.md

@@ -9,6 +9,17 @@
   - Supports `all` argument to download/upload all datasets at once
   - Use `--force` flag to overwrite existing databases
   - Avoids lengthy database rebuild times for users running benchmarks
+- **Stable Citation Registry**: Citation indices now persist across tool calls within a session
+  - Same `chunk_id` always returns the same citation index (first-occurrence-wins)
+  - New `citation_registry: dict[str, int]` field on `ChatSessionState`
+  - New `get_or_assign_index(chunk_id)` method for stable index assignment
+  - Registry serialized/restored via AG-UI state protocol
+- **Recall Tool**: Check conversation history before running research
+  - New `recall` tool on chat agent searches previous Q&A pairs by semantic similarity
+  - Uses embedding similarity matching with 0.8 cosine similarity threshold
+  - Returns previous answer with citations if found, avoiding redundant research calls
+  - Emits `StateSnapshotEvent` so frontend can display recalled citations
+  - Updated system prompt with routing guidance: use `recall` FIRST for follow-up questions
 - **Dynamic Session Context**: Compressed conversation history for multi-turn chat
   - New `SessionContext` model stores summarized conversation state instead of raw Q&A history
   - Background LLM-based summarization runs after each `ask` tool call (non-blocking)

docs/agents.md

@@ -61,12 +61,15 @@ Key features:
 
 ### Tools
 
-The chat agent uses three tools:
+The chat agent uses four tools:
 
+- `recall` — Search conversation history for previous answers (use FIRST for follow-up questions)
 - `search` — Hybrid search with optional document filter
 - `ask` — Answer questions using the conversational research graph
 - `get_document` — Retrieve a specific document by title or URI
 
+The `recall` tool uses embedding similarity to find semantically matching questions from conversation history. If a match is found (above 0.8 cosine similarity threshold), it returns the previous answer with citations, avoiding redundant research calls.
+
 ### CLI Usage
 
 ```bash
@@ -104,11 +107,24 @@ The `ChatSessionState` maintains:
 - `session_id` — Unique identifier for the session
 - `qa_history` — List of previous Q/A pairs (FIFO, max 50)
 - `session_context` — Automatically maintained session context summary
+- `document_filter` — List of document titles/URIs to restrict searches
+- `citation_registry` — Stable mapping of chunk IDs to citation indices
+
+**Citation Registry**: Citation indices persist across tool calls within a session. The same `chunk_id` always returns the same citation index (first-occurrence-wins). This ensures consistent citation numbering in multi-turn conversations — `[1]` always refers to the same source.
+
+```python
+# Example: citation indices are stable across calls
+state = ChatSessionState()
+
+# First call returns citations [1], [2], [3]
+# Second call reuses [1] if same chunk, assigns [4], [5] for new chunks
+# User can reference [1] in follow-up and it still refers to original source
+```
 
 Q/A history is used to:
 
 1. Provide context for follow-up questions
-2. Avoid repeating previous answers
+2. Avoid repeating previous answers via the `recall` tool
 3. Enable semantic ranking of relevant past answers
 
 ### AG-UI Integration
@@ -135,7 +151,9 @@ The emitted state structure:
   "haiku.rag.chat": {
     "session_id": "",
     "citations": [...],
-    "qa_history": [...]
+    "qa_history": [...],
+    "document_filter": [...],
+    "citation_registry": {"chunk-id-1": 1, "chunk-id-2": 2}
   }
 }
 ```

haiku_rag_slim/haiku/rag/agents/chat/agent.py

@@ -385,7 +385,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
     async def recall(
         ctx: RunContext[ChatDeps],
         topic: str,
-    ) -> str:
+    ) -> ToolReturn:
         """Search conversation history for a previous answer on this topic.
 
         Use this FIRST when the user asks about something that may have been
@@ -396,11 +396,11 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             topic: The topic or question to search for in conversation history
         """
         if ctx.deps.session_state is None:
-            return "No conversation history available."
+            return ToolReturn(return_value="No conversation history available.")
 
         qa_history = ctx.deps.session_state.qa_history
         if not qa_history:
-            return "No previous answers found."
+            return ToolReturn(return_value="No previous answers found.")
 
         # Get embedder and embed the topic
         embedder = get_embedder(ctx.deps.config)
@@ -421,9 +421,9 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
 
         # Check if similarity exceeds threshold
         if best_similarity < RECALL_SIMILARITY_THRESHOLD:
-            return "No previous answer found on this topic."
+            return ToolReturn(return_value="No previous answer found on this topic.")
 
-        # Return the matching answer with citations
+        # Build result with the matching answer
         matched_qa = qa_history[best_match_idx]
         result = f"**Previous answer found** (similarity: {best_similarity:.2f}):\n\n"
         result += f"**Question:** {matched_qa.question}\n\n"
@@ -433,6 +433,31 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             citation_refs = " ".join(f"[{c.index}]" for c in matched_qa.citations)
             result += f"Sources: {citation_refs}"
 
-        return result
+        # Emit state with citations so frontend can display them
+        session_id = ctx.deps.session_state.session_id
+        new_state = ChatSessionState(
+            session_id=session_id,
+            citations=matched_qa.citations,
+            qa_history=ctx.deps.session_state.qa_history,
+            session_context=get_cached_session_context(session_id)
+            if session_id
+            else None,
+            document_filter=ctx.deps.session_state.document_filter,
+            citation_registry=ctx.deps.session_state.citation_registry,
+        )
+
+        snapshot = new_state.model_dump()
+        if ctx.deps.state_key:
+            snapshot = {ctx.deps.state_key: snapshot}
+
+        return ToolReturn(
+            return_value=result,
+            metadata=[
+                StateSnapshotEvent(
+                    type=EventType.STATE_SNAPSHOT,
+                    snapshot=snapshot,
+                )
+            ],
+        )
 
     return agent