diff --git a/CHANGELOG.md b/CHANGELOG.md index e0d2cccd..abd095be 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,17 @@ - Supports `all` argument to download/upload all datasets at once - Use `--force` flag to overwrite existing databases - Avoids lengthy database rebuild times for users running benchmarks +- **Stable Citation Registry**: Citation indices now persist across tool calls within a session + - Same `chunk_id` always returns the same citation index (first-occurrence-wins) + - New `citation_registry: dict[str, int]` field on `ChatSessionState` + - New `get_or_assign_index(chunk_id)` method for stable index assignment + - Registry serialized/restored via AG-UI state protocol +- **Recall Tool**: Check conversation history before running research + - New `recall` tool on chat agent searches previous Q&A pairs by semantic similarity + - Uses embedding similarity matching with 0.8 cosine similarity threshold + - Returns previous answer with citations if found, avoiding redundant research calls + - Emits `StateSnapshotEvent` so frontend can display recalled citations + - Updated system prompt with routing guidance: use `recall` FIRST for follow-up questions - **Dynamic Session Context**: Compressed conversation history for multi-turn chat - New `SessionContext` model stores summarized conversation state instead of raw Q&A history - Background LLM-based summarization runs after each `ask` tool call (non-blocking) diff --git a/docs/agents.md b/docs/agents.md index 069c0c7d..3e516eef 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -61,12 +61,15 @@ Key features: ### Tools -The chat agent uses three tools: +The chat agent uses four tools: +- `recall` — Search conversation history for previous answers (use FIRST for follow-up questions) - `search` — Hybrid search with optional document filter - `ask` — Answer questions using the conversational research graph - `get_document` — Retrieve a specific document by title or URI +The `recall` tool uses embedding similarity to find semantically matching questions from conversation history. If a match is found (above 0.8 cosine similarity threshold), it returns the previous answer with citations, avoiding redundant research calls. + ### CLI Usage ```bash @@ -104,11 +107,24 @@ The `ChatSessionState` maintains: - `session_id` — Unique identifier for the session - `qa_history` — List of previous Q/A pairs (FIFO, max 50) - `session_context` — Automatically maintained session context summary +- `document_filter` — List of document titles/URIs to restrict searches +- `citation_registry` — Stable mapping of chunk IDs to citation indices + +**Citation Registry**: Citation indices persist across tool calls within a session. The same `chunk_id` always returns the same citation index (first-occurrence-wins). This ensures consistent citation numbering in multi-turn conversations — `[1]` always refers to the same source. + +```python +# Example: citation indices are stable across calls +state = ChatSessionState() + +# First call returns citations [1], [2], [3] +# Second call reuses [1] if same chunk, assigns [4], [5] for new chunks +# User can reference [1] in follow-up and it still refers to original source +``` Q/A history is used to: 1. Provide context for follow-up questions -2. Avoid repeating previous answers +2. Avoid repeating previous answers via the `recall` tool 3. Enable semantic ranking of relevant past answers ### AG-UI Integration @@ -135,7 +151,9 @@ The emitted state structure: "haiku.rag.chat": { "session_id": "", "citations": [...], - "qa_history": [...] + "qa_history": [...], + "document_filter": [...], + "citation_registry": {"chunk-id-1": 1, "chunk-id-2": 2} } } ``` diff --git a/haiku_rag_slim/haiku/rag/agents/chat/agent.py b/haiku_rag_slim/haiku/rag/agents/chat/agent.py index ad072383..ea6e6a9a 100644 --- a/haiku_rag_slim/haiku/rag/agents/chat/agent.py +++ b/haiku_rag_slim/haiku/rag/agents/chat/agent.py @@ -385,7 +385,7 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: async def recall( ctx: RunContext[ChatDeps], topic: str, - ) -> str: + ) -> ToolReturn: """Search conversation history for a previous answer on this topic. Use this FIRST when the user asks about something that may have been @@ -396,11 +396,11 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: topic: The topic or question to search for in conversation history """ if ctx.deps.session_state is None: - return "No conversation history available." + return ToolReturn(return_value="No conversation history available.") qa_history = ctx.deps.session_state.qa_history if not qa_history: - return "No previous answers found." + return ToolReturn(return_value="No previous answers found.") # Get embedder and embed the topic embedder = get_embedder(ctx.deps.config) @@ -421,9 +421,9 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: # Check if similarity exceeds threshold if best_similarity < RECALL_SIMILARITY_THRESHOLD: - return "No previous answer found on this topic." + return ToolReturn(return_value="No previous answer found on this topic.") - # Return the matching answer with citations + # Build result with the matching answer matched_qa = qa_history[best_match_idx] result = f"**Previous answer found** (similarity: {best_similarity:.2f}):\n\n" result += f"**Question:** {matched_qa.question}\n\n" @@ -433,6 +433,31 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]: citation_refs = " ".join(f"[{c.index}]" for c in matched_qa.citations) result += f"Sources: {citation_refs}" - return result + # Emit state with citations so frontend can display them + session_id = ctx.deps.session_state.session_id + new_state = ChatSessionState( + session_id=session_id, + citations=matched_qa.citations, + qa_history=ctx.deps.session_state.qa_history, + session_context=get_cached_session_context(session_id) + if session_id + else None, + document_filter=ctx.deps.session_state.document_filter, + citation_registry=ctx.deps.session_state.citation_registry, + ) + + snapshot = new_state.model_dump() + if ctx.deps.state_key: + snapshot = {ctx.deps.state_key: snapshot} + + return ToolReturn( + return_value=result, + metadata=[ + StateSnapshotEvent( + type=EventType.STATE_SNAPSHOT, + snapshot=snapshot, + ) + ], + ) return agent