From 3e66d6b21d70ff4cadd26f51a4ceaf66bcd25965 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 12 Jan 2026 17:02:23 +0200 Subject: [PATCH] Make qa list a FIFO with 50 max, keep a cache of embeddings --- haiku_rag_slim/haiku/rag/agents/chat/agent.py | 6 ++++ haiku_rag_slim/haiku/rag/agents/chat/state.py | 35 +++++++++++++++++-- tests/agents/chat/test_chat_agent.py | 35 +++++++++++++++++++ tests/agents/chat/test_state.py | 19 +++++++++- 4 files changed, 91 insertions(+), 4 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index 2ae7cb2c..ffae5408 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/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( diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py index 3e49bfca..4a0425a5 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/state.py +++ b/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]] = [] diff --git a/tests/agents/chat/test_chat_agent.py b/tests/agents/chat/test_chat_agent.py index 3f85df1d..b022aac2 100644 --- a/tests/agents/chat/test_chat_agent.py +++ b/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}" diff --git a/tests/agents/chat/test_state.py b/tests/agents/chat/test_state.py index a75534a8..df37728d 100644 --- a/tests/agents/chat/test_state.py +++ b/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