Recall tool in chat agent

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

@@ -1,4 +1,5 @@
 import asyncio
+import math
 
 from ag_ui.core import EventType, StateSnapshotEvent
 from pydantic_ai import Agent, RunContext, ToolReturn
@@ -23,8 +24,23 @@ from haiku.rag.agents.research.graph import build_conversational_graph
 from haiku.rag.agents.research.models import Citation
 from haiku.rag.agents.research.state import ResearchDeps, ResearchState
 from haiku.rag.config.models import AppConfig
+from haiku.rag.embeddings import get_embedder
 from haiku.rag.utils import get_model
 
+# Similarity threshold for recall matching
+RECALL_SIMILARITY_THRESHOLD = 0.8
+
+
+def _cosine_similarity(vec1: list[float], vec2: list[float]) -> float:
+    """Compute cosine similarity between two vectors."""
+    dot_product = sum(a * b for a, b in zip(vec1, vec2))
+    norm1 = math.sqrt(sum(a * a for a in vec1))
+    norm2 = math.sqrt(sum(b * b for b in vec2))
+    if norm1 == 0 or norm2 == 0:
+        return 0.0
+    return dot_product / (norm1 * norm2)
+
+
 # Track summarization tasks per session to allow cancellation
 _summarization_tasks: dict[str, asyncio.Task[None]] = {}
 
@@ -365,4 +381,58 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             f"**Content:**\n{doc.content}"
         )
 
+    @agent.tool
+    async def recall(
+        ctx: RunContext[ChatDeps],
+        topic: str,
+    ) -> str:
+        """Search conversation history for a previous answer on this topic.
+
+        Use this FIRST when the user asks about something that may have been
+        discussed before. Returns the previous answer with citations if found,
+        or indicates no match exists.
+
+        Args:
+            topic: The topic or question to search for in conversation history
+        """
+        if ctx.deps.session_state is None:
+            return "No conversation history available."
+
+        qa_history = ctx.deps.session_state.qa_history
+        if not qa_history:
+            return "No previous answers found."
+
+        # Get embedder and embed the topic
+        embedder = get_embedder(ctx.deps.config)
+        topic_embedding = await embedder.embed_query(topic)
+
+        # Embed all previous questions
+        questions = [qa.question for qa in qa_history]
+        question_embeddings = await embedder.embed_documents(questions)
+
+        # Find best match by cosine similarity
+        best_match_idx = -1
+        best_similarity = 0.0
+        for i, q_embedding in enumerate(question_embeddings):
+            similarity = _cosine_similarity(topic_embedding, q_embedding)
+            if similarity > best_similarity:
+                best_similarity = similarity
+                best_match_idx = i
+
+        # Check if similarity exceeds threshold
+        if best_similarity < RECALL_SIMILARITY_THRESHOLD:
+            return "No previous answer found on this topic."
+
+        # Return the matching answer with citations
+        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"
+        result += f"**Answer:** {matched_qa.answer}\n\n"
+
+        if matched_qa.citations:
+            citation_refs = " ".join(f"[{c.index}]" for c in matched_qa.citations)
+            result += f"Sources: {citation_refs}"
+
+        return result
+
     return agent

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

@@ -4,14 +4,16 @@ You have access to a knowledge base of documents. Use your tools to search and a
 
 CRITICAL RULES:
 1. For greetings or casual chat: respond directly WITHOUT using any tools
-2. For questions: Use the "ask" tool EXACTLY ONCE - it handles query expansion internally
-3. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally
-4. NEVER call the same tool multiple times for a single user message
-5. NEVER make up information - always use tools to get facts from the knowledge base
+2. For follow-up questions about topics already discussed: Use "recall" FIRST to check conversation history
+3. For new questions: Use the "ask" tool EXACTLY ONCE - it handles query expansion internally
+4. For searches: Use the "search" tool EXACTLY ONCE - it handles multi-query expansion internally
+5. NEVER call the same tool multiple times for a single user message
+6. NEVER make up information - always use tools to get facts from the knowledge base
 
 How to decide which tool to use:
+- "recall" - Use FIRST when the user asks about a topic that may have been discussed before (e.g., "remind me about X", "what did you say about Y", "tell me again about Z"). If recall finds a previous answer, use it directly. If recall returns "no match", proceed with "ask".
 - "get_document" - Use when the user references a SPECIFIC document by name, title, or URI (e.g., "summarize document X", "get the paper about Y", "fetch 2412.00566"). Retrieves the full document content.
-- "ask" - Use for general questions about topics in the knowledge base when no specific document is named. It searches across all documents and returns answers with citations.
+- "ask" - Use for NEW questions about topics in the knowledge base when no specific document is named. It searches across all documents and returns answers with citations.
 - "search" - Use when the user explicitly asks to search/find/explore documents. Call it ONCE. After calling search, copy the ENTIRE tool response to your output INCLUDING the content snippets. Do NOT shorten, summarize, or omit any part of the results.
 
 IMPORTANT - When user mentions a document in search/ask:

tests/agents/chat/test_chat_agent.py

@@ -774,3 +774,91 @@ def test_search_tool_citation_registry_logic():
         "chunk-c": 3,
         "chunk-d": 4,
     }
+
+
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_recall_tool_finds_previous_answer(allow_model_requests, temp_db_path):
+    """Test recall tool finds a previous answer on similar topic."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        agent = create_chat_agent(Config)
+
+        # Pre-populate qa_history with a previous answer
+        session_state = ChatSessionState(
+            session_id="test-recall",
+            qa_history=[
+                QAResponse(
+                    question="What are the class labels in DocLayNet?",
+                    answer="DocLayNet defines 11 class labels including Caption, Footnote, Formula, List-item, Page-footer, Page-header, Picture, Section-header, Table, Text, and Title.",
+                    confidence=0.9,
+                    citations=[
+                        Citation(
+                            index=1,
+                            document_id="doc-1",
+                            chunk_id="chunk-1",
+                            document_uri="doclaynet.md",
+                            document_title="DocLayNet",
+                            content="DocLayNet class labels...",
+                        )
+                    ],
+                ),
+            ],
+        )
+        deps = ChatDeps(
+            client=client,
+            config=Config,
+            session_state=session_state,
+        )
+
+        # Ask about a similar topic - should find the previous answer
+        result = await agent.run(
+            "Remind me about the DocLayNet class labels",
+            deps=deps,
+        )
+
+        # The agent should use recall and find the previous answer
+        assert result.output is not None
+        # Should mention the class labels from the cached answer
+        assert "11" in result.output or "class" in result.output.lower()
+
+
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_recall_tool_no_match(allow_model_requests, temp_db_path):
+    """Test recall tool returns no match for unrelated topic."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        agent = create_chat_agent(Config)
+
+        # Pre-populate qa_history with an unrelated answer
+        session_state = ChatSessionState(
+            session_id="test-recall-nomatch",
+            qa_history=[
+                QAResponse(
+                    question="What is the capital of France?",
+                    answer="The capital of France is Paris.",
+                    confidence=0.9,
+                    citations=[],
+                ),
+            ],
+        )
+        deps = ChatDeps(
+            client=client,
+            config=Config,
+            session_state=session_state,
+        )
+
+        # Add a document so ask tool can find something
+        await client.create_document(
+            content=DOCLAYNET_CLASS_LABELS,
+            uri="doclaynet-labels",
+            title="DocLayNet Class Labels",
+        )
+
+        # Ask about an unrelated topic - recall should not match
+        result = await agent.run(
+            "What are the class labels in DocLayNet?",
+            deps=deps,
+        )
+
+        # The agent should use ask (not recall) since no match
+        assert result.output is not None