Concolidate tests

tests/agents/chat/test_chat_agent.py

@@ -251,40 +251,6 @@ async def test_chat_agent_search_tool(allow_model_requests, temp_db_path):
         assert len(result.output) > 0
 
 
-@pytest.mark.asyncio
-@pytest.mark.vcr()
-async def test_chat_agent_search_with_state_key(allow_model_requests, temp_db_path):
-    """Test search tool emits keyed state when state_key is set."""
-    async with HaikuRAG(temp_db_path, create=True) as client:
-        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,
-            state_key=AGUI_STATE_KEY,
-        )
-
-        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):
@@ -501,35 +467,6 @@ async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path)
         assert len(session_state.qa_history) >= 1
 
 
-@pytest.mark.asyncio
-@pytest.mark.vcr()
-async def test_chat_agent_ask_with_state_key(allow_model_requests, temp_db_path):
-    """Test ask tool emits keyed state when state_key is set."""
-    async with HaikuRAG(temp_db_path, create=True) as client:
-        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-ask-keyed")
-        deps = ChatDeps(
-            client=client,
-            config=Config,
-            session_state=session_state,
-            state_key=AGUI_STATE_KEY,
-        )
-
-        result = await agent.run(
-            "What is the highest count class in the DocLayNet dataset?",
-            deps=deps,
-        )
-
-        assert result.output is not None
-        assert len(session_state.qa_history) >= 1
-
-
 @pytest.mark.asyncio
 @pytest.mark.vcr()
 async def test_chat_agent_ask_triggers_background_summarization(
@@ -801,3 +738,186 @@ def test_search_tool_citation_registry_logic():
         "chunk-c": 3,
         "chunk-d": 4,
     }
