From d808c6c4257079f4af4c38ca3dacdc4c9cf57617 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 2 Dec 2025 10:09:05 +0200 Subject: [PATCH] Simplify citations in qa & graph agents --- evaluations/evaluations/benchmark.py | 3 +- haiku_rag_slim/haiku/rag/app.py | 24 ++++-- haiku_rag_slim/haiku/rag/client.py | 15 ++-- haiku_rag_slim/haiku/rag/graph/agui/server.py | 3 +- .../haiku/rag/graph/common/models.py | 56 +++++++++++-- .../haiku/rag/graph/common/nodes.py | 35 ++++----- .../haiku/rag/graph/common/prompts.py | 43 ++++------ .../haiku/rag/graph/deep_qa/dependencies.py | 13 +++- .../haiku/rag/graph/deep_qa/graph.py | 35 +++++---- .../haiku/rag/graph/deep_qa/models.py | 9 ++- .../haiku/rag/graph/deep_qa/prompts.py | 25 +----- .../haiku/rag/graph/research/common.py | 3 +- .../haiku/rag/graph/research/dependencies.py | 12 ++- haiku_rag_slim/haiku/rag/mcp.py | 10 ++- haiku_rag_slim/haiku/rag/qa/__init__.py | 6 +- haiku_rag_slim/haiku/rag/qa/agent.py | 76 ++++++++---------- haiku_rag_slim/haiku/rag/qa/prompts.py | 78 ++++--------------- haiku_rag_slim/haiku/rag/utils.py | 33 +++++++- tests/graph/test_deep_qa.py | 47 +---------- tests/test_app.py | 42 ++++++---- tests/test_client.py | 31 +------- tests/test_mcp.py | 8 +- tests/test_qa.py | 8 +- 23 files changed, 287 insertions(+), 328 deletions(-) diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index e71643d7..754c4ef0 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -248,7 +248,8 @@ async def run_qa_benchmark( qa = get_qa_agent(rag, system_prompt=system_prompt) async def answer_question(question: str) -> str: - return await qa.answer(question) + answer, _ = await qa.answer(question) + return answer eval_name = name if name is not None else f"{spec.key}_qa_evaluation" diff --git a/haiku_rag_slim/haiku/rag/app.py b/haiku_rag_slim/haiku/rag/app.py index eed988e8..47c66bf8 100644 --- a/haiku_rag_slim/haiku/rag/app.py +++ b/haiku_rag_slim/haiku/rag/app.py @@ -29,7 +29,7 @@ from haiku.rag.store.models.document import Document if TYPE_CHECKING: from haiku.rag.store.models import SearchResult -from haiku.rag.utils import format_bytes +from haiku.rag.utils import format_bytes, format_citations logger = logging.getLogger(__name__) @@ -291,39 +291,49 @@ class HaikuRAGApp: """ async with HaikuRAG(db_path=self.db_path, config=self.config) as self.client: try: + citations = [] if deep: from haiku.rag.graph.deep_qa.dependencies import DeepQAContext from haiku.rag.graph.deep_qa.graph import build_deep_qa_graph from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState graph = build_deep_qa_graph(config=self.config) - context = DeepQAContext( - original_question=question, use_citations=cite - ) + context = DeepQAContext(original_question=question) state = DeepQAState.from_config(context=context, config=self.config) deps = DeepQADeps(client=self.client) if verbose: # Use AG-UI renderer to process and display events - from haiku.rag.graph.agui import AGUIConsoleRenderer + from haiku.rag.graph.common.models import Citation renderer = AGUIConsoleRenderer(self.console) result_dict = await renderer.render( stream_graph(graph, state, deps) ) - # Result should be a dict with 'answer' key + # Result should be a dict with 'answer' and 'citations' keys answer = result_dict.get("answer", "") if result_dict else "" + if cite and result_dict: + # Convert dicts to Citation objects + raw_citations = result_dict.get("citations", []) + citations = [ + Citation(**c) if isinstance(c, dict) else c + for c in raw_citations + ] else: # Run without rendering events, just get the result result = await graph.run(state=state, deps=deps) answer = result.answer + if cite: + citations = result.citations else: - answer = await self.client.ask(question, cite=cite) + answer, citations = await self.client.ask(question) self.console.print(f"[bold blue]Question:[/bold blue] {question}") self.console.print() self.console.print("[bold green]Answer:[/bold green]") self.console.print(Markdown(answer)) + if cite and citations: + self.console.print(Markdown(format_citations(citations))) except Exception as e: self.console.print(f"[red]Error: {e}[/red]") diff --git a/haiku_rag_slim/haiku/rag/client.py b/haiku_rag_slim/haiku/rag/client.py index ba33e999..064a8b0e 100644 --- a/haiku_rag_slim/haiku/rag/client.py +++ b/haiku_rag_slim/haiku/rag/client.py @@ -8,6 +8,7 @@ from collections.abc import AsyncGenerator from dataclasses import dataclass from enum import Enum from pathlib import Path +from typing import TYPE_CHECKING from urllib.parse import urlparse import httpx @@ -22,6 +23,9 @@ from haiku.rag.store.repositories.chunk import ChunkRepository from haiku.rag.store.repositories.document import DocumentRepository from haiku.rag.store.repositories.settings import SettingsRepository +if TYPE_CHECKING: + from haiku.rag.graph.common.models import Citation + logger = logging.getLogger(__name__) @@ -850,23 +854,20 @@ class HaikuRAG: return final_results + passthrough async def ask( - self, question: str, cite: bool = False, system_prompt: str | None = None - ) -> str: + self, question: str, system_prompt: str | None = None + ) -> "tuple[str, list[Citation]]": """Ask a question using the configured QA agent. Args: question: The question to ask. - cite: Whether to include citations in the response. system_prompt: Optional custom system prompt for the QA agent. Returns: - The generated answer as a string. + Tuple of (answer text, list of resolved citations). """ from haiku.rag.qa import get_qa_agent - qa_agent = get_qa_agent( - self, config=self._config, use_citations=cite, system_prompt=system_prompt - ) + qa_agent = get_qa_agent(self, config=self._config, system_prompt=system_prompt) return await qa_agent.answer(question) async def rebuild_database( diff --git a/haiku_rag_slim/haiku/rag/graph/agui/server.py b/haiku_rag_slim/haiku/rag/graph/agui/server.py index b569540d..5dadb5a8 100644 --- a/haiku_rag_slim/haiku/rag/graph/agui/server.py +++ b/haiku_rag_slim/haiku/rag/graph/agui/server.py @@ -210,8 +210,7 @@ def create_agui_server(config: "AppConfig", db_path: Path | None = None) -> Star messages = input_state.get("messages", []) if messages: question = messages[0].get("content", "") - use_citations = input_state.get("use_citations", False) - context = DeepQAContext(original_question=question, use_citations=use_citations) + context = DeepQAContext(original_question=question) return DeepQAState.from_config(context=context, config=config) def deep_ask_deps_factory(input_config: dict[str, Any]) -> DeepQADeps: diff --git a/haiku_rag_slim/haiku/rag/graph/common/models.py b/haiku_rag_slim/haiku/rag/graph/common/models.py index 407cb6bb..0cff7057 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/models.py +++ b/haiku_rag_slim/haiku/rag/graph/common/models.py @@ -1,7 +1,14 @@ """Common models used across different graph implementations.""" +from typing import TYPE_CHECKING + from pydantic import BaseModel, Field, field_validator +from haiku.rag.store.models import BoundingBox + +if TYPE_CHECKING: + from haiku.rag.store.models import SearchResult + class ResearchPlan(BaseModel): """A structured research plan with sub-questions to explore.""" @@ -21,18 +28,25 @@ class ResearchPlan(BaseModel): return v +class Citation(BaseModel): + """Resolved citation with full metadata for display/visual grounding.""" + + document_uri: str + document_title: str | None = None + page_numbers: list[int] = Field(default_factory=list) + headings: list[str] | None = None + content: str + bounding_boxes: list[BoundingBox] | None = None + + class SearchAnswer(BaseModel): - """Answer from a search operation with sources.""" + """Structured answer from a search operation.""" query: str = Field(..., description="The question that was answered") - answer: str = Field(..., description="The comprehensive answer to the question") - context: list[str] = Field( + answer: str = Field(..., description="The answer to the question") + cited_chunks: list[str] = Field( default_factory=list, - description="Relevant snippets that directly support the answer", - ) - sources: list[str] = Field( - default_factory=list, - description="Source URIs or titles that contributed to this answer", + description="IDs of chunks used to form the answer", ) confidence: float = Field( default=1.0, @@ -40,3 +54,29 @@ class SearchAnswer(BaseModel): ge=0.0, le=1.0, ) + + +def resolve_citations( + cited_chunk_ids: list[str], + search_results: "list[SearchResult]", +) -> list[Citation]: + """Resolve chunk IDs to full Citation objects with metadata.""" + # Build lookup by chunk_id + by_id = {r.chunk_id: r for r in search_results if r.chunk_id} + + citations = [] + for chunk_id in cited_chunk_ids: + r = by_id.get(chunk_id) + if not r: + continue + citations.append( + Citation( + document_uri=r.document_uri or "", + document_title=r.document_title, + page_numbers=r.page_numbers, + headings=r.headings, + content=r.content, + bounding_boxes=r.bounding_boxes, + ) + ) + return citations diff --git a/haiku_rag_slim/haiku/rag/graph/common/nodes.py b/haiku_rag_slim/haiku/rag/graph/common/nodes.py index 62720abd..13cd7ff2 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/nodes.py +++ b/haiku_rag_slim/haiku/rag/graph/common/nodes.py @@ -5,7 +5,6 @@ from collections.abc import Awaitable, Callable from typing import Any, Protocol from pydantic_ai import Agent, RunContext -from pydantic_ai.format_prompt import format_as_xml from pydantic_ai.output import ToolOutput from pydantic_graph.beta import StepContext @@ -16,6 +15,7 @@ from haiku.rag.graph.agui.emitter import AGUIEmitter from haiku.rag.graph.common import get_model from haiku.rag.graph.common.models import ResearchPlan, SearchAnswer from haiku.rag.graph.common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT +from haiku.rag.store.models import SearchResult class GraphContext(Protocol): @@ -24,7 +24,9 @@ class GraphContext(Protocol): original_question: str sub_questions: list[str] - def add_qa_response(self, qa: SearchAnswer) -> None: + def add_qa_response( + self, qa: SearchAnswer, search_results: list[SearchResult] + ) -> None: """Add a QA response to context.""" ... @@ -49,6 +51,7 @@ class GraphAgentDeps(Protocol): client: HaikuRAG context: GraphContext + search_results: list[SearchResult] def create_plan_node[AgentDepsT: GraphAgentDeps]( @@ -228,28 +231,24 @@ async def _do_search[AgentDepsT: GraphAgentDeps]( async def search_and_answer( ctx2: RunContext[AgentDepsT], query: str, limit: int = 5 ) -> str: + """Search the knowledge base for relevant documents. + + Returns results with chunk IDs and relevance scores. + Reference results by their chunk_id in cited_chunks. + """ results = await ctx2.deps.client.search(query, limit=limit) results = await ctx2.deps.client.expand_context(results) + # Store results for citation resolution + ctx2.deps.search_results = results - entries: list[dict[str, Any]] = [] + parts = [] for r in results: - entry: dict[str, Any] = { - "text": r.content, - "score": r.score, - "document_uri": (r.document_uri or ""), - } - if r.document_title: - entry["document_title"] = r.document_title - if r.page_numbers: - entry["page_numbers"] = r.page_numbers - if r.headings: - entry["headings"] = r.headings - entries.append(entry) + parts.append(f"[{r.chunk_id}] (score: {r.score:.2f}) {r.content}") - if not entries: + if not parts: return f"No relevant information found in the knowledge base for: {query}" - return format_as_xml(entries, root_tag="snippets") + return "\n\n".join(parts) # Tool is registered via decorator above _ = search_and_answer @@ -260,7 +259,7 @@ async def _do_search[AgentDepsT: GraphAgentDeps]( result = await agent.run(sub_q, deps=agent_deps) answer = result.output if answer: - state.context.add_qa_response(answer) + state.context.add_qa_response(answer, agent_deps.search_results) # State updated with new answer - emit state update and narrate if deps.agui_emitter: deps.agui_emitter.update_state(state) diff --git a/haiku_rag_slim/haiku/rag/graph/common/prompts.py b/haiku_rag_slim/haiku/rag/graph/common/prompts.py index 3b497511..8073ba4a 100644 --- a/haiku_rag_slim/haiku/rag/graph/common/prompts.py +++ b/haiku_rag_slim/haiku/rag/graph/common/prompts.py @@ -20,32 +20,21 @@ Plan requirements: SEARCH_AGENT_PROMPT = """You are a search and question-answering specialist. -Tasks: -1. Search the knowledge base for relevant evidence. -2. Analyze retrieved snippets. -3. Provide an answer strictly grounded in that evidence. +Process: +1. Call search_and_answer with relevant keywords from the question. +2. Review the results and their relevance scores. +3. If needed, perform follow-up searches with different keywords (max 3 total). +4. Provide a concise answer based strictly on the retrieved content. -Tool usage: -- Always call search_and_answer before drafting any answer. -- The tool returns snippets with: - - `text`: verbatim content - - `score`: relevance score - - `document_uri`: full path to the source document - - `document_title`: title if available - - `page_numbers`: list of page numbers where content appears (if available) - - `headings`: section heading hierarchy (if available) -- You may call the tool multiple times to refine or broaden context, but do not - exceed 3 total calls. Favor precision over volume. -- Use scores to prioritize evidence, but include only the minimal subset of - snippet texts (verbatim) in SearchAnswer.context (typically 1-4). -- Set SearchAnswer.sources to include document_uri, page numbers, and headings - for each snippet used. Format: "document_uri (p. X, Section: Y)" or just - "document_uri" if no page/heading info. One source per context snippet. -- If no relevant information is found, clearly say so and return an empty - context list and sources list. +The search tool returns results like: +[chunk_abc123] (score: 0.85) Content text here... +[chunk_def456] (score: 0.72) More content... -Answering rules: -- Be direct and specific; avoid meta commentary about the process. -- Do not include any claims not supported by the provided snippets. -- Prefer concise phrasing; avoid copying long passages. -- When evidence is partial, state the limits explicitly in the answer.""" +In your response, include the chunk IDs you used in cited_chunks. + +Guidelines: +- Base answers strictly on retrieved content - do not use external knowledge. +- If multiple results are relevant, synthesize them coherently. +- If information is insufficient, say so clearly. +- Be concise and direct; avoid meta commentary about the process. +- Higher scores indicate more relevant results.""" diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py index 1c33262b..0238b927 100644 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py +++ b/haiku_rag_slim/haiku/rag/graph/deep_qa/dependencies.py @@ -2,6 +2,7 @@ from pydantic import BaseModel, Field from haiku.rag.client import HaikuRAG from haiku.rag.graph.common.models import SearchAnswer +from haiku.rag.store.models import SearchResult class DeepQAContext(BaseModel): @@ -12,11 +13,12 @@ class DeepQAContext(BaseModel): qa_responses: list[SearchAnswer] = Field( default_factory=list, description="QA pairs collected during answering" ) - use_citations: bool = Field( - default=False, description="Whether to include citations in the answer" - ) - def add_qa_response(self, qa: SearchAnswer) -> None: + def add_qa_response( + self, qa: SearchAnswer, search_results: list[SearchResult] + ) -> None: + """Add a QA response.""" + del search_results # Not needed with chunk ID-based citations self.qa_responses.append(qa) @@ -25,3 +27,6 @@ class DeepQADependencies(BaseModel): client: HaikuRAG = Field(description="RAG client for document operations") context: DeepQAContext = Field(description="Shared QA context") + search_results: list[SearchResult] = Field( + default_factory=list, description="Search results for citation resolution" + ) diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py index 567c0cf3..1df7b34f 100644 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py +++ b/haiku_rag_slim/haiku/rag/graph/deep_qa/graph.py @@ -6,16 +6,13 @@ from pydantic_graph.beta.join import reduce_list_append from haiku.rag.config import Config from haiku.rag.config.models import AppConfig from haiku.rag.graph.common import get_model -from haiku.rag.graph.common.models import SearchAnswer +from haiku.rag.graph.common.models import SearchAnswer, resolve_citations from haiku.rag.graph.common.nodes import create_plan_node, create_search_node from haiku.rag.graph.deep_qa.dependencies import DeepQADependencies from haiku.rag.graph.deep_qa.models import DeepQAAnswer, DeepQAEvaluation -from haiku.rag.graph.deep_qa.prompts import ( - DECISION_PROMPT, - SYNTHESIS_PROMPT, - SYNTHESIS_PROMPT_WITH_CITATIONS, -) +from haiku.rag.graph.deep_qa.prompts import DECISION_PROMPT, SYNTHESIS_PROMPT from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState +from haiku.rag.store.models import SearchResult def build_deep_qa_graph( @@ -102,7 +99,7 @@ def build_deep_qa_graph( { "question": qa.query, "answer": qa.answer, - "sources": qa.sources, + "confidence": qa.confidence, } for qa in state.context.qa_responses ], @@ -163,16 +160,10 @@ def build_deep_qa_graph( ) try: - prompt_template = ( - SYNTHESIS_PROMPT_WITH_CITATIONS - if state.context.use_citations - else SYNTHESIS_PROMPT - ) - agent = Agent( model=get_model(model_config, config), - output_type=DeepQAAnswer, - instructions=prompt_template, + output_type=SearchAnswer, + instructions=SYNTHESIS_PROMPT, retries=3, deps_type=DeepQADependencies, ) @@ -183,7 +174,8 @@ def build_deep_qa_graph( { "question": qa.query, "answer": qa.answer, - "sources": qa.sources, + "confidence": qa.confidence, + "cited_chunks": qa.cited_chunks, } for qa in state.context.qa_responses ], @@ -197,13 +189,22 @@ def build_deep_qa_graph( context=state.context, ) result = await agent.run(prompt, deps=agent_deps) + llm_answer = result.output + + # Resolve citations by fetching chunks by ID + search_results = [] + for chunk_id in llm_answer.cited_chunks: + chunk = await deps.client.chunk_repository.get_by_id(chunk_id) + if chunk: + search_results.append(SearchResult.from_chunk(chunk, score=1.0)) + citations = resolve_citations(llm_answer.cited_chunks, search_results) if deps.agui_emitter: deps.agui_emitter.update_activity( "synthesizing", {"message": "Answer complete"} ) - return result.output + return DeepQAAnswer(answer=llm_answer.answer, citations=citations) finally: if deps.agui_emitter: deps.agui_emitter.finish_step() diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/models.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/models.py index e33dfc95..ba4d2b89 100644 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/models.py +++ b/haiku_rag_slim/haiku/rag/graph/deep_qa/models.py @@ -1,5 +1,7 @@ from pydantic import BaseModel, Field +from haiku.rag.graph.common.models import Citation + class DeepQAEvaluation(BaseModel): is_sufficient: bool = Field( @@ -13,8 +15,9 @@ class DeepQAEvaluation(BaseModel): class DeepQAAnswer(BaseModel): + """Final deep QA answer with resolved citations.""" + answer: str = Field(description="The comprehensive answer to the question") - sources: list[str] = Field( - description="Document titles or URIs used to generate the answer", - default_factory=list, + citations: list[Citation] = Field( + default_factory=list, description="Resolved citations for the answer" ) diff --git a/haiku_rag_slim/haiku/rag/graph/deep_qa/prompts.py b/haiku_rag_slim/haiku/rag/graph/deep_qa/prompts.py index 75f27fd7..56d4d4d9 100644 --- a/haiku_rag_slim/haiku/rag/graph/deep_qa/prompts.py +++ b/haiku_rag_slim/haiku/rag/graph/deep_qa/prompts.py @@ -10,34 +10,15 @@ Task: Output format: - answer: The complete answer to the original question (2-4 paragraphs) -- sources: List of document titles/URIs used (extract from the sub-answers) +- cited_chunks: List of chunk IDs (from sub-answers) that directly support your answer Guidelines: - Start directly with the answer - no preamble like "Based on the research..." - Use a clear, professional tone - Organize information logically - If evidence is incomplete, state limitations clearly -- Do not include any claims not supported by the gathered information""" - -SYNTHESIS_PROMPT_WITH_CITATIONS = """You are an expert at synthesizing information into clear, concise answers with proper citations. - -Task: -- Combine the gathered information from sub-questions into a single comprehensive answer -- Answer the original question directly and completely -- Base your answer strictly on the provided evidence -- Include inline citations using [Source Title] format - -Output format: -- answer: The complete answer with inline citations (2-4 paragraphs) -- sources: List of document titles/URIs used (extract from the sub-answers) - -Guidelines: -- Start directly with the answer - no preamble like "Based on the research..." -- Add citations after each claim: [Source Title] -- Use a clear, professional tone -- Organize information logically -- If evidence is incomplete, state limitations clearly -- Do not include any claims not supported by the gathered information""" +- Do not include any claims not supported by the gathered information +- Each sub-answer includes cited_chunks IDs - include the relevant ones in your response""" DECISION_PROMPT = """You are an expert at evaluating whether gathered information is sufficient to answer a question. diff --git a/haiku_rag_slim/haiku/rag/graph/research/common.py b/haiku_rag_slim/haiku/rag/graph/research/common.py index 28c1aada..a9d96c86 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/common.py +++ b/haiku_rag_slim/haiku/rag/graph/research/common.py @@ -14,8 +14,7 @@ def format_context_for_prompt(context: ResearchContext) -> str: { "question": qa.query, "answer": qa.answer, - "context_snippets": qa.context, - "sources": qa.sources, # pyright: ignore[reportAttributeAccessIssue] + "confidence": qa.confidence, } for qa in context.qa_responses ], diff --git a/haiku_rag_slim/haiku/rag/graph/research/dependencies.py b/haiku_rag_slim/haiku/rag/graph/research/dependencies.py index 79a92cae..d1b5ddaa 100644 --- a/haiku_rag_slim/haiku/rag/graph/research/dependencies.py +++ b/haiku_rag_slim/haiku/rag/graph/research/dependencies.py @@ -9,6 +9,7 @@ from haiku.rag.graph.research.models import ( InsightAnalysis, InsightRecord, ) +from haiku.rag.store.models import SearchResult class ResearchContext(BaseModel): @@ -37,8 +38,12 @@ class ResearchContext(BaseModel): self._insights_by_id = {ins.id: ins for ins in self.insights} self._gaps_by_id = {gap.id: gap for gap in self.gaps} - def add_qa_response(self, qa: SearchAnswer) -> None: - """Add a structured QA response (minimal context already included).""" + def add_qa_response( + self, qa: SearchAnswer, search_results: list[SearchResult] + ) -> None: + """Add a structured QA response.""" + # Research doesn't accumulate search_results for citation resolution + del search_results self.qa_responses.append(qa) def upsert_insights(self, records: Iterable[InsightRecord]) -> list[InsightRecord]: @@ -144,6 +149,9 @@ class ResearchDependencies(BaseModel): client: HaikuRAG = Field(description="RAG client for document operations") context: ResearchContext = Field(description="Shared research context") + search_results: list[SearchResult] = Field( + default_factory=list, description="Search results for citation resolution" + ) def _merge_unique(existing: list[str], incoming: Iterable[str]) -> list[str]: diff --git a/haiku_rag_slim/haiku/rag/mcp.py b/haiku_rag_slim/haiku/rag/mcp.py index 90dd30ba..676fa82d 100644 --- a/haiku_rag_slim/haiku/rag/mcp.py +++ b/haiku_rag_slim/haiku/rag/mcp.py @@ -8,6 +8,7 @@ from haiku.rag.client import HaikuRAG from haiku.rag.config import AppConfig, Config from haiku.rag.graph.research.models import ResearchReport from haiku.rag.store.models import SearchResult +from haiku.rag.utils import format_citations class DocumentResult(BaseModel): @@ -176,16 +177,17 @@ def create_mcp_server(db_path: Path, config: AppConfig = Config) -> FastMCP: from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState graph = build_deep_qa_graph(config=config) - context = DeepQAContext( - original_question=question, use_citations=cite - ) + context = DeepQAContext(original_question=question) state = DeepQAState.from_config(context=context, config=config) deps = DeepQADeps(client=rag) result = await graph.run(state=state, deps=deps) answer = result.answer + citations = result.citations else: - answer = await rag.ask(question, cite=cite) + answer, citations = await rag.ask(question) + if cite and citations: + answer += "\n\n" + format_citations(citations) return answer except Exception as e: return f"Error answering question: {e!s}" diff --git a/haiku_rag_slim/haiku/rag/qa/__init__.py b/haiku_rag_slim/haiku/rag/qa/__init__.py index faf1fca8..1ec754d6 100644 --- a/haiku_rag_slim/haiku/rag/qa/__init__.py +++ b/haiku_rag_slim/haiku/rag/qa/__init__.py @@ -6,16 +6,13 @@ from haiku.rag.qa.agent import QuestionAnswerAgent def get_qa_agent( client: HaikuRAG, config: AppConfig = Config, - use_citations: bool = False, system_prompt: str | None = None, ) -> QuestionAnswerAgent: - """ - Factory function to get a QA agent based on the configuration. + """Factory function to get a QA agent based on the configuration. Args: client: HaikuRAG client instance. config: Configuration to use. Defaults to global Config. - use_citations: Whether to include citations in responses. system_prompt: Optional custom system prompt. Returns: @@ -24,6 +21,5 @@ def get_qa_agent( return QuestionAnswerAgent( client=client, model_config=config.qa.model, - use_citations=use_citations, system_prompt=system_prompt, ) diff --git a/haiku_rag_slim/haiku/rag/qa/agent.py b/haiku_rag_slim/haiku/rag/qa/agent.py index 513c2b99..b4d7f8e0 100644 --- a/haiku_rag_slim/haiku/rag/qa/agent.py +++ b/haiku_rag_slim/haiku/rag/qa/agent.py @@ -1,33 +1,24 @@ -from pydantic import BaseModel, Field +from pydantic import BaseModel from pydantic_ai import Agent, RunContext +from pydantic_ai.output import ToolOutput from haiku.rag.client import HaikuRAG from haiku.rag.config import Config from haiku.rag.config.models import ModelConfig from haiku.rag.graph.common import get_model -from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIONS - - -class ToolSearchResult(BaseModel): - """Search result model exposed to the LLM tool.""" - - content: str = Field(description="The document text content") - score: float = Field(description="Relevance score (higher is more relevant)") - document_uri: str = Field(description="The URI/path of the source document") - document_title: str | None = Field( - default=None, description="The title of the document (if available)" - ) - page_numbers: list[int] = Field( - default=[], description="Page numbers where this content appears" - ) - headings: list[str] | None = Field( - default=None, description="Section heading hierarchy for this content" - ) +from haiku.rag.graph.common.models import ( + Citation, + SearchAnswer, + resolve_citations, +) +from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT +from haiku.rag.store.models import SearchResult class Dependencies(BaseModel): model_config = {"arbitrary_types_allowed": True} client: HaikuRAG + search_results: list[SearchResult] = [] class QuestionAnswerAgent: @@ -35,22 +26,16 @@ class QuestionAnswerAgent: self, client: HaikuRAG, model_config: ModelConfig, - use_citations: bool = False, - q: float = 0.0, system_prompt: str | None = None, ): self._client = client - - if system_prompt is None: - system_prompt = ( - QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT - ) model_obj = get_model(model_config, Config) self._agent = Agent( model=model_obj, deps_type=Dependencies, - system_prompt=system_prompt, + output_type=ToolOutput(SearchAnswer, max_retries=3), + instructions=system_prompt or QA_SYSTEM_PROMPT, retries=3, ) @@ -59,24 +44,29 @@ class QuestionAnswerAgent: ctx: RunContext[Dependencies], query: str, limit: int = 5, - ) -> list[ToolSearchResult]: - """Search the knowledge base for relevant documents.""" + ) -> str: + """Search the knowledge base for relevant documents. + + Returns results with chunk IDs and relevance scores. + Reference results by their chunk_id in cited_chunks. + """ results = await ctx.deps.client.search(query, limit=limit) results = await ctx.deps.client.expand_context(results) - return [ - ToolSearchResult( - content=r.content, - score=r.score, - document_uri=(r.document_uri or ""), - document_title=r.document_title, - page_numbers=r.page_numbers, - headings=r.headings, - ) - for r in results - ] + # Store results for citation resolution + ctx.deps.search_results = results + # Format with chunk IDs + parts = [] + for r in results: + parts.append(f"[{r.chunk_id}] (score: {r.score:.2f}) {r.content}") + return "\n\n".join(parts) if parts else "No results found." - async def answer(self, question: str) -> str: - """Answer a question using the RAG system.""" + async def answer(self, question: str) -> tuple[str, list[Citation]]: + """Answer a question using the RAG system. + + Returns: + Tuple of (answer text, list of resolved citations) + """ deps = Dependencies(client=self._client) result = await self._agent.run(question, deps=deps) - return result.output + citations = resolve_citations(result.output.cited_chunks, deps.search_results) + return result.output.answer, citations diff --git a/haiku_rag_slim/haiku/rag/qa/prompts.py b/haiku_rag_slim/haiku/rag/qa/prompts.py index 57aec4b3..edc70cbc 100644 --- a/haiku_rag_slim/haiku/rag/qa/prompts.py +++ b/haiku_rag_slim/haiku/rag/qa/prompts.py @@ -1,65 +1,21 @@ -QA_SYSTEM_PROMPT = """ -You are a knowledgeable assistant that helps users find information from a document knowledge base. +QA_SYSTEM_PROMPT = """You are a knowledgeable assistant that answers questions using a document knowledge base. -Your process: -1. When a user asks a question, use the search_documents tool to find relevant information -2. Search with specific keywords and phrases from the user's question -3. Review the search results and their relevance scores -4. If you need additional context, perform follow-up searches with different keywords -5. Provide a short and to the point comprehensive answer based only on the retrieved documents +Process: +1. Call search_documents with relevant keywords from the question +2. Review the results and their relevance scores +3. If needed, perform follow-up searches with different keywords (max 3 total) +4. Provide a concise answer based strictly on the retrieved content + +The search tool returns results like: +[chunk_abc123] (score: 0.85) Content text here... +[chunk_def456] (score: 0.72) More content... + +In your response, include the chunk IDs you used in cited_chunks. Guidelines: -- Base your answers strictly on the provided document content -- Quote or reference specific information when possible -- If multiple documents contain relevant information, synthesize them coherently -- Indicate when information is incomplete or when you need to search for additional context -- If the retrieved documents don't contain sufficient information, clearly state: "I cannot find enough information in the knowledge base to answer this question." -- For complex questions, consider breaking them down and performing multiple searches -- Stick to the answer, do not ellaborate or provide context unless explicitly asked for it. - -Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents. -/no_think -""" - -QA_SYSTEM_PROMPT_WITH_CITATIONS = """ -You are a knowledgeable assistant that helps users find information from a document knowledge base. - -IMPORTANT: You MUST use the search_documents tool for every question. Do not answer any question without first searching the knowledge base. - -Your process: -1. IMMEDIATELY call the search_documents tool with relevant keywords from the user's question -2. Review the search results and their relevance scores -3. If you need additional context, perform follow-up searches with different keywords -4. Provide a short and to the point comprehensive answer based only on the retrieved documents -5. Always include citations for the sources used in your answer - -Guidelines: -- Base your answers strictly on the provided document content -- If multiple documents contain relevant information, synthesize them coherently -- Indicate when information is incomplete or when you need to search for additional context -- If the retrieved documents don't contain sufficient information, clearly state: "I cannot find enough information in the knowledge base to answer this question." -- For complex questions, consider breaking them down and performing multiple searches -- Stick to the answer, do not ellaborate or provide context unless explicitly asked for it. -- ALWAYS include citations at the end of your response using the format below - -Citation Format: -After your answer, include a "Citations:" section that lists: -- The document URI (from the document_uri field) - always include the full path -- The document title if available (from the document_title field) -- Page number(s) if available (from the page_numbers field) -- Section heading if available (from the headings field) -- A VERBATIM excerpt (100-200 characters, copy-paste exact text) from the content field - do NOT summarize, truncate, or paraphrase - -Example response format: -[Your answer here] - -Citations: - -- **/docs/user-guide.pdf** - "User Manual" (p. 5, Section: Introduction) - The system requires Python 3.10 or higher to run properly... - -- **/reports/quarterly-analysis.md** (pp. 12-13, Section: Results) - Revenue increased by 15% compared to the previous quarter... - -Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents. +- Base answers strictly on retrieved content - do not use external knowledge +- If multiple results are relevant, synthesize them coherently +- If information is insufficient, say: "I cannot find enough information in the knowledge base to answer this question." +- Be concise and direct - avoid elaboration unless asked +- Higher scores indicate more relevant results """ diff --git a/haiku_rag_slim/haiku/rag/utils.py b/haiku_rag_slim/haiku/rag/utils.py index 2eec29b6..899ababc 100644 --- a/haiku_rag_slim/haiku/rag/utils.py +++ b/haiku_rag_slim/haiku/rag/utils.py @@ -4,10 +4,13 @@ import sys from importlib import metadata from pathlib import Path from types import ModuleType -from typing import Any +from typing import TYPE_CHECKING, Any from packaging.version import Version, parse +if TYPE_CHECKING: + from haiku.rag.graph.common.models import Citation + def apply_common_settings( settings: Any | None, @@ -271,6 +274,34 @@ def format_bytes(num_bytes: int) -> str: return f"{size:.1f} PB" +def format_citations(citations: "list[Citation]") -> str: + """Format citations as markdown string.""" + if not citations: + return "" + lines = ["## Citations\n"] + for c in citations: + # Build citation header + parts = [f"- **{c.document_uri}**"] + if c.document_title: + parts.append(f' - "{c.document_title}"') + location_parts = [] + if c.page_numbers: + if len(c.page_numbers) == 1: + location_parts.append(f"p. {c.page_numbers[0]}") + else: + location_parts.append(f"pp. {c.page_numbers[0]}-{c.page_numbers[-1]}") + if c.headings: + location_parts.append(f"Section: {c.headings[-1]}") + if location_parts: + parts.append(f" ({', '.join(location_parts)})") + lines.append("".join(parts)) + # Add truncated content excerpt + excerpt = c.content[:500] + "…" if len(c.content) > 500 else c.content + excerpt = excerpt.replace("\r\n", " ").replace("\n", " ").replace("\r", " ") + lines.append(f"\n {excerpt}\n") + return "\n".join(lines) + + def get_default_data_dir() -> Path: """Get the user data directory for the current system platform. diff --git a/tests/graph/test_deep_qa.py b/tests/graph/test_deep_qa.py index 2bcaa221..05288ec6 100644 --- a/tests/graph/test_deep_qa.py +++ b/tests/graph/test_deep_qa.py @@ -25,9 +25,7 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path): graph = build_deep_qa_graph() state = DeepQAState( - context=DeepQAContext( - original_question="What is haiku.rag?", use_citations=False - ), + context=DeepQAContext(original_question="What is haiku.rag?"), max_sub_questions=3, ) @@ -40,42 +38,6 @@ async def test_deep_qa_graph_end_to_end(monkeypatch, temp_db_path): # TestModel will generate valid structured output based on schemas assert result.answer is not None assert isinstance(result.answer, str) - assert isinstance(result.sources, list) - - client.close() - - -@pytest.mark.asyncio -async def test_deep_qa_with_citations(monkeypatch, temp_db_path): - """Test deep Q&A with citations enabled using TestModel.""" - - # Mock get_model to return TestModel - def test_model_factory(provider, model, config=None): - return TestModel() - - # Patch all locations where get_model is imported - monkeypatch.setattr("haiku.rag.utils.get_model", test_model_factory) - monkeypatch.setattr("haiku.rag.graph.common.get_model", test_model_factory) - monkeypatch.setattr("haiku.rag.graph.common.nodes.get_model", test_model_factory) - monkeypatch.setattr("haiku.rag.graph.deep_qa.graph.get_model", test_model_factory) - - graph = build_deep_qa_graph() - - state = DeepQAState( - context=DeepQAContext(original_question="What is Python?", use_citations=True), - max_sub_questions=2, - ) - - # Use real client but with TestModel for LLM calls - client = HaikuRAG(temp_db_path, create=True) - deps = DeepQADeps(client=client) - - result = await graph.run(state=state, deps=deps) - - # Verify citations flag was used - assert state.context.use_citations is True - assert result.answer is not None - assert isinstance(result.sources, list) client.close() @@ -87,7 +49,6 @@ async def test_deep_qa_context_operations(): assert context.original_question == "Test question?" assert context.sub_questions == [] assert context.qa_responses == [] - assert context.use_citations is False context.sub_questions = ["Sub Q1", "Sub Q2"] assert len(context.sub_questions) == 2 @@ -95,10 +56,10 @@ async def test_deep_qa_context_operations(): qa = SearchAnswer( query="Sub Q1", answer="Answer 1", - context=["Context 1"], - sources=["source1.md"], + cited_chunks=["chunk_1", "chunk_2"], + confidence=0.9, ) - context.add_qa_response(qa) + context.add_qa_response(qa, search_results=[]) assert len(context.qa_responses) == 1 assert context.qa_responses[0].query == "Sub Q1" diff --git a/tests/test_app.py b/tests/test_app.py index e189c022..357ba874 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -293,8 +293,9 @@ async def test_serve_all_services(app: HaikuRAGApp, monkeypatch): async def test_ask_without_cite(app: HaikuRAGApp, monkeypatch): """Test asking a question without citations.""" mock_answer = "Test answer" + mock_citations = [] mock_client = AsyncMock() - mock_client.ask.return_value = mock_answer + mock_client.ask.return_value = (mock_answer, mock_citations) mock_client.__aenter__.return_value = mock_client mock_print = MagicMock() @@ -303,15 +304,25 @@ async def test_ask_without_cite(app: HaikuRAGApp, monkeypatch): with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): await app.ask("test question") - mock_client.ask.assert_called_once_with("test question", cite=False) + mock_client.ask.assert_called_once_with("test question") @pytest.mark.asyncio async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch): """Test asking a question with citations.""" + from haiku.rag.graph.common.models import Citation + mock_answer = "Test answer with citations" + mock_citations = [ + Citation( + document_uri="test.md", + document_title="Test Document", + page_numbers=[1], + content="Test content", + ) + ] mock_client = AsyncMock() - mock_client.ask.return_value = mock_answer + mock_client.ask.return_value = (mock_answer, mock_citations) mock_client.__aenter__.return_value = mock_client mock_print = MagicMock() @@ -320,15 +331,18 @@ async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch): with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): await app.ask("test question", cite=True) - mock_client.ask.assert_called_once_with("test question", cite=True) + mock_client.ask.assert_called_once_with("test question") + # Verify print was called (once for answer, once for citations) + assert mock_print.call_count >= 1 @pytest.mark.asyncio async def test_ask_with_verbose(app: HaikuRAGApp, monkeypatch): """Test asking a question with verbose (should be ignored for non-deep).""" mock_answer = "Test answer" + mock_citations = [] mock_client = AsyncMock() - mock_client.ask.return_value = mock_answer + mock_client.ask.return_value = (mock_answer, mock_citations) mock_client.__aenter__.return_value = mock_client mock_print = MagicMock() @@ -337,7 +351,7 @@ async def test_ask_with_verbose(app: HaikuRAGApp, monkeypatch): with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): await app.ask("test question", verbose=True) - mock_client.ask.assert_called_once_with("test question", cite=False) + mock_client.ask.assert_called_once_with("test question") @pytest.mark.asyncio @@ -345,7 +359,7 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch): """Test asking a question with deep QA.""" from haiku.rag.graph.deep_qa.models import DeepQAAnswer - mock_output = DeepQAAnswer(answer="Deep QA answer", sources=["test.md"]) + mock_output = DeepQAAnswer(answer="Deep QA answer") mock_graph = AsyncMock() mock_graph.run.return_value = mock_output @@ -365,17 +379,14 @@ async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch): mock_graph.run.assert_called_once() call_kwargs = mock_graph.run.call_args[1] assert call_kwargs["state"].context.original_question == "test question" - assert call_kwargs["state"].context.use_citations is False @pytest.mark.asyncio async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch): - """Test asking a question with deep QA and citations.""" + """Test asking a question with deep QA and citations (cite ignored for deep).""" from haiku.rag.graph.deep_qa.models import DeepQAAnswer - mock_output = DeepQAAnswer( - answer="Deep QA answer with citations [test.md]", sources=["test.md"] - ) + mock_output = DeepQAAnswer(answer="Deep QA answer") mock_graph = AsyncMock() mock_graph.run.return_value = mock_output @@ -395,14 +406,13 @@ async def test_ask_with_deep_and_cite(app: HaikuRAGApp, monkeypatch): mock_graph.run.assert_called_once() call_kwargs = mock_graph.run.call_args[1] assert call_kwargs["state"].context.original_question == "test question" - assert call_kwargs["state"].context.use_citations is True @pytest.mark.asyncio async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch): """Test asking a question with deep QA and verbose output.""" - mock_output = {"answer": "Deep QA answer", "sources": ["test.md"]} + mock_output = {"answer": "Deep QA answer", "citations": []} mock_renderer = AsyncMock() mock_renderer.render.return_value = mock_output @@ -419,9 +429,7 @@ async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch): with patch( "haiku.rag.graph.deep_qa.graph.build_deep_qa_graph", return_value=mock_graph ): - with patch( - "haiku.rag.graph.agui.AGUIConsoleRenderer", return_value=mock_renderer - ): + with patch("haiku.rag.app.AGUIConsoleRenderer", return_value=mock_renderer): await app.ask("test question", deep=True, verbose=True) # With verbose, it should use AGUIConsoleRenderer.render, not graph.run diff --git a/tests/test_client.py b/tests/test_client.py index b583633f..2578bb85 100644 --- a/tests/test_client.py +++ b/tests/test_client.py @@ -732,8 +732,8 @@ async def test_client_create_document_with_custom_chunks(temp_db_path): @pytest.mark.asyncio -async def test_client_ask_without_cite(monkeypatch, temp_db_path): - """Test asking questions without citations.""" +async def test_client_ask(monkeypatch, temp_db_path): + """Test asking questions returns answer and citations.""" from pydantic_ai.models.test import TestModel # Mock get_model to return TestModel @@ -748,35 +748,12 @@ async def test_client_ask_without_cite(monkeypatch, temp_db_path): ) # Use real QA agent with TestModel - answer = await client.ask("What is Python?") - - # TestModel will generate a valid string response - assert answer is not None - assert isinstance(answer, str) - - -@pytest.mark.asyncio -async def test_client_ask_with_cite(monkeypatch, temp_db_path): - """Test asking questions with citations.""" - from pydantic_ai.models.test import TestModel - - # Mock get_model to return TestModel - monkeypatch.setattr( - "haiku.rag.utils.get_model", lambda *args, **kwargs: TestModel() - ) - - async with HaikuRAG(temp_db_path, create=True) as client: - # Create a test document - await client.create_document( - content="Python is a high-level programming language.", uri="test.txt" - ) - - # Use real QA agent with TestModel - answer = await client.ask("What is Python?", cite=True) + answer, citations = await client.ask("What is Python?") # TestModel will generate a valid string response assert answer is not None assert isinstance(answer, str) + assert isinstance(citations, list) @pytest.mark.asyncio diff --git a/tests/test_mcp.py b/tests/test_mcp.py index 8c60b558..8ddf0d7a 100644 --- a/tests/test_mcp.py +++ b/tests/test_mcp.py @@ -45,7 +45,7 @@ async def test_mcp_ask_question(): with patch("haiku.rag.mcp.HaikuRAG") as mock_rag_class: mock_rag = AsyncMock() - mock_rag.ask = AsyncMock(return_value="This is the answer") + mock_rag.ask = AsyncMock(return_value=("This is the answer", [])) mock_rag_class.return_value.__aenter__ = AsyncMock(return_value=mock_rag) mock_rag_class.return_value.__aexit__ = AsyncMock(return_value=None) @@ -57,7 +57,7 @@ async def test_mcp_ask_question(): ) assert result == "This is the answer" - mock_rag.ask.assert_called_once_with("What is this?", cite=False) + mock_rag.ask.assert_called_once_with("What is this?") @pytest.mark.asyncio @@ -261,14 +261,16 @@ async def test_mcp_ask_question_deep(): mock_graph = AsyncMock() mock_result = AsyncMock() mock_result.answer = "Deep answer" + mock_result.citations = [] mock_graph.run = AsyncMock(return_value=mock_result) mock_graph_builder.return_value = mock_graph tools = await mcp.get_tools() ask_tool = next(t for t in tools.values() if t.name == "ask_question") + # cite=False to avoid citation formatting in output result = await ask_tool.fn( # type: ignore[attr-defined] - question="Deep question?", cite=True, deep=True + question="Deep question?", cite=False, deep=True ) assert result == "Deep answer" diff --git a/tests/test_qa.py b/tests/test_qa.py index 113f5219..9473c648 100644 --- a/tests/test_qa.py +++ b/tests/test_qa.py @@ -31,7 +31,7 @@ async def test_qa_ollama(qa_corpus: Dataset, temp_db_path): question = doc["question"] expected_answer = doc["answer"] - answer = await qa.answer(question) + answer, _ = await qa.answer(question) is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer) assert is_equivalent, ( @@ -55,7 +55,7 @@ async def test_qa_openai(qa_corpus: Dataset, temp_db_path): question = doc["question"] expected_answer = doc["answer"] - answer = await qa.answer(question) + answer, _ = await qa.answer(question) is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer) assert is_equivalent, ( @@ -81,7 +81,7 @@ async def test_qa_anthropic(qa_corpus: Dataset, temp_db_path): question = doc["question"] expected_answer = doc["answer"] - answer = await qa.answer(question) + answer, _ = await qa.answer(question) is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer) assert is_equivalent, ( @@ -104,7 +104,7 @@ async def test_qa_vllm(qa_corpus: Dataset, temp_db_path): question = doc["question"] expected_answer = doc["answer"] - answer = await qa.answer(question) + answer, _ = await qa.answer(question) is_equivalent = await llm_judge.judge_answers(question, answer, expected_answer) assert is_equivalent, (