From 8f35033a4254c2a24a6972a2d0f644abb9d37ec8 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 26 Jan 2026 12:17:36 +0200 Subject: [PATCH] Update search(), ask() to use get_or_assign_index() --- haiku_rag_slim/haiku/rag/agents/chat/agent.py | 96 +++++++----- haiku_rag_slim/haiku/rag/agents/chat/state.py | 17 ++ tests/agents/chat/test_chat_agent.py | 97 ++++++++++++ tests/agents/chat/test_state.py | 147 ++++++++++++++++++ 4 files changed, 320 insertions(+), 37 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index abb82b6d..63247dc2 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -97,22 +97,28 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: if not results: return ToolReturn(return_value="No results found.") - # Build citation infos for frontend display - citation_infos = [ - Citation( - index=i + 1, - document_id=r.document_id or "", - chunk_id=r.chunk_id or "", - document_uri=r.document_uri or "", - document_title=r.document_title, - page_numbers=r.page_numbers or [], - headings=r.headings, - content=r.content, + # Build citation infos using stable registry indices + citation_infos = [] + for r in results: + chunk_id = r.chunk_id or "" + if ctx.deps.session_state is not None and chunk_id: + index = ctx.deps.session_state.get_or_assign_index(chunk_id) + else: + index = len(citation_infos) + 1 + citation_infos.append( + Citation( + index=index, + document_id=r.document_id or "", + chunk_id=chunk_id, + document_uri=r.document_uri or "", + document_title=r.document_title, + page_numbers=r.page_numbers or [], + headings=r.headings, + content=r.content, + ) ) - for i, r in enumerate(results) - ] - # Build new state with citations + # Build new state with citations and registry session_id = ctx.deps.session_state.session_id if ctx.deps.session_state else "" new_state = ChatSessionState( session_id=session_id, @@ -126,20 +132,25 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: document_filter=( ctx.deps.session_state.document_filter if ctx.deps.session_state else [] ), + citation_registry=( + ctx.deps.session_state.citation_registry + if ctx.deps.session_state + else {} + ), ) # Return detailed results for the agent to present result_lines = [] - for i, r in enumerate(results): - title = r.document_title or r.document_uri or "Unknown" + for c in citation_infos: + title = c.document_title or c.document_uri or "Unknown" # Truncate content for display - snippet = r.content[:300].replace("\n", " ").strip() - if len(r.content) > 300: + snippet = c.content[:300].replace("\n", " ").strip() + if len(c.content) > 300: snippet += "..." - line = f"[{i + 1}] **{title}**" - if r.page_numbers: - line += f" (pages {', '.join(map(str, r.page_numbers))})" + line = f"[{c.index}] **{title}**" + if c.page_numbers: + line += f" (pages {', '.join(map(str, c.page_numbers))})" line += f"\n {snippet}" result_lines.append(line) @@ -215,20 +226,26 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: result = await graph.run(state=state, deps=deps) - # Build citation infos for frontend and history - citation_infos = [ - Citation( - index=i + 1, - document_id=c.document_id, - chunk_id=c.chunk_id, - document_uri=c.document_uri, - document_title=c.document_title, - page_numbers=c.page_numbers, - headings=c.headings, - content=c.content, + # Build citation infos using stable registry indices + citation_infos = [] + for c in result.citations: + # Use registry for stable indices across calls + if ctx.deps.session_state is not None: + index = ctx.deps.session_state.get_or_assign_index(c.chunk_id) + else: + index = len(citation_infos) + 1 + citation_infos.append( + Citation( + index=index, + document_id=c.document_id, + chunk_id=c.chunk_id, + document_uri=c.document_uri, + document_title=c.document_title, + page_numbers=c.page_numbers, + headings=c.headings, + content=c.content, + ) ) - for i, c in enumerate(result.citations) - ] # Accumulate Q&A in session state with full citation metadata if ctx.deps.session_state is not None: @@ -260,7 +277,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: _summarization_tasks[session_id] = task task.add_done_callback(lambda t: _summarization_tasks.pop(session_id, None)) - # Build new state with citations AND accumulated qa_history + # Build new state with citations, qa_history, and registry new_state = ChatSessionState( session_id=session_id, citations=citation_infos, @@ -273,12 +290,17 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: document_filter=( ctx.deps.session_state.document_filter if ctx.deps.session_state else [] ), + citation_registry=( + ctx.deps.session_state.citation_registry + if ctx.deps.session_state + else {} + ), ) - # Format answer with citation references and confidence + # Format answer with citation references using stable indices answer_text = result.answer if citation_infos: - citation_refs = " ".join(f"[{i + 1}]" for i in range(len(citation_infos))) + citation_refs = " ".join(f"[{c.index}]" for c in citation_infos) answer_text = f"{answer_text}\n\nSources: {citation_refs}" snapshot = new_state.model_dump() diff --git a/haiku_rag_slim/haiku/rag/agents/chat/state.py b/haiku_rag_slim/haiku/rag/agents/chat/state.py index 7cff4129..24ea42a0 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/state.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/state.py @@ -49,6 +49,21 @@ class ChatSessionState(BaseModel): qa_history: list[QAResponse] = [] session_context: SessionContext | None = None document_filter: list[str] = [] + citation_registry: dict[str, int] = {} + + def get_or_assign_index(self, chunk_id: str) -> int: + """Get or assign a stable citation index for a chunk_id. + + Citation indices persist across tool calls within a session. + The first chunk gets index 1, subsequent new chunks get incrementing indices. + Same chunk_id always returns the same index. + """ + if chunk_id in self.citation_registry: + return self.citation_registry[chunk_id] + + new_index = len(self.citation_registry) + 1 + self.citation_registry[chunk_id] = new_index + return new_index @dataclass @@ -103,6 +118,8 @@ class ChatDeps: self.session_state.document_filter = state_data.get( "document_filter", [] ) + if "citation_registry" in state_data: + self.session_state.citation_registry = state_data["citation_registry"] # NOTE: session_context intentionally NOT updated from client # The agent owns this via server-side cache diff --git a/tests/agents/chat/test_chat_agent.py b/tests/agents/chat/test_chat_agent.py index d3b6aeb0..ea3a0057 100644 --- a/tests/agents/chat/test_chat_agent.py +++ b/tests/agents/chat/test_chat_agent.py @@ -677,3 +677,100 @@ async def test_search_agent_with_session_filter(allow_model_requests, temp_db_pa assert "labels" in (r.document_uri or "").lower() or "Labels" in ( r.document_title or "" ) + + +def test_ask_tool_citation_registry_logic(): + """Test the citation index assignment logic used by the ask tool. + + Verifies that: + 1. First chunk gets index 1 + 2. Second unique chunk gets index 2 + 3. Same chunk_id always gets same index + 4. Indices don't reset between calls + """ + session_state = ChatSessionState(session_id="test-registry") + + # Simulate first ask tool building citations + first_ask_chunks = ["chunk-a", "chunk-b"] + first_citations = [] + for chunk_id in first_ask_chunks: + index = session_state.get_or_assign_index(chunk_id) + first_citations.append( + Citation( + index=index, + document_id="doc-1", + chunk_id=chunk_id, + document_uri="test.md", + content="test", + ) + ) + + assert first_citations[0].index == 1 + assert first_citations[1].index == 2 + + # Simulate second ask tool - overlapping chunk_id should keep same index + second_ask_chunks = ["chunk-b", "chunk-c"] # chunk-b was in first ask + second_citations = [] + for chunk_id in second_ask_chunks: + index = session_state.get_or_assign_index(chunk_id) + second_citations.append( + Citation( + index=index, + document_id="doc-1", + chunk_id=chunk_id, + document_uri="test.md", + content="test", + ) + ) + + # chunk-b should have same index as before + assert second_citations[0].index == 2 + # chunk-c is new, gets next index + assert second_citations[1].index == 3 + + # Registry should have all three chunks + assert len(session_state.citation_registry) == 3 + assert session_state.citation_registry == {"chunk-a": 1, "chunk-b": 2, "chunk-c": 3} + + +def test_search_tool_citation_registry_logic(): + """Test the citation index assignment logic used by the search tool. + + Verifies that search and ask tools share the same registry, + maintaining stable indices across different tool calls. + """ + session_state = ChatSessionState(session_id="test-registry") + + # Simulate ask tool first (assigns indices 1, 2) + for chunk_id in ["chunk-a", "chunk-b"]: + session_state.get_or_assign_index(chunk_id) + + # Simulate search tool returning overlapping + new chunks + search_chunks = ["chunk-b", "chunk-c", "chunk-d"] # chunk-b already exists + search_citations = [] + for chunk_id in search_chunks: + index = session_state.get_or_assign_index(chunk_id) + search_citations.append( + Citation( + index=index, + document_id="doc-1", + chunk_id=chunk_id, + document_uri="test.md", + content="test", + ) + ) + + # chunk-b should have same index as assigned by ask (2) + assert search_citations[0].index == 2 + # New chunks get incrementing indices + assert search_citations[1].index == 3 + assert search_citations[2].index == 4 + + # Registry should have all four chunks + assert len(session_state.citation_registry) == 4 + assert session_state.citation_registry == { + "chunk-a": 1, + "chunk-b": 2, + "chunk-c": 3, + "chunk-d": 4, + } diff --git a/tests/agents/chat/test_state.py b/tests/agents/chat/test_state.py index 418b1231..15613479 100644 --- a/tests/agents/chat/test_state.py +++ b/tests/agents/chat/test_state.py @@ -376,6 +376,153 @@ 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.""" + 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") + index1 = session_state.get_or_assign_index("chunk-abc") + index2 = session_state.get_or_assign_index("chunk-def") + assert index1 == 1 + 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 + + +def test_citation_registry_get_or_assign_index_stability(): + """Test citation indices are stable across multiple calls.""" + from haiku.rag.agents.chat.state import ChatSessionState + + session_state = ChatSessionState(session_id="test") + + # First call 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 + 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 + + # New chunk gets next index + idx_d = session_state.get_or_assign_index("chunk-d") + assert idx_d == 4 + + +def test_citation_registry_serialization(): + """Test citation_registry is included in model_dump 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") + + state_dict = session_state.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}, + } + ) + + # 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 + # New chunk should get next index + assert session_state.get_or_assign_index("chunk-c") == 3 + + +def test_chat_deps_state_getter_includes_citation_registry(): + """Test ChatDeps.state getter includes citation_registry.""" + from unittest.mock import MagicMock + + from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState + + mock_client = MagicMock() + mock_config = MagicMock() + + session_state = ChatSessionState(session_id="test") + session_state.get_or_assign_index("chunk-a") + + deps = ChatDeps( + client=mock_client, + config=mock_config, + session_state=session_state, + state_key=AGUI_STATE_KEY, + ) + + state = deps.state + assert state is not None + assert AGUI_STATE_KEY in state + assert state[AGUI_STATE_KEY]["citation_registry"] == {"chunk-a": 1} + + +def test_chat_deps_state_setter_restores_citation_registry(): + """Test ChatDeps.state setter restores citation_registry from incoming state.""" + from unittest.mock import MagicMock + + from haiku.rag.agents.chat.state import AGUI_STATE_KEY, ChatDeps, ChatSessionState + + mock_client = MagicMock() + mock_config = MagicMock() + + session_state = ChatSessionState(session_id="test") + deps = ChatDeps( + client=mock_client, + config=mock_config, + session_state=session_state, + state_key=AGUI_STATE_KEY, + ) + + # Simulate incoming AG-UI state with citation_registry + incoming_state = { + AGUI_STATE_KEY: { + "session_id": "test", + "qa_history": [], + "citations": [], + "citation_registry": {"chunk-x": 1, "chunk-y": 2}, + } + } + + deps.state = incoming_state + + assert deps.session_state is not None + # Registry should be restored + assert deps.session_state.get_or_assign_index("chunk-x") == 1 + assert deps.session_state.get_or_assign_index("chunk-y") == 2 + # New chunk gets next index + assert deps.session_state.get_or_assign_index("chunk-z") == 3 + + def test_chat_deps_state_setter_restores_document_filter(): """Test ChatDeps.state setter restores document_filter from incoming state.""" from unittest.mock import MagicMock