Update search(), ask() to use get_or_assign_index()

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()

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
 

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,
+    }

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