Add tests for chat agent & friends

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

@@ -1,4 +1,3 @@
-import logging
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
@@ -14,8 +13,6 @@ from haiku.rag.store.models import SearchResult
 if TYPE_CHECKING:
     from haiku.rag.embeddings import EmbedderWrapper
 
-logger = logging.getLogger(__name__)
-
 
 class CitationInfo(BaseModel):
     """Citation info for frontend display."""
@@ -104,31 +101,25 @@ async def rank_qa_history_by_similarity(
     if len(qa_history) <= top_k:
         return qa_history
 
-    try:
-        # Embed current question
-        question_embedding = np.array(await embedder.embed_query(current_question))
+    # Embed current question
+    question_embedding = np.array(await embedder.embed_query(current_question))
 
-        # Embed Q&A pairs as "Q: {question}\nA: {answer}"
-        qa_texts = [f"Q: {qa.question}\nA: {qa.answer}" for qa in qa_history]
-        qa_embeddings = await embedder.embed_documents(qa_texts)
+    # Embed Q&A pairs as "Q: {question}\nA: {answer}"
+    qa_texts = [f"Q: {qa.question}\nA: {qa.answer}" for qa in qa_history]
+    qa_embeddings = await embedder.embed_documents(qa_texts)
 
-        # Compute similarities
-        similarities: list[tuple[int, float]] = []
-        for i, qa_emb in enumerate(qa_embeddings):
-            sim = _cosine_similarity(question_embedding, np.array(qa_emb))
-            similarities.append((i, sim))
+    # Compute similarities
+    similarities: list[tuple[int, float]] = []
+    for i, qa_emb in enumerate(qa_embeddings):
+        sim = _cosine_similarity(question_embedding, np.array(qa_emb))
+        similarities.append((i, sim))
 
-        # Sort by similarity (descending) and take top-K
-        similarities.sort(key=lambda x: x[1], reverse=True)
-        top_indices = sorted([idx for idx, _ in similarities[:top_k]])
+    # Sort by similarity (descending) and take top-K
+    similarities.sort(key=lambda x: x[1], reverse=True)
+    top_indices = sorted([idx for idx, _ in similarities[:top_k]])
 
-        # Return in original order
-        return [qa_history[i] for i in top_indices]
-
-    except Exception as e:
-        logger.warning(f"Failed to rank qa_history by similarity: {e}")
-        # Fallback: return last top_k entries
-        return qa_history[-top_k:]
+    # Return in original order
+    return [qa_history[i] for i in top_indices]
 
 
 @dataclass

tests/agents/chat/test_chat_agent.py

@@ -227,3 +227,254 @@ async def test_chat_agent_with_qa_history_ranking(allow_model_requests, temp_db_
 
         # Verify qa_history was updated (new Q&A was added)
         assert len(session_state.qa_history) == 7  # 6 original + 1 new
+
+
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_chat_agent_search_tool(allow_model_requests, temp_db_path):
+    """Test the chat agent's search tool functionality."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        # Add test documents
+        await client.create_document(
+            content=DOCLAYNET_CLASS_LABELS,
+            uri="doclaynet-labels",
+            title="DocLayNet Class Labels",
+        )
+        await client.create_document(
+            content=DOCLAYNET_ANNOTATION,
+            uri="doclaynet-annotation",
+            title="DocLayNet Annotation",
+        )
+
+        agent = create_chat_agent(Config)
+        session_state = ChatSessionState(session_id="test-search")
+        deps = ChatDeps(
+            client=client,
+            config=Config,
+            session_state=session_state,
+        )
+
+        # Ask something that should trigger the search tool
+        result = await agent.run(
+            "Search for documents about class labels",
+            deps=deps,
+        )
+
+        assert result.output is not None
+        assert len(result.output) > 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_chat_agent_search_tool_with_filter(allow_model_requests, temp_db_path):
+    """Test the chat agent's search tool with document filter."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        # Add test documents
+        await client.create_document(
+            content=DOCLAYNET_CLASS_LABELS,
+            uri="doclaynet-labels",
+            title="DocLayNet Class Labels",
+        )
+        await client.create_document(
+            content=DOCLAYNET_DATA_SOURCES,
+            uri="doclaynet-sources",
+            title="DocLayNet Sources",
+        )
+
+        agent = create_chat_agent(Config)
+        session_state = ChatSessionState(session_id="test-search-filter")
+        deps = ChatDeps(
+            client=client,
+            config=Config,
+            session_state=session_state,
+        )
+
+        # Ask to search within a specific document
+        result = await agent.run(
+            "Search for information about class labels in the DocLayNet Class Labels document",
+            deps=deps,
+        )
+
+        assert result.output is not None
+
+
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_chat_agent_get_document_tool(allow_model_requests, temp_db_path):
+    """Test the chat agent's get_document tool."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        # Add a test document
+        await client.create_document(
+            content=DOCLAYNET_CLASS_LABELS,
+            uri="doclaynet-labels",
+            title="DocLayNet Class Labels",
+        )
+
+        agent = create_chat_agent(Config)
+        deps = ChatDeps(
+            client=client,
+            config=Config,
+        )
+
+        # Ask to get a specific document
+        result = await agent.run(
+            "Get me the DocLayNet Class Labels document",
+            deps=deps,
+        )
+
+        assert result.output is not None
+        # The response should contain info about the document
+        assert "DocLayNet" in result.output or "class" in result.output.lower()
+
+
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_chat_agent_get_document_not_found(allow_model_requests, temp_db_path):
+    """Test the chat agent's get_document tool when document is not found."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        agent = create_chat_agent(Config)
+        deps = ChatDeps(
+            client=client,
+            config=Config,
+        )
+
+        # Ask for a document that doesn't exist
+        result = await agent.run(
+            "Get me the nonexistent document",
+            deps=deps,
+        )
+
+        assert result.output is not None
+
+
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_search_agent_with_context(allow_model_requests, temp_db_path):
+    """Test SearchAgent's search method with context."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        # Add test documents
+        await client.create_document(
+            content=DOCLAYNET_CLASS_LABELS,
+            uri="doclaynet-labels",
+            title="DocLayNet Class Labels",
+        )
+        await client.create_document(
+            content=DOCLAYNET_ANNOTATION,
+            uri="doclaynet-annotation",
+            title="DocLayNet Annotation",
+        )
+
+        search_agent = SearchAgent(client, Config)
+
+        # Search with context
+        results = await search_agent.search(
+            query="What are the class labels?",
+            context="We're discussing document layout analysis",
+        )
+
+        assert isinstance(results, list)
+
+
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_search_agent_with_filter(allow_model_requests, temp_db_path):
+    """Test SearchAgent's search method with document filter."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        # Add test documents
+        await client.create_document(
+            content=DOCLAYNET_CLASS_LABELS,
+            uri="doclaynet-labels",
+            title="DocLayNet Class Labels",
+        )
+        await client.create_document(
+            content=DOCLAYNET_DATA_SOURCES,
+            uri="doclaynet-sources",
+            title="DocLayNet Sources",
+        )
+
+        search_agent = SearchAgent(client, Config)
+
+        # Search with filter - only the labels document
+        results = await search_agent.search(
+            query="What information is available?",
+            filter="uri LIKE '%labels%'",
+        )
+
+        assert isinstance(results, list)
+        # Results should only come from the labels document
+        for r in results:
+            assert "labels" in (r.document_uri or "")
+
+
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_search_agent_deduplication(allow_model_requests, temp_db_path):
+    """Test SearchAgent deduplicates results by chunk_id."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        # Add test documents
+        await client.create_document(
+            content=DOCLAYNET_CLASS_LABELS,
+            uri="doclaynet-labels",
+            title="DocLayNet Class Labels",
+        )
+
+        search_agent = SearchAgent(client, Config)
+
+        # Search - the search agent will likely run multiple queries
+        # that could return the same chunk, which should be deduplicated
+        results = await search_agent.search(
+            query="Tell me about class labels and their counts",
+        )
+
+        assert isinstance(results, list)
+
+        # Verify no duplicate chunk_ids
+        chunk_ids = [r.chunk_id for r in results if r.chunk_id]
+        assert len(chunk_ids) == len(set(chunk_ids)), "Found duplicate chunk_ids"
+
+
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_search_agent_no_results(allow_model_requests, temp_db_path):
+    """Test SearchAgent handles no results gracefully."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        search_agent = SearchAgent(client, Config)
+
+        # Search in empty database
+        results = await search_agent.search(
+            query="Find information about nonexistent topic xyz123",
+        )
+
+        assert isinstance(results, list)
+        assert len(results) == 0
+
+
+@pytest.mark.asyncio
+@pytest.mark.vcr()
+async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path):
+    """Test that the ask tool adds citations to the response."""
+    async with HaikuRAG(temp_db_path, create=True) as client:
+        # Add a document with specific content
+        await client.create_document(
+            content=DOCLAYNET_CLASS_LABELS,
+            uri="doclaynet-labels",
+            title="DocLayNet Class Labels",
+        )
+
+        agent = create_chat_agent(Config)
+        session_state = ChatSessionState(session_id="test-citations")
+        deps = ChatDeps(
+            client=client,
+            config=Config,
+            session_state=session_state,
+        )
+
+        # Ask a question that should use the ask tool with citations
+        result = await agent.run(
+            "What is the highest count class in the DocLayNet dataset?",
+            deps=deps,
+        )
+
+        assert result.output is not None
+        # The qa_history should have been updated with the new Q&A
+        assert len(session_state.qa_history) >= 1

tests/agents/chat/test_state.py

@@ -5,6 +5,8 @@ import pytest
 from haiku.rag.agents.chat.state import (
     CitationInfo,
     QAResponse,
+    build_document_filter,
+    format_conversation_context,
     rank_qa_history_by_similarity,
 )
 from haiku.rag.client import HaikuRAG
@@ -166,3 +168,66 @@ async def test_rank_qa_history_preserves_order(temp_db_path, allow_model_request
         result_questions = [qa.question for qa in result]
         assert "What is Python?" in result_questions
         assert "How to use Python for data science?" in result_questions
+
+
+def test_format_conversation_context_empty():
+    """Test format_conversation_context with empty history."""
+    result = format_conversation_context([])
+    assert result == ""
+
+
+def test_format_conversation_context_with_history():
+    """Test format_conversation_context formats qa_history as XML."""
+    citation = CitationInfo(
+        index=1,
+        document_id="doc-123",
+        chunk_id="chunk-456",
+        document_uri="test.md",
+        document_title="Test Document",
+        content="Test content",
+    )
+    qa_history = [
+        QAResponse(
+            question="What is Python?",
+            answer="A programming language",
+            citations=[citation],
+        ),
+        QAResponse(
+            question="What is Java?",
+            answer="Another programming language",
+        ),
+    ]
+
+    result = format_conversation_context(qa_history)
+
+    assert "<conversation_context>" in result
+    assert "previous_qa" in result
+    assert "What is Python?" in result
+    assert "A programming language" in result
+    assert "What is Java?" in result
+    assert "Another programming language" in result
+    assert "Test Document" in result  # source from first citation
+
+
+def test_build_document_filter_simple():
+    """Test build_document_filter with simple name."""
+    result = build_document_filter("mytest")
+    assert "LOWER(uri) LIKE LOWER('%mytest%')" in result
+    assert "LOWER(title) LIKE LOWER('%mytest%')" in result
+
+
+def test_build_document_filter_with_spaces():
+    """Test build_document_filter handles spaces correctly."""
+    result = build_document_filter("TB MED 593")
+    # Should include both the original (with spaces) and without spaces
+    assert "LOWER(uri) LIKE LOWER('%TB MED 593%')" in result
+    assert "LOWER(uri) LIKE LOWER('%TBMED593%')" in result
+    assert "LOWER(title) LIKE LOWER('%TB MED 593%')" in result
+    assert "LOWER(title) LIKE LOWER('%TBMED593%')" in result
+
+
+def test_build_document_filter_escapes_quotes():
+    """Test build_document_filter escapes single quotes."""
+    result = build_document_filter("O'Reilly")
+    # Single quotes should be doubled for SQL escaping
+    assert "O''Reilly" in result

tests/cassettes/test_chat_agent/test_chat_agent_get_document_not_found.yaml

@@ -0,0 +1,154 @@
+interactions:
+- request:
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate, zstd
+      connection:
+      - keep-alive
+      content-length:
+      - '4073'
+      content-type:
+      - application/json
+      host:
+      - localhost:11434
+    method: POST
+    parsed_body:
+      messages:
+      - content: |-
+          You are a helpful research assistant powered by haiku.rag, a knowledge base system.
+
+          You have access to a knowledge base of documents. Use your tools to search and answer questions.
+
+          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
+
+          How to decide which tool to use:
+          - "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.
+          - "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:
+          - If user says "search in <doc>", "find in <doc>", "answer from <doc>", or "<topic> in <doc>":
+            - Extract the TOPIC as `query`/`question`
+            - Extract the DOCUMENT NAME as `document_name`
+          - Examples for search:
+            - "search for embeddings in the ML paper" → query="embeddings", document_name="ML paper"
+            - "find transformer architecture in 2412.00566" → query="transformer architecture", document_name="2412.00566"
+          - Examples for ask:
+            - "what does the ML paper say about embeddings?" → question="what are the embedding methods?", document_name="ML paper"
+            - "answer from 2412.00566 about model training" → question="how is the model trained?", document_name="2412.00566"
+
+          Be friendly and conversational. When you use the "ask" tool, summarize the key findings for the user.
+        role: system
+      - content: Get me the nonexistent document
+        role: user
+      model: gpt-oss
+      reasoning_effort: low
+      stream: false
+      tool_choice: auto
+      tools:
+      - function:
+          description: |-
+            Search the knowledge base for relevant documents.
+
+            Use this when you need to find documents or explore the knowledge base.
+            Results are displayed to the user - just list the titles found.
+          name: search
+          parameters:
+            additionalProperties: false
+            properties:
+              document_name:
+                anyOf:
+                - type: string
+                - type: 'null'
+                default: null
+                description: Optional document name/title to search within
+              limit:
+                anyOf:
+                - type: integer
+                - type: 'null'
+                default: null
+                description: 'Number of results to return (default: 5)'
+              query:
+                description: The search query (what to search for)
+                type: string
+            required:
+            - query
+            type: object
+        type: function
+      - function:
+          description: |-
+            Answer a specific question using the knowledge base.
+
+            Use this for direct questions that need a focused answer with citations.
+            Uses a research graph for planning, searching, and synthesis.
+          name: ask
+          parameters:
+            additionalProperties: false
+            properties:
+              document_name:
+                anyOf:
+                - type: string
+                - type: 'null'
+                default: null
+                description: Optional document name/title to search within (e.g., "tbmed593", "army manual")
+              question:
+                description: The question to answer
+                type: string
+            required:
+            - question
+            type: object
+        type: function
+      - function:
+          description: |-
+            Retrieve a specific document by title or URI.
+
+            Use this when the user wants to fetch/get/retrieve a specific document.
+          name: get_document
+          parameters:
+            additionalProperties: false
+            properties:
+              query:
+                description: The document title or URI to look up
+                type: string
+            required:
+            - query
+            type: object
+          strict: true
+        type: function
+    uri: http://localhost:11434/v1/chat/completions
+  response:
+    headers:
+      content-length:
+      - '585'
+      content-type:
+      - application/json
+    parsed_body:
+      choices:
+      - finish_reason: stop
+        index: 0
+        message:
+          content: I’m sorry—I couldn’t find a document matching that name in the knowledge base. If you have any other request
+            or need help locating a different resource, just let me know!
+          reasoning: User asks for nonexistent document. Use get_document? but tool should not fabricate. We can explain not
+            found.
+          role: assistant
+      created: 1768225927
+      id: chatcmpl-215
+      model: gpt-oss
+      object: chat.completion
+      system_fingerprint: fp_ollama
+      usage:
+        completion_tokens: 68
+        prompt_tokens: 842
+        total_tokens: 910
+    status:
+      code: 200
+      message: OK
+version: 1