diff --git a/README.md b/README.md index 23c1c63e..c8361a80 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,9 @@ haiku-rag ask "Who is the author of haiku.rag?" # Ask questions with citations haiku-rag ask "Who is the author of haiku.rag?" --cite +# Deep QA (multi-agent question decomposition) +haiku-rag ask "Who is the author of haiku.rag?" --deep --cite + # Multi‑agent research (iterative plan/search/evaluate) haiku-rag research \ "What are the main drivers and trends of global temperature anomalies since 1990?" \ diff --git a/docs/agents.md b/docs/agents.md index 59c9f8f8..d926258f 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -1,8 +1,9 @@ ## Agents -Two agentic flows are provided by haiku.rag: +Three agentic flows are provided by haiku.rag: - Simple QA Agent — a focused question answering agent +- Deep QA Agent — multi-agent question decomposition for complex questions - Research Multi‑Agent — a multi‑step, analyzable research workflow @@ -37,6 +38,81 @@ answer = await agent.answer("What is climate change?") print(answer) ``` +### Deep QA Agent + +Deep QA is a multi-agent system that decomposes complex questions into sub-questions, answers them in batches, evaluates sufficiency, and iterates if needed before synthesizing a final answer. It's lighter than the full research workflow but more powerful than the simple QA agent. + +```mermaid +--- +title: Deep QA graph +--- +stateDiagram-v2 + DeepPlanNode --> DeepSearchDispatchNode + DeepSearchDispatchNode --> DeepSearchDispatchNode + DeepSearchDispatchNode --> DeepDecisionNode + DeepDecisionNode --> DeepSearchDispatchNode + DeepDecisionNode --> DeepSynthesizeNode + DeepSynthesizeNode --> [*] +``` + +Key nodes: + +- **Plan**: Decomposes the question into focused sub-questions +- **Search (batched)**: Answers sub-questions in parallel batches (respects max_concurrency) +- **Decision**: Evaluates if we have sufficient information or need another iteration +- **Synthesize**: Generates the final comprehensive answer + +Key differences from Research: + +- **Simpler evaluation**: Uses sufficiency check (not confidence + insight analysis) +- **Direct answers**: Returns just the answer (not a full research report) +- **Question-focused**: Optimized for answering specific questions, not open-ended research +- **Supports citations**: Can include inline source citations like `[document.md]` +- **Configurable iterations**: Control max_iterations (default: 2) and max_concurrency (default: 3) + +CLI usage: + +```bash +# Deep QA without citations +haiku-rag ask "What are the main features of haiku.rag?" --deep + +# Deep QA with citations +haiku-rag ask "What are the main features of haiku.rag?" --deep --cite +``` + +Python usage: + +```python +from haiku.rag.client import HaikuRAG +from haiku.rag.qa.deep.dependencies import DeepQAContext +from haiku.rag.qa.deep.graph import build_deep_qa_graph +from haiku.rag.qa.deep.nodes import DeepPlanNode +from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState + +async with HaikuRAG(path_to_db) as client: + graph = build_deep_qa_graph() + context = DeepQAContext( + original_question="What are the main features of haiku.rag?", + use_citations=True + ) + state = DeepQAState( + context=context, + max_sub_questions=3, + max_iterations=2, + max_concurrency=3 + ) + deps = DeepQADeps(client=client) + + result = await graph.run( + start_node=DeepPlanNode(provider="openai", model="gpt-4o-mini"), + state=state, + deps=deps + ) + + print(result.output.answer) + print(result.output.sources) +``` + ### Research Graph The research workflow is implemented as a typed pydantic‑graph. It plans, searches (in parallel batches), evaluates, and synthesizes into a final report — with clear stop conditions and shared state. diff --git a/docs/cli.md b/docs/cli.md index 751d40e7..c21386a5 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -92,7 +92,12 @@ Ask questions with citations showing source documents: haiku-rag ask "Who is the author of haiku.rag?" --cite ``` -The QA agent will search your documents for relevant information and provide a comprehensive answer. With `--cite`, responses include citations showing which documents were used. +Use deep QA for complex questions (multi-agent decomposition): +```bash +haiku-rag ask "What are the main features and architecture of haiku.rag?" --deep --cite +``` + +The QA agent will search your documents for relevant information and provide a comprehensive answer. With `--cite`, responses include citations showing which documents were used. With `--deep`, the question is decomposed into sub-questions that are answered in parallel before synthesizing a final answer. When available, citations use the document title; otherwise they fall back to the URI. ## Research diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 0a0302f9..247ac882 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -194,10 +194,34 @@ class HaikuRAGApp: for chunk, score in results: self._rich_print_search_result(chunk, score) - async def ask(self, question: str, cite: bool = False): + async def ask(self, question: str, cite: bool = False, deep: bool = False): async with HaikuRAG(db_path=self.db_path) as self.client: try: - answer = await self.client.ask(question, cite=cite) + if deep: + from haiku.rag.qa.deep.dependencies import DeepQAContext + from haiku.rag.qa.deep.graph import build_deep_qa_graph + from haiku.rag.qa.deep.nodes import DeepPlanNode + from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState + + graph = build_deep_qa_graph() + context = DeepQAContext( + original_question=question, use_citations=cite + ) + state = DeepQAState(context=context) + deps = DeepQADeps(client=self.client) + + start_node = DeepPlanNode( + provider=Config.QA_PROVIDER, + model=Config.QA_MODEL, + ) + + result = await graph.run( + start_node=start_node, state=state, deps=deps + ) + answer = result.output.answer + else: + answer = await self.client.ask(question, cite=cite) + self.console.print(f"[bold blue]Question:[/bold blue] {question}") self.console.print() self.console.print("[bold green]Answer:[/bold green]") diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index ca8349fe..fc4f7d94 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -299,11 +299,16 @@ def ask( "--cite", help="Include citations in the response", ), + deep: bool = typer.Option( + False, + "--deep", + help="Use deep multi-agent QA for complex questions", + ), ): from haiku.rag.app import HaikuRAGApp app = HaikuRAGApp(db_path=db) - asyncio.run(app.ask(question=question, cite=cite)) + asyncio.run(app.ask(question=question, cite=cite, deep=deep)) @cli.command("research", help="Run multi-agent research and output a concise report") diff --git a/src/haiku/rag/graph/common.py b/src/haiku/rag/graph/common.py index 22009ef8..2b2b721a 100644 --- a/src/haiku/rag/graph/common.py +++ b/src/haiku/rag/graph/common.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Any +from typing import Any, Protocol from pydantic_ai.models.openai import OpenAIChatModel from pydantic_ai.providers.ollama import OllamaProvider @@ -6,8 +6,9 @@ from pydantic_ai.providers.openai import OpenAIProvider from haiku.rag.config import Config -if TYPE_CHECKING: # pragma: no cover - from haiku.rag.research.state import ResearchDeps, ResearchState + +class HasEmitLog(Protocol): + def emit_log(self, message: str, state: Any = None) -> None: ... def get_model(provider: str, model: str) -> Any: @@ -28,5 +29,5 @@ def get_model(provider: str, model: str) -> Any: return f"{provider}:{model}" -def log(deps: "ResearchDeps", state: "ResearchState", msg: str) -> None: - deps.emit_log(msg, state) +def log(deps: HasEmitLog, state: Any, message: str) -> None: + deps.emit_log(message, state) diff --git a/src/haiku/rag/qa/deep/__init__.py b/src/haiku/rag/qa/deep/__init__.py new file mode 100644 index 00000000..38782a45 --- /dev/null +++ b/src/haiku/rag/qa/deep/__init__.py @@ -0,0 +1 @@ +from haiku.rag.qa.deep.models import DeepAnswer diff --git a/src/haiku/rag/qa/deep/dependencies.py b/src/haiku/rag/qa/deep/dependencies.py new file mode 100644 index 00000000..c6017a61 --- /dev/null +++ b/src/haiku/rag/qa/deep/dependencies.py @@ -0,0 +1,29 @@ +from pydantic import BaseModel, Field +from rich.console import Console + +from haiku.rag.client import HaikuRAG +from haiku.rag.graph.models import SearchAnswer + + +class DeepQAContext(BaseModel): + original_question: str = Field(description="The original question") + sub_questions: list[str] = Field( + default_factory=list, description="Decomposed sub-questions" + ) + 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: + self.qa_responses.append(qa) + + +class DeepQADependencies(BaseModel): + model_config = {"arbitrary_types_allowed": True} + + client: HaikuRAG = Field(description="RAG client for document operations") + context: DeepQAContext = Field(description="Shared QA context") + console: Console | None = None diff --git a/src/haiku/rag/qa/deep/graph.py b/src/haiku/rag/qa/deep/graph.py new file mode 100644 index 00000000..a5a90afd --- /dev/null +++ b/src/haiku/rag/qa/deep/graph.py @@ -0,0 +1,21 @@ +from pydantic_graph import Graph + +from haiku.rag.qa.deep.models import DeepAnswer +from haiku.rag.qa.deep.nodes import ( + DeepDecisionNode, + DeepPlanNode, + DeepSearchDispatchNode, + DeepSynthesizeNode, +) +from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState + + +def build_deep_qa_graph() -> Graph[DeepQAState, DeepQADeps, DeepAnswer]: + return Graph( + nodes=[ + DeepPlanNode, + DeepSearchDispatchNode, + DeepDecisionNode, + DeepSynthesizeNode, + ] + ) diff --git a/src/haiku/rag/qa/deep/models.py b/src/haiku/rag/qa/deep/models.py new file mode 100644 index 00000000..933c3616 --- /dev/null +++ b/src/haiku/rag/qa/deep/models.py @@ -0,0 +1,20 @@ +from pydantic import BaseModel, Field + + +class DeepEvaluation(BaseModel): + is_sufficient: bool = Field( + description="Whether we have sufficient information to answer the question" + ) + reasoning: str = Field(description="Explanation of the sufficiency assessment") + new_questions: list[str] = Field( + description="Additional sub-questions needed if insufficient", + default_factory=list, + ) + + +class DeepAnswer(BaseModel): + 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, + ) diff --git a/src/haiku/rag/qa/deep/nodes.py b/src/haiku/rag/qa/deep/nodes.py new file mode 100644 index 00000000..d97284c8 --- /dev/null +++ b/src/haiku/rag/qa/deep/nodes.py @@ -0,0 +1,303 @@ +import asyncio +from dataclasses import dataclass +from typing import Any + +from pydantic_ai import Agent, RunContext +from pydantic_ai.format_prompt import format_as_xml +from pydantic_ai.output import ToolOutput +from pydantic_graph import BaseNode, End, GraphRunContext + +from haiku.rag.graph.common import get_model, log +from haiku.rag.graph.models import ResearchPlan, SearchAnswer +from haiku.rag.graph.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT +from haiku.rag.qa.deep.dependencies import DeepQADependencies +from haiku.rag.qa.deep.models import DeepAnswer, DeepEvaluation +from haiku.rag.qa.deep.prompts import ( + DECISION_PROMPT, + SYNTHESIS_PROMPT, + SYNTHESIS_PROMPT_WITH_CITATIONS, +) +from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState + + +@dataclass +class DeepPlanNode(BaseNode[DeepQAState, DeepQADeps, DeepAnswer]): + provider: str + model: str + + async def run( + self, ctx: GraphRunContext[DeepQAState, DeepQADeps] + ) -> BaseNode[DeepQAState, DeepQADeps, DeepAnswer]: + state = ctx.state + deps = ctx.deps + + log(deps, state, "\n[bold cyan]📋 Planning approach...[/bold cyan]") + + plan_agent = Agent( + model=get_model(self.provider, self.model), + output_type=ResearchPlan, + instructions=( + PLAN_PROMPT + + "\n\nUse the gather_context tool once on the main question before planning." + ), + retries=3, + deps_type=DeepQADependencies, + ) + + @plan_agent.tool + async def gather_context( + ctx2: RunContext[DeepQADependencies], query: str, limit: int = 6 + ) -> str: + results = await ctx2.deps.client.search(query, limit=limit) + expanded = await ctx2.deps.client.expand_context(results) + return "\n\n".join(chunk.content for chunk, _ in expanded) + + prompt = ( + "Plan a focused approach for answering the main question.\n\n" + f"Main question: {state.context.original_question}" + ) + + agent_deps = DeepQADependencies( + client=deps.client, + context=state.context, + console=deps.console, + ) + plan_result = await plan_agent.run(prompt, deps=agent_deps) + state.context.sub_questions = list(plan_result.output.sub_questions)[ + : state.max_sub_questions + ] + + log(deps, state, "\n[bold green]✅ Plan Created:[/bold green]") + log( + deps, + state, + f" [bold]Main Question:[/bold] {state.context.original_question}", + ) + log(deps, state, " [bold]Sub-questions:[/bold]") + for i, sq in enumerate(state.context.sub_questions, 1): + log(deps, state, f" {i}. {sq}") + + return DeepSearchDispatchNode(self.provider, self.model) + + +@dataclass +class DeepSearchDispatchNode(BaseNode[DeepQAState, DeepQADeps, DeepAnswer]): + provider: str + model: str + + async def run( + self, ctx: GraphRunContext[DeepQAState, DeepQADeps] + ) -> BaseNode[DeepQAState, DeepQADeps, DeepAnswer]: + state = ctx.state + deps = ctx.deps + + if not state.context.sub_questions: + return DeepDecisionNode(self.provider, self.model) + + # Take up to max_concurrency questions and answer them concurrently + take = max(1, state.max_concurrency) + batch: list[str] = [] + while state.context.sub_questions and len(batch) < take: + batch.append(state.context.sub_questions.pop(0)) + + async def answer_one(sub_q: str) -> SearchAnswer | None: + log( + deps, + state, + f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}", + ) + agent = Agent( + model=get_model(self.provider, self.model), + output_type=ToolOutput(SearchAnswer, max_retries=3), + instructions=SEARCH_AGENT_PROMPT, + retries=3, + deps_type=DeepQADependencies, + ) + + @agent.tool + async def search_and_answer( + ctx2: RunContext[DeepQADependencies], query: str, limit: int = 5 + ) -> str: + search_results = await ctx2.deps.client.search(query, limit=limit) + expanded = await ctx2.deps.client.expand_context(search_results) + + entries: list[dict[str, Any]] = [ + { + "text": chunk.content, + "score": score, + "document_uri": ( + chunk.document_title or chunk.document_uri or "" + ), + } + for chunk, score in expanded + ] + if not entries: + return f"No relevant information found in the knowledge base for: {query}" + + return format_as_xml(entries, root_tag="snippets") + + agent_deps = DeepQADependencies( + client=deps.client, + context=state.context, + console=deps.console, + ) + try: + result = await agent.run(sub_q, deps=agent_deps) + except Exception as e: + log(deps, state, f"[red]Search failed:[/red] {e}") + return None + + return result.output + + answers = await asyncio.gather(*(answer_one(q) for q in batch)) + for ans in answers: + if ans is None: + continue + state.context.add_qa_response(ans) + preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "") + log(deps, state, f" [green]✓[/green] {preview}") + + return DeepSearchDispatchNode(self.provider, self.model) + + +@dataclass +class DeepDecisionNode(BaseNode[DeepQAState, DeepQADeps, DeepAnswer]): + provider: str + model: str + + async def run( + self, ctx: GraphRunContext[DeepQAState, DeepQADeps] + ) -> BaseNode[DeepQAState, DeepQADeps, DeepAnswer]: + state = ctx.state + deps = ctx.deps + + log( + deps, + state, + "\n[bold cyan]📊 Evaluating information sufficiency...[/bold cyan]", + ) + + agent = Agent( + model=get_model(self.provider, self.model), + output_type=DeepEvaluation, + instructions=DECISION_PROMPT, + retries=3, + deps_type=DeepQADependencies, + ) + + context_data = { + "original_question": state.context.original_question, + "gathered_answers": [ + { + "question": qa.query, + "answer": qa.answer, + "sources": qa.sources, + } + for qa in state.context.qa_responses + ], + } + context_xml = format_as_xml(context_data, root_tag="gathered_information") + + prompt = ( + "Evaluate whether we have sufficient information to answer the question.\n\n" + f"{context_xml}" + ) + + agent_deps = DeepQADependencies( + client=deps.client, + context=state.context, + console=deps.console, + ) + result = await agent.run(prompt, deps=agent_deps) + evaluation = result.output + + state.iterations += 1 + + log(deps, state, f" [bold]Assessment:[/bold] {evaluation.reasoning}") + status = "[green]Yes[/green]" if evaluation.is_sufficient else "[red]No[/red]" + log(deps, state, f" Sufficient: {status}") + + # Add new questions if not sufficient + for new_q in evaluation.new_questions: + if new_q not in state.context.sub_questions: + state.context.sub_questions.append(new_q) + + if evaluation.new_questions: + log(deps, state, " [cyan]New questions:[/cyan]") + for question in evaluation.new_questions: + log(deps, state, f" • {question}") + + # Decide next step + if evaluation.is_sufficient or state.iterations >= state.max_iterations: + if state.iterations >= state.max_iterations: + log( + deps, + state, + f"\n[bold yellow]⚠️ Reached max iterations ({state.max_iterations})[/bold yellow]", + ) + log(deps, state, "\n[bold green]✅ Moving to synthesis.[/bold green]") + return DeepSynthesizeNode(self.provider, self.model) + + log( + deps, + state, + f"\n[bold cyan]🔄 Starting iteration {state.iterations + 1}...[/bold cyan]", + ) + return DeepSearchDispatchNode(self.provider, self.model) + + +@dataclass +class DeepSynthesizeNode(BaseNode[DeepQAState, DeepQADeps, DeepAnswer]): + provider: str + model: str + + async def run( + self, ctx: GraphRunContext[DeepQAState, DeepQADeps] + ) -> End[DeepAnswer]: + state = ctx.state + deps = ctx.deps + + log( + deps, + state, + "\n[bold cyan]📝 Synthesizing final answer...[/bold cyan]", + ) + + prompt_template = ( + SYNTHESIS_PROMPT_WITH_CITATIONS + if state.context.use_citations + else SYNTHESIS_PROMPT + ) + + agent = Agent( + model=get_model(self.provider, self.model), + output_type=DeepAnswer, + instructions=prompt_template, + retries=3, + deps_type=DeepQADependencies, + ) + + context_data = { + "original_question": state.context.original_question, + "sub_answers": [ + { + "question": qa.query, + "answer": qa.answer, + "sources": qa.sources, + } + for qa in state.context.qa_responses + ], + } + context_xml = format_as_xml(context_data, root_tag="gathered_information") + + prompt = f"Synthesize a comprehensive answer to the original question.\n\n{context_xml}" + + agent_deps = DeepQADependencies( + client=deps.client, + context=state.context, + console=deps.console, + ) + result = await agent.run(prompt, deps=agent_deps) + + log(deps, state, "[bold green]✅ Answer complete![/bold green]") + return End(result.output) diff --git a/src/haiku/rag/qa/deep/prompts.py b/src/haiku/rag/qa/deep/prompts.py new file mode 100644 index 00000000..2b0d16b7 --- /dev/null +++ b/src/haiku/rag/qa/deep/prompts.py @@ -0,0 +1,57 @@ +SYNTHESIS_PROMPT = """You are an expert at synthesizing information into clear, concise answers. + +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 +- Be clear, accurate, and well-structured + +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) + +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""" + +DECISION_PROMPT = """You are an expert at evaluating whether gathered information is sufficient to answer a question. + +Task: +- Review the original question and all gathered sub-question answers +- Determine if we have enough information to provide a comprehensive answer +- If insufficient, suggest specific new sub-questions to fill the gaps + +Output format: +- is_sufficient: Boolean indicating if we can answer the question comprehensively +- reasoning: Clear explanation of your assessment +- new_questions: List of specific follow-up questions needed (empty if sufficient) + +Guidelines: +- Be strict but reasonable in your assessment +- Focus on whether core aspects of the question are addressed +- New questions should be specific and distinct from what's been asked +- Limit new questions to 2-3 maximum +- Consider whether additional searches would meaningfully improve the answer""" diff --git a/src/haiku/rag/qa/deep/state.py b/src/haiku/rag/qa/deep/state.py new file mode 100644 index 00000000..f76a99a3 --- /dev/null +++ b/src/haiku/rag/qa/deep/state.py @@ -0,0 +1,25 @@ +from dataclasses import dataclass + +from rich.console import Console + +from haiku.rag.client import HaikuRAG +from haiku.rag.qa.deep.dependencies import DeepQAContext + + +@dataclass +class DeepQADeps: + client: HaikuRAG + console: Console | None = None + + def emit_log(self, message: str, state: "DeepQAState | None" = None) -> None: + if self.console: + self.console.print(message) + + +@dataclass +class DeepQAState: + context: DeepQAContext + max_sub_questions: int = 3 + max_iterations: int = 2 + max_concurrency: int = 3 + iterations: int = 0 diff --git a/tests/test_app.py b/tests/test_app.py index f596afa0..cb6a1f65 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -246,3 +246,65 @@ async def test_ask_with_cite(app: HaikuRAGApp, monkeypatch): await app.ask("test question", cite=True) mock_client.ask.assert_called_once_with("test question", cite=True) + + +@pytest.mark.asyncio +async def test_ask_with_deep(app: HaikuRAGApp, monkeypatch): + """Test asking a question with deep QA.""" + from haiku.rag.qa.deep.models import DeepAnswer + + mock_output = DeepAnswer(answer="Deep QA answer", sources=["test.md"]) + mock_result = MagicMock() + mock_result.output = mock_output + + mock_graph = AsyncMock() + mock_graph.run.return_value = mock_result + + mock_client = AsyncMock() + mock_client.__aenter__.return_value = mock_client + + mock_print = MagicMock() + monkeypatch.setattr(app.console, "print", mock_print) + + with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): + with patch( + "haiku.rag.qa.deep.graph.build_deep_qa_graph", return_value=mock_graph + ): + await app.ask("test question", deep=True) + + 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.""" + from haiku.rag.qa.deep.models import DeepAnswer + + mock_output = DeepAnswer( + answer="Deep QA answer with citations [test.md]", sources=["test.md"] + ) + mock_result = MagicMock() + mock_result.output = mock_output + + mock_graph = AsyncMock() + mock_graph.run.return_value = mock_result + + mock_client = AsyncMock() + mock_client.__aenter__.return_value = mock_client + + mock_print = MagicMock() + monkeypatch.setattr(app.console, "print", mock_print) + + with patch("haiku.rag.app.HaikuRAG", return_value=mock_client): + with patch( + "haiku.rag.qa.deep.graph.build_deep_qa_graph", return_value=mock_graph + ): + await app.ask("test question", deep=True, cite=True) + + 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 diff --git a/tests/test_cli.py b/tests/test_cli.py index 19ea95a1..3d70cde6 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -207,7 +207,7 @@ def test_ask(): assert result.exit_code == 0 mock_app_instance.ask.assert_called_once_with( - question="What is Python?", cite=False + question="What is Python?", cite=False, deep=False ) @@ -220,7 +220,37 @@ def test_ask_with_cite(): result = runner.invoke(cli, ["ask", "What is Python?", "--cite"]) assert result.exit_code == 0 - mock_app_instance.ask.assert_called_once_with(question="What is Python?", cite=True) + mock_app_instance.ask.assert_called_once_with( + question="What is Python?", cite=True, deep=False + ) + + +def test_ask_with_deep(): + with patch("haiku.rag.app.HaikuRAGApp") as mock_app: + mock_app_instance = MagicMock() + mock_app_instance.ask = AsyncMock() + mock_app.return_value = mock_app_instance + + result = runner.invoke(cli, ["ask", "What is Python?", "--deep"]) + + assert result.exit_code == 0 + mock_app_instance.ask.assert_called_once_with( + question="What is Python?", cite=False, deep=True + ) + + +def test_ask_with_deep_and_cite(): + with patch("haiku.rag.app.HaikuRAGApp") as mock_app: + mock_app_instance = MagicMock() + mock_app_instance.ask = AsyncMock() + mock_app.return_value = mock_app_instance + + result = runner.invoke(cli, ["ask", "What is Python?", "--deep", "--cite"]) + + assert result.exit_code == 0 + mock_app_instance.ask.assert_called_once_with( + question="What is Python?", cite=True, deep=True + ) def test_info(): diff --git a/tests/test_deep_qa.py b/tests/test_deep_qa.py new file mode 100644 index 00000000..1f6f8587 --- /dev/null +++ b/tests/test_deep_qa.py @@ -0,0 +1,168 @@ +from typing import Any, cast + +import pytest + +from haiku.rag.graph.models import SearchAnswer +from haiku.rag.qa.deep.dependencies import DeepQAContext +from haiku.rag.qa.deep.graph import build_deep_qa_graph +from haiku.rag.qa.deep.models import DeepAnswer +from haiku.rag.qa.deep.nodes import ( + DeepDecisionNode, + DeepPlanNode, + DeepSearchDispatchNode, + DeepSynthesizeNode, +) +from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState + + +@pytest.mark.asyncio +async def test_deep_qa_graph_end_to_end(monkeypatch): + graph = build_deep_qa_graph() + + state = DeepQAState( + context=DeepQAContext( + original_question="What is haiku.rag?", use_citations=False + ), + max_sub_questions=3, + ) + deps = DeepQADeps(client=cast(Any, None), console=None) + + async def fake_plan_run(self, ctx) -> Any: + ctx.state.context.sub_questions = [ + "Describe haiku.rag in one sentence", + "List core components of haiku.rag", + ] + return DeepSearchDispatchNode(self.provider, self.model) + + async def fake_search_dispatch_run(self, ctx) -> Any: + if not ctx.state.context.sub_questions: + return DeepDecisionNode(self.provider, self.model) + + batch = ctx.state.context.sub_questions[:] + ctx.state.context.sub_questions.clear() + + for question in batch: + ctx.state.context.add_qa_response( + SearchAnswer( + query=question, + answer=f"Answer to: {question}", + context=["Context snippet"], + sources=["test.md"], + ) + ) + return DeepSearchDispatchNode(self.provider, self.model) + + async def fake_decision_run(self, ctx) -> Any: + ctx.state.iterations += 1 + return DeepSynthesizeNode(self.provider, self.model) + + async def fake_synthesize_run(self, ctx) -> Any: + from pydantic_graph import End + + return End( + DeepAnswer( + answer="haiku.rag is a RAG system with components A, B, C.", + sources=["test.md"], + ) + ) + + monkeypatch.setattr(DeepPlanNode, "run", fake_plan_run) + monkeypatch.setattr(DeepSearchDispatchNode, "run", fake_search_dispatch_run) + monkeypatch.setattr(DeepDecisionNode, "run", fake_decision_run) + monkeypatch.setattr(DeepSynthesizeNode, "run", fake_synthesize_run) + + start = DeepPlanNode(provider="ollama", model="test") + result = await graph.run(start_node=start, state=state, deps=deps) + + assert result.output.answer == "haiku.rag is a RAG system with components A, B, C." + assert result.output.sources == ["test.md"] + assert len(state.context.qa_responses) == 2 + + +@pytest.mark.asyncio +async def test_deep_qa_with_citations(monkeypatch): + graph = build_deep_qa_graph() + + state = DeepQAState( + context=DeepQAContext(original_question="What is Python?", use_citations=True), + max_sub_questions=2, + ) + deps = DeepQADeps(client=cast(Any, None), console=None) + + async def fake_plan_run(self, ctx) -> Any: + ctx.state.context.sub_questions = ["What is Python used for?"] + return DeepSearchDispatchNode(self.provider, self.model) + + async def fake_search_dispatch_run(self, ctx) -> Any: + if not ctx.state.context.sub_questions: + return DeepDecisionNode(self.provider, self.model) + + batch = ctx.state.context.sub_questions[:] + ctx.state.context.sub_questions.clear() + + for question in batch: + ctx.state.context.add_qa_response( + SearchAnswer( + query=question, + answer="Python is used for web development and data science.", + context=["Python snippet"], + sources=["python.md"], + ) + ) + return DeepSearchDispatchNode(self.provider, self.model) + + async def fake_decision_run(self, ctx) -> Any: + ctx.state.iterations += 1 + return DeepSynthesizeNode(self.provider, self.model) + + async def fake_synthesize_run(self, ctx) -> Any: + from pydantic_graph import End + + return End( + DeepAnswer( + answer="Python is a programming language [python.md].", + sources=["python.md"], + ) + ) + + monkeypatch.setattr(DeepPlanNode, "run", fake_plan_run) + monkeypatch.setattr(DeepSearchDispatchNode, "run", fake_search_dispatch_run) + monkeypatch.setattr(DeepDecisionNode, "run", fake_decision_run) + monkeypatch.setattr(DeepSynthesizeNode, "run", fake_synthesize_run) + + start = DeepPlanNode(provider="ollama", model="test") + result = await graph.run(start_node=start, state=state, deps=deps) + + assert "[python.md]" in result.output.answer + assert state.context.use_citations is True + + +@pytest.mark.asyncio +async def test_deep_qa_context_operations(): + context = DeepQAContext(original_question="Test question?") + + 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 + + qa = SearchAnswer( + query="Sub Q1", + answer="Answer 1", + context=["Context 1"], + sources=["source1.md"], + ) + context.add_qa_response(qa) + assert len(context.qa_responses) == 1 + assert context.qa_responses[0].query == "Sub Q1" + + +def test_deep_qa_state_initialization(): + context = DeepQAContext(original_question="Test?") + state = DeepQAState(context=context, max_sub_questions=5) + + assert state.context.original_question == "Test?" + assert state.max_sub_questions == 5