+
+
+# =============================================================================
+# Prior Answer Recall Tests
+# =============================================================================
+
+
+def test_cosine_similarity_identical_vectors():
+    """Test cosine similarity returns 1.0 for identical vectors."""
+    from haiku.rag.agents.chat.agent import _cosine_similarity
+
+    vec = [1.0, 2.0, 3.0]
+    assert _cosine_similarity(vec, vec) == pytest.approx(1.0)
+
+
+def test_cosine_similarity_orthogonal_vectors():
+    """Test cosine similarity returns 0.0 for orthogonal vectors."""
+    from haiku.rag.agents.chat.agent import _cosine_similarity
+
+    vec1 = [1.0, 0.0, 0.0]
+    vec2 = [0.0, 1.0, 0.0]
+    assert _cosine_similarity(vec1, vec2) == pytest.approx(0.0)
+
+
+def test_cosine_similarity_opposite_vectors():
+    """Test cosine similarity returns -1.0 for opposite vectors."""
+    from haiku.rag.agents.chat.agent import _cosine_similarity
+
+    vec1 = [1.0, 2.0, 3.0]
+    vec2 = [-1.0, -2.0, -3.0]
+    assert _cosine_similarity(vec1, vec2) == pytest.approx(-1.0)
+
+
+def test_cosine_similarity_zero_vector():
+    """Test cosine similarity handles zero vectors gracefully."""
+    from haiku.rag.agents.chat.agent import _cosine_similarity
+
+    vec = [1.0, 2.0, 3.0]
+    zero = [0.0, 0.0, 0.0]
+    assert _cosine_similarity(vec, zero) == 0.0
+    assert _cosine_similarity(zero, vec) == 0.0
+    assert _cosine_similarity(zero, zero) == 0.0
+
+
+def test_prior_answer_relevance_threshold_constant():
+    """Test PRIOR_ANSWER_RELEVANCE_THRESHOLD is set to expected value."""
+    from haiku.rag.agents.chat.agent import PRIOR_ANSWER_RELEVANCE_THRESHOLD
+
+    assert PRIOR_ANSWER_RELEVANCE_THRESHOLD == 0.7
+
+
+def test_prior_answer_matching_above_threshold():
+    """Test that similar questions (above threshold) are matched."""
+    from haiku.rag.agents.chat.agent import (
+        PRIOR_ANSWER_RELEVANCE_THRESHOLD,
+        _cosine_similarity,
+    )
+
+    # Simulate two nearly identical question embeddings
+    question_embedding = [0.5, 0.5, 0.5, 0.5]
+    prior_embedding = [0.51, 0.49, 0.5, 0.5]  # Very similar
+
+    similarity = _cosine_similarity(question_embedding, prior_embedding)
+    assert similarity >= PRIOR_ANSWER_RELEVANCE_THRESHOLD
+
+
+def test_prior_answer_matching_below_threshold():
+    """Test that dissimilar questions (below threshold) are not matched."""
+    from haiku.rag.agents.chat.agent import (
+        PRIOR_ANSWER_RELEVANCE_THRESHOLD,
+        _cosine_similarity,
+    )
+
+    # Simulate two different question embeddings
+    question_embedding = [1.0, 0.0, 0.0, 0.0]
+    prior_embedding = [0.0, 1.0, 0.0, 0.0]  # Orthogonal = very different
+
+    similarity = _cosine_similarity(question_embedding, prior_embedding)
+    assert similarity < PRIOR_ANSWER_RELEVANCE_THRESHOLD
+
+
+def test_qa_response_embedding_cache():
+    """Test that QAResponse stores and retrieves question_embedding correctly."""
+    embedding = [0.1, 0.2, 0.3, 0.4]
+    qa = QAResponse(
+        question="What is X?",
+        answer="X is Y.",
+        confidence=0.9,
+        question_embedding=embedding,
+    )
+
+    assert qa.question_embedding == embedding
+    # Embedding should be excluded from serialization (AG-UI state)
+    serialized = qa.model_dump()
+    assert "question_embedding" not in serialized
+
+
+def test_qa_response_embedding_default_none():
+    """Test that QAResponse.question_embedding defaults to None."""
+    qa = QAResponse(
+        question="What is X?",
+        answer="X is Y.",
+        confidence=0.9,
+    )
+
+    assert qa.question_embedding is None
+
+
+# =============================================================================
+# Background Task Cancellation Tests
+# =============================================================================
+
+
+@pytest.mark.asyncio
+async def test_summarization_task_cancellation():
+    """Test that new summarization tasks cancel previous ones for same session."""
+    import asyncio
+
+    from haiku.rag.agents.chat.agent import _summarization_tasks
+
+    # Clear any existing tasks
+    _summarization_tasks.clear()
+
+    session_id = "test-cancel-session"
+
+    # Create a slow task that simulates summarization
+    async def slow_task():
+        await asyncio.sleep(10)  # Would take 10 seconds
+
+    # Start first task
+    task1 = asyncio.create_task(slow_task())
+    _summarization_tasks[session_id] = task1
+
+    # Simulate what happens when second ask comes in - cancel first task
+    if session_id in _summarization_tasks:
+        _summarization_tasks[session_id].cancel()
+
+    # Yield to let cancellation propagate
+    await asyncio.sleep(0)
+
+    # Start second task
+    task2 = asyncio.create_task(slow_task())
+    _summarization_tasks[session_id] = task2
+
+    # First task should be cancelled
+    assert task1.cancelled() or task1.done()
+
+    # Second task should be running
+    assert not task2.done()
+
+    # Cleanup
+    task2.cancel()
+    try:
+        await task2
+    except asyncio.CancelledError:
+        pass
+    _summarization_tasks.clear()
+
+
+@pytest.mark.asyncio
+async def test_summarization_task_cleanup_on_completion():
+    """Test that completed tasks are cleaned up from _summarization_tasks."""
+    import asyncio
+
+    from haiku.rag.agents.chat.agent import _summarization_tasks
+
+    _summarization_tasks.clear()
+
+    session_id = "test-cleanup-session"
+
+    # Create a fast task
+    async def fast_task():
+        await asyncio.sleep(0.01)
+
+    task = asyncio.create_task(fast_task())
+    _summarization_tasks[session_id] = task
+    task.add_done_callback(lambda t: _summarization_tasks.pop(session_id, None))
+
+    # Wait for completion
+    await task
+
+    # Task should be cleaned up
+    assert session_id not in _summarization_tasks

tests/agents/chat/test_context.py

@@ -66,8 +66,7 @@ class TestSummarizeSession:
     """Tests for summarize_session function."""
 
     @pytest.mark.asyncio
-    @pytest.mark.vcr()
-    async def test_summarize_session_empty_history(self, allow_model_requests):
+    async def test_summarize_session_empty_history(self):
         """Test summarize_session with empty qa_history returns empty string."""
         from haiku.rag.agents.chat.context import summarize_session
 

tests/agents/chat/test_state.py

@@ -376,48 +376,43 @@ def test_chat_deps_state_setter_ignores_session_context():
     assert deps.session_state.session_context.summary == "Server-side context"
 
 
