Make qa list a FIFO with 50 max, keep a cache of embeddings

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

@@ -4,6 +4,7 @@ from pydantic_ai import Agent, RunContext, ToolReturn
 from haiku.rag.agents.chat.prompts import CHAT_SYSTEM_PROMPT
 from haiku.rag.agents.chat.search import SearchAgent
 from haiku.rag.agents.chat.state import (
+    MAX_QA_HISTORY,
     ChatDeps,
     ChatSessionState,
     CitationInfo,
@@ -217,6 +218,11 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
                 citations=citation_infos,
             )
             ctx.deps.session_state.qa_history.append(qa_response)
+            # Enforce FIFO limit
+            if len(ctx.deps.session_state.qa_history) > MAX_QA_HISTORY:
+                ctx.deps.session_state.qa_history = ctx.deps.session_state.qa_history[
+                    -MAX_QA_HISTORY:
+                ]
 
         # Build new state with citations AND accumulated qa_history
         new_state = ChatSessionState(

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

@@ -1,3 +1,4 @@
+import hashlib
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING
 
@@ -13,6 +14,15 @@ from haiku.rag.store.models import SearchResult
 if TYPE_CHECKING:
     from haiku.rag.embeddings import EmbedderWrapper
 
+MAX_QA_HISTORY = 50
+
+_embedding_cache: dict[str, list[float]] = {}
+
+
+def _qa_cache_key(question: str, answer: str) -> str:
+    """Generate cache key from Q/A content."""
+    return hashlib.sha256(f"Q: {question}\nA: {answer}".encode()).hexdigest()
+
 
 class CitationInfo(BaseModel):
     """Citation info for frontend display."""
@@ -104,9 +114,28 @@ async def rank_qa_history_by_similarity(
     # 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)
+    # Check cache and collect uncached entries
+    qa_embeddings: list[list[float]] = []
+    uncached_indices: list[int] = []
+    uncached_texts: list[str] = []
+
+    for i, qa in enumerate(qa_history):
+        cache_key = _qa_cache_key(qa.question, qa.answer)
+        if cache_key in _embedding_cache:
+            qa_embeddings.append(_embedding_cache[cache_key])
+        else:
+            qa_embeddings.append([])  # placeholder
+            uncached_indices.append(i)
+            uncached_texts.append(f"Q: {qa.question}\nA: {qa.answer}")
+
+    # Embed only uncached entries
+    if uncached_texts:
+        new_embeddings = await embedder.embed_documents(uncached_texts)
+        for idx, embedding in zip(uncached_indices, new_embeddings):
+            qa = qa_history[idx]
+            cache_key = _qa_cache_key(qa.question, qa.answer)
+            _embedding_cache[cache_key] = embedding
+            qa_embeddings[idx] = embedding
 
     # Compute similarities
     similarities: list[tuple[int, float]] = []

tests/agents/chat/test_chat_agent.py

@@ -10,6 +10,7 @@ from haiku.rag.agents.chat import (
     SearchAgent,
     create_chat_agent,
 )
+from haiku.rag.agents.chat.state import MAX_QA_HISTORY
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
 
@@ -478,3 +479,37 @@ async def test_chat_agent_ask_adds_citations(allow_model_requests, temp_db_path)
         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
+
+
+def test_fifo_limit_enforcement():
+    """Test that FIFO limit enforcement logic works correctly.
+
+    This tests the FIFO trimming logic used in the ask() tool:
+    if len(qa_history) > MAX_QA_HISTORY:
+        qa_history = qa_history[-MAX_QA_HISTORY:]
+    """
+    # Create a session state with MAX_QA_HISTORY + 1 entries
+    qa_history = [
+        QAResponse(
+            question=f"Question {i}",
+            answer=f"Answer {i}",
+            confidence=0.9,
+        )
+        for i in range(MAX_QA_HISTORY + 1)
+    ]
+
+    session_state = ChatSessionState(
+        session_id="test-fifo",
+        qa_history=qa_history,
+    )
+
+    # Simulate the FIFO enforcement from agent.py
+    if len(session_state.qa_history) > MAX_QA_HISTORY:
+        session_state.qa_history = session_state.qa_history[-MAX_QA_HISTORY:]
+
+    # History should be trimmed to MAX_QA_HISTORY
+    assert len(session_state.qa_history) == MAX_QA_HISTORY
+    # The first entry should now be "Question 1" (Question 0 was dropped)
+    assert session_state.qa_history[0].question == "Question 1"
+    # The last entry should be the last added question
+    assert session_state.qa_history[-1].question == f"Question {MAX_QA_HISTORY}"

tests/agents/chat/test_state.py

@@ -3,8 +3,11 @@ from pathlib import Path
 import pytest
 
 from haiku.rag.agents.chat.state import (
+    MAX_QA_HISTORY,
     CitationInfo,
     QAResponse,
+    _embedding_cache,
+    _qa_cache_key,
     build_document_filter,
     format_conversation_context,
     rank_qa_history_by_similarity,
@@ -66,7 +69,7 @@ async def test_rank_qa_history_small_list(temp_db_path, allow_model_requests):
 @pytest.mark.asyncio
 @pytest.mark.vcr()
 async def test_rank_qa_history_returns_top_k(temp_db_path, allow_model_requests):
-    """Test ranking returns top-K most similar entries."""
+    """Test ranking returns top-K most similar entries and populates cache."""
     async with HaikuRAG(temp_db_path, create=True) as client:
         embedder = client.chunk_repository.embedder
 
@@ -106,6 +109,9 @@ async def test_rank_qa_history_returns_top_k(temp_db_path, allow_model_requests)
             ),
         ]
 
+        # Clear cache to verify it gets populated
+        _embedding_cache.clear()
+
         # Ask a question related to class labels (Q1)
         result = await rank_qa_history_by_similarity(
             current_question="Which class label has the highest count in DocLayNet?",
@@ -121,6 +127,12 @@ async def test_rank_qa_history_returns_top_k(temp_db_path, allow_model_requests)
         result_questions = [qa.question for qa in result]
         assert "What are the 11 class labels in DocLayNet?" in result_questions
 
+        # Verify cache is populated for all Q/A pairs
+        for qa in qa_history:
+            cache_key = _qa_cache_key(qa.question, qa.answer)
+            assert cache_key in _embedding_cache
+            assert len(_embedding_cache[cache_key]) > 0
+
 
 @pytest.mark.asyncio
 @pytest.mark.vcr()
@@ -231,3 +243,8 @@ def test_build_document_filter_escapes_quotes():
     result = build_document_filter("O'Reilly")
     # Single quotes should be doubled for SQL escaping
     assert "O''Reilly" in result
+
+
+def test_max_qa_history_constant():
+    """Test MAX_QA_HISTORY constant value."""
+    assert MAX_QA_HISTORY == 50