-def test_citation_registry_get_or_assign_index_first_chunk():
-    """Test get_or_assign_index assigns index 1 to first chunk."""
+def test_citation_registry_index_assignment():
+    """Test get_or_assign_index basic index assignment behavior.
+
+    Verifies:
+    - First chunk gets index 1
+    - Second unique chunk gets index 2
+    - Same chunk_id always returns same index
+    """
     from haiku.rag.agents.chat.state import ChatSessionState
 
     session_state = ChatSessionState(session_id="test")
-    index = session_state.get_or_assign_index("chunk-abc")
-    assert index == 1
 
-
-def test_citation_registry_get_or_assign_index_second_chunk():
-    """Test get_or_assign_index assigns incremental indices to new chunks."""
-    from haiku.rag.agents.chat.state import ChatSessionState
-
-    session_state = ChatSessionState(session_id="test")
+    # First chunk gets index 1
     index1 = session_state.get_or_assign_index("chunk-abc")
-    index2 = session_state.get_or_assign_index("chunk-def")
     assert index1 == 1
+
+    # Second unique chunk gets index 2
+    index2 = session_state.get_or_assign_index("chunk-def")
     assert index2 == 2
 
-
-def test_citation_registry_get_or_assign_index_same_chunk():
-    """Test get_or_assign_index returns same index for same chunk_id."""
-    from haiku.rag.agents.chat.state import ChatSessionState
-
-    session_state = ChatSessionState(session_id="test")
-    index1 = session_state.get_or_assign_index("chunk-abc")
-    index2 = session_state.get_or_assign_index("chunk-abc")
-    assert index1 == index2 == 1
+    # Same chunk_id returns same index (not incremented)
+    index1_again = session_state.get_or_assign_index("chunk-abc")
+    assert index1_again == 1
 
 
-def test_citation_registry_get_or_assign_index_stability():
-    """Test citation indices are stable across multiple calls."""
+def test_citation_registry_stability():
+    """Test citation indices are stable across multiple calls in any order."""
     from haiku.rag.agents.chat.state import ChatSessionState
 
     session_state = ChatSessionState(session_id="test")
 
-    # First call assigns indices 1, 2, 3
+    # First round assigns indices 1, 2, 3
     idx_a = session_state.get_or_assign_index("chunk-a")
     idx_b = session_state.get_or_assign_index("chunk-b")
     idx_c = session_state.get_or_assign_index("chunk-c")
 
-    # Second round - existing chunks keep their indices
+    # Second round - existing chunks keep their indices regardless of order
     assert session_state.get_or_assign_index("chunk-b") == idx_b
     assert session_state.get_or_assign_index("chunk-a") == idx_a
     assert session_state.get_or_assign_index("chunk-c") == idx_c
@@ -427,38 +422,28 @@ def test_citation_registry_get_or_assign_index_stability():
     assert idx_d == 4
 
 
-def test_citation_registry_serialization():
-    """Test citation_registry is included in model_dump for AG-UI state."""
+def test_citation_registry_serialization_roundtrip():
+    """Test citation_registry serializes and deserializes correctly for AG-UI state."""
     from haiku.rag.agents.chat.state import ChatSessionState
 
-    session_state = ChatSessionState(session_id="test")
-    session_state.get_or_assign_index("chunk-a")
-    session_state.get_or_assign_index("chunk-b")
+    # Create state and assign indices
+    original = ChatSessionState(session_id="test")
+    original.get_or_assign_index("chunk-a")
+    original.get_or_assign_index("chunk-b")
 
-    state_dict = session_state.model_dump()
+    # Serialize
+    state_dict = original.model_dump()
     assert "citation_registry" in state_dict
     assert state_dict["citation_registry"] == {"chunk-a": 1, "chunk-b": 2}
 
-
-def test_citation_registry_deserialization():
-    """Test citation_registry is restored from dict."""
-    from haiku.rag.agents.chat.state import ChatSessionState
-
-    # Simulate state from AG-UI using model_validate (proper Pydantic deserialization)
-    session_state = ChatSessionState.model_validate(
-        {
-            "session_id": "test",
-            "citations": [],
-            "qa_history": [],
-            "citation_registry": {"chunk-a": 1, "chunk-b": 2},
-        }
-    )
+    # Deserialize (simulating AG-UI state restoration)
+    restored = ChatSessionState.model_validate(state_dict)
 
     # Existing chunks should return their persisted indices
-    assert session_state.get_or_assign_index("chunk-a") == 1
-    assert session_state.get_or_assign_index("chunk-b") == 2
+    assert restored.get_or_assign_index("chunk-a") == 1
+    assert restored.get_or_assign_index("chunk-b") == 2
     # New chunk should get next index
-    assert session_state.get_or_assign_index("chunk-c") == 3
+    assert restored.get_or_assign_index("chunk-c") == 3
 
 
 def test_chat_deps_state_getter_includes_citation_registry():

tests/agents/research/test_research_graph.py

@@ -62,3 +62,125 @@ def test_research_plan_rejects_too_many_sub_questions():
 
     with pytest.raises(ValidationError, match="Cannot have more than 12"):
         ResearchPlan(sub_questions=[f"q{i}" for i in range(13)])
+
+
+# =============================================================================
+# Conversational Graph Tests
+# =============================================================================
+
+
+def test_build_conversational_graph_returns_graph():
+    """Test build_conversational_graph returns a valid Graph instance."""
+    from pydantic_graph.beta import Graph
+
+    from haiku.rag.agents.research.graph import build_conversational_graph
+
+    graph = build_conversational_graph()
+    assert graph is not None
+    assert isinstance(graph, Graph)
+
+
+def test_conversational_answer_model():
+    """Test ConversationalAnswer model can be created with all fields."""
+    from haiku.rag.agents.research.models import Citation, ConversationalAnswer
+
+    citation = Citation(
+        index=1,
+        document_id="doc-1",
+        chunk_id="chunk-1",
+        document_uri="test.md",
+        document_title="Test Doc",
+        content="Test content",
+    )
+
+    answer = ConversationalAnswer(
+        answer="The answer is 42.",
+        citations=[citation],
+        confidence=0.95,
+    )
+
+    assert answer.answer == "The answer is 42."
+    assert len(answer.citations) == 1
+    assert answer.confidence == 0.95
+
+
+def test_conversational_answer_default_values():
+    """Test ConversationalAnswer uses correct default values."""
+    from haiku.rag.agents.research.models import ConversationalAnswer
+
+    answer = ConversationalAnswer(answer="Just the answer.")
+
+    assert answer.answer == "Just the answer."
+    assert answer.citations == []
+    assert answer.confidence == 1.0
+
+
+def test_format_context_for_prompt_basic():
+    """Test format_context_for_prompt with basic context."""
+    from haiku.rag.agents.research.dependencies import ResearchContext
+    from haiku.rag.agents.research.graph import format_context_for_prompt
+
+    context = ResearchContext(original_question="What is X?")
+    result = format_context_for_prompt(context)
+
+    assert "<context>" in result
+    assert "What is X?" in result
+
+
+def test_format_context_for_prompt_with_session_context():
+    """Test format_context_for_prompt includes session_context as background."""
+    from haiku.rag.agents.research.dependencies import ResearchContext
+    from haiku.rag.agents.research.graph import format_context_for_prompt
+
+    context = ResearchContext(
+        original_question="What is Y?",
+        session_context="Previous discussion about topic Z.",
+    )
+    result = format_context_for_prompt(context)
+
+    assert "<background>" in result
+    assert "Previous discussion" in result
+    assert "What is Y?" in result
+
+
+def test_format_context_for_prompt_excludes_pending_questions():
+    """Test format_context_for_prompt can exclude pending questions."""
+    from haiku.rag.agents.research.dependencies import ResearchContext
+    from haiku.rag.agents.research.graph import format_context_for_prompt
+
+    context = ResearchContext(
+        original_question="Main question?",
+        sub_questions=["Sub Q1?", "Sub Q2?"],
+    )
+
+    # With pending questions (default)
+    with_pending = format_context_for_prompt(context, include_pending_questions=True)
+    assert "Sub Q1?" in with_pending
+
+    # Without pending questions (for synthesis)
+    without_pending = format_context_for_prompt(
+        context, include_pending_questions=False
+    )
+    assert "Sub Q1?" not in without_pending
+
+
+def test_format_context_for_prompt_with_prior_answers():
+    """Test format_context_for_prompt includes prior_answers."""
+    from haiku.rag.agents.research.dependencies import ResearchContext
+    from haiku.rag.agents.research.graph import format_context_for_prompt
+    from haiku.rag.agents.research.models import SearchAnswer
+
+    context = ResearchContext(original_question="Main question?")
+    context.add_qa_response(
+        SearchAnswer(
+            query="Sub question?",
+            answer="The answer is here.",
+            confidence=0.9,
+        )
+    )
+
+    result = format_context_for_prompt(context)
+
+    assert "<prior_answers>" in result
+    assert "Sub question?" in result
+    assert "The answer is here." in result