diff --git a/README.md b/README.md index 23c1c63e..c97c15c7 100644 --- a/README.md +++ b/README.md @@ -40,6 +40,12 @@ 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 + +# Deep QA with verbose output +haiku-rag ask "Who is the author of haiku.rag?" --deep --verbose + # 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..9e9c5372 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,80 @@ 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 + DeepQAPlanNode --> DeepQASearchDispatchNode + DeepQASearchDispatchNode --> DeepQADecisionNode + DeepQADecisionNode --> DeepQASearchDispatchNode + DeepQADecisionNode --> DeepQASynthesizeNode + DeepQASynthesizeNode --> [*] +``` + +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 DeepQAPlanNode +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=DeepQAPlanNode(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..033faa5a 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -92,7 +92,17 @@ 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 +``` + +Show verbose output with deep QA: +```bash +haiku-rag ask "What are the main features and architecture of haiku.rag?" --deep --verbose +``` + +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. With `--verbose` (only with `--deep`), you'll see the planning, searching, evaluation, and synthesis steps as they happen. 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..ebf18515 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -194,10 +194,44 @@ 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, + verbose: 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 rich.console import Console + + 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 DeepQAPlanNode + 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, console=Console() if verbose else None + ) + + start_node = DeepQAPlanNode( + 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..62ecf660 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -299,11 +299,21 @@ def ask( "--cite", help="Include citations in the response", ), + deep: bool = typer.Option( + False, + "--deep", + help="Use deep multi-agent QA for complex questions", + ), + verbose: bool = typer.Option( + False, + "--verbose", + help="Show verbose progress output (only with --deep)", + ), ): 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, verbose=verbose)) @cli.command("research", help="Run multi-agent research and output a concise report") diff --git a/src/haiku/rag/graph/__init__.py b/src/haiku/rag/graph/__init__.py new file mode 100644 index 00000000..18013005 --- /dev/null +++ b/src/haiku/rag/graph/__init__.py @@ -0,0 +1 @@ +from haiku.rag.graph.models import ResearchPlan, SearchAnswer diff --git a/src/haiku/rag/graph/base.py b/src/haiku/rag/graph/base.py new file mode 100644 index 00000000..1d78fa7a --- /dev/null +++ b/src/haiku/rag/graph/base.py @@ -0,0 +1,31 @@ +from typing import Protocol, runtime_checkable + +from pydantic import BaseModel, Field +from rich.console import Console + +from haiku.rag.client import HaikuRAG +from haiku.rag.graph.models import SearchAnswer + + +@runtime_checkable +class GraphContext(Protocol): + """Protocol for graph context objects.""" + + original_question: str + sub_questions: list[str] + qa_responses: list[SearchAnswer] + + def add_qa_response(self, qa: SearchAnswer) -> None: ... + + +class BaseGraphDeps(BaseModel): + """Base dependencies for graph nodes.""" + + model_config = {"arbitrary_types_allowed": True} + + client: HaikuRAG = Field(description="RAG client for document operations") + console: Console | None = None + + def emit_log(self, message: str) -> None: + if self.console: + self.console.print(message) diff --git a/src/haiku/rag/graph/common.py b/src/haiku/rag/graph/common.py new file mode 100644 index 00000000..2b2b721a --- /dev/null +++ b/src/haiku/rag/graph/common.py @@ -0,0 +1,33 @@ +from typing import Any, Protocol + +from pydantic_ai.models.openai import OpenAIChatModel +from pydantic_ai.providers.ollama import OllamaProvider +from pydantic_ai.providers.openai import OpenAIProvider + +from haiku.rag.config import Config + + +class HasEmitLog(Protocol): + def emit_log(self, message: str, state: Any = None) -> None: ... + + +def get_model(provider: str, model: str) -> Any: + if provider == "ollama": + return OpenAIChatModel( + model_name=model, + provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), + ) + elif provider == "vllm": + return OpenAIChatModel( + model_name=model, + provider=OpenAIProvider( + base_url=f"{Config.VLLM_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1", + api_key="none", + ), + ) + else: + return f"{provider}:{model}" + + +def log(deps: HasEmitLog, state: Any, message: str) -> None: + deps.emit_log(message, state) diff --git a/src/haiku/rag/graph/models.py b/src/haiku/rag/graph/models.py new file mode 100644 index 00000000..475b6e43 --- /dev/null +++ b/src/haiku/rag/graph/models.py @@ -0,0 +1,24 @@ +from pydantic import BaseModel, Field + + +class ResearchPlan(BaseModel): + main_question: str + sub_questions: list[str] + + +class SearchAnswer(BaseModel): + query: str = Field(description="The search query that was performed") + answer: str = Field(description="The answer generated based on the context") + context: list[str] = Field( + description=( + "Only the minimal set of relevant snippets (verbatim) that directly " + "support the answer" + ) + ) + sources: list[str] = Field( + description=( + "Document titles (if available) or URIs corresponding to the" + " snippets actually used in the answer (one per snippet; omit if none)" + ), + default_factory=list, + ) diff --git a/src/haiku/rag/graph/nodes/__init__.py b/src/haiku/rag/graph/nodes/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/src/haiku/rag/research/nodes/analysis.py b/src/haiku/rag/graph/nodes/analysis.py similarity index 96% rename from src/haiku/rag/research/nodes/analysis.py rename to src/haiku/rag/graph/nodes/analysis.py index 18057089..1d29d935 100644 --- a/src/haiku/rag/research/nodes/analysis.py +++ b/src/haiku/rag/graph/nodes/analysis.py @@ -3,15 +3,13 @@ from dataclasses import dataclass from pydantic_ai import Agent from pydantic_graph import BaseNode, GraphRunContext +from haiku.rag.graph.common import get_model, log from haiku.rag.research.common import ( format_analysis_for_prompt, format_context_for_prompt, - get_model, - log, ) from haiku.rag.research.dependencies import ResearchDependencies from haiku.rag.research.models import EvaluationResult, InsightAnalysis, ResearchReport -from haiku.rag.research.nodes.synthesize import SynthesizeNode from haiku.rag.research.prompts import DECISION_AGENT_PROMPT, INSIGHT_AGENT_PROMPT from haiku.rag.research.state import ResearchDeps, ResearchState @@ -89,6 +87,8 @@ class AnalyzeInsightsNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]) for question in analysis.new_questions: log(deps, state, f" • {question}") + from haiku.rag.graph.nodes.analysis import DecisionNode + return DecisionNode(self.provider, self.model) @@ -169,7 +169,8 @@ class DecisionNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]" log(deps, state, f" Sufficient: {status}") - from haiku.rag.research.nodes.search import SearchDispatchNode + from haiku.rag.graph.nodes.search import SearchDispatchNode + from haiku.rag.graph.nodes.synthesize import SynthesizeNode if ( output.is_sufficient diff --git a/src/haiku/rag/research/nodes/plan.py b/src/haiku/rag/graph/nodes/plan.py similarity index 89% rename from src/haiku/rag/research/nodes/plan.py rename to src/haiku/rag/graph/nodes/plan.py index 63612a55..5f2cade9 100644 --- a/src/haiku/rag/research/nodes/plan.py +++ b/src/haiku/rag/graph/nodes/plan.py @@ -3,11 +3,11 @@ from dataclasses import dataclass from pydantic_ai import Agent, RunContext from pydantic_graph import BaseNode, GraphRunContext -from haiku.rag.research.common import get_model, log +from haiku.rag.graph.common import get_model, log +from haiku.rag.graph.models import ResearchPlan +from haiku.rag.graph.prompts import PLAN_PROMPT from haiku.rag.research.dependencies import ResearchDependencies -from haiku.rag.research.models import ResearchPlan, ResearchReport -from haiku.rag.research.nodes.search import SearchDispatchNode -from haiku.rag.research.prompts import PLAN_PROMPT +from haiku.rag.research.models import ResearchReport from haiku.rag.research.state import ResearchDeps, ResearchState @@ -67,4 +67,6 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): for i, sq in enumerate(state.context.sub_questions, 1): log(deps, state, f" {i}. {sq}") + from haiku.rag.graph.nodes.search import SearchDispatchNode + return SearchDispatchNode(self.provider, self.model) diff --git a/src/haiku/rag/research/nodes/search.py b/src/haiku/rag/graph/nodes/search.py similarity index 92% rename from src/haiku/rag/research/nodes/search.py rename to src/haiku/rag/graph/nodes/search.py index 664a01cc..b7e12bd3 100644 --- a/src/haiku/rag/research/nodes/search.py +++ b/src/haiku/rag/graph/nodes/search.py @@ -7,10 +7,11 @@ from pydantic_ai.format_prompt import format_as_xml from pydantic_ai.output import ToolOutput from pydantic_graph import BaseNode, GraphRunContext -from haiku.rag.research.common import get_model, log +from haiku.rag.graph.common import get_model, log +from haiku.rag.graph.models import SearchAnswer +from haiku.rag.graph.prompts import SEARCH_AGENT_PROMPT from haiku.rag.research.dependencies import ResearchDependencies -from haiku.rag.research.models import ResearchReport, SearchAnswer -from haiku.rag.research.prompts import SEARCH_AGENT_PROMPT +from haiku.rag.research.models import ResearchReport from haiku.rag.research.state import ResearchDeps, ResearchState @@ -25,7 +26,7 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): state = ctx.state deps = ctx.deps if not state.context.sub_questions: - from haiku.rag.research.nodes.analysis import AnalyzeInsightsNode + from haiku.rag.graph.nodes.analysis import AnalyzeInsightsNode return AnalyzeInsightsNode(self.provider, self.model) diff --git a/src/haiku/rag/research/nodes/synthesize.py b/src/haiku/rag/graph/nodes/synthesize.py similarity index 90% rename from src/haiku/rag/research/nodes/synthesize.py rename to src/haiku/rag/graph/nodes/synthesize.py index 4f8eee13..309a933d 100644 --- a/src/haiku/rag/research/nodes/synthesize.py +++ b/src/haiku/rag/graph/nodes/synthesize.py @@ -3,10 +3,9 @@ from dataclasses import dataclass from pydantic_ai import Agent from pydantic_graph import BaseNode, End, GraphRunContext -from haiku.rag.research.common import format_context_for_prompt, get_model, log -from haiku.rag.research.dependencies import ( - ResearchDependencies, -) +from haiku.rag.graph.common import get_model, log +from haiku.rag.research.common import format_context_for_prompt +from haiku.rag.research.dependencies import ResearchDependencies from haiku.rag.research.models import ResearchReport from haiku.rag.research.prompts import SYNTHESIS_AGENT_PROMPT from haiku.rag.research.state import ResearchDeps, ResearchState diff --git a/src/haiku/rag/graph/prompts.py b/src/haiku/rag/graph/prompts.py new file mode 100644 index 00000000..97f42aa8 --- /dev/null +++ b/src/haiku/rag/graph/prompts.py @@ -0,0 +1,45 @@ +PLAN_PROMPT = """You are the research orchestrator for a focused, iterative +workflow. + +Responsibilities: +1. Understand and decompose the main question +2. Propose a minimal, high‑leverage plan +3. Coordinate specialized agents to gather evidence +4. Iterate based on gaps and new findings + +Plan requirements: +- Produce at most 3 sub_questions that together cover the main question. +- Each sub_question must be a standalone, self‑contained query that can run + without extra context. Include concrete entities, scope, timeframe, and any + qualifiers. Avoid ambiguous pronouns (it/they/this/that). +- Prioritize the highest‑value aspects first; avoid redundancy and overlap. +- Prefer questions that are likely answerable from the current knowledge base; + if coverage is uncertain, make scopes narrower and specific. +- Order sub_questions by execution priority (most valuable first).""" + +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. + +Tool usage: +- Always call search_and_answer before drafting any answer. +- The tool returns snippets with verbatim `text`, a relevance `score`, and the + originating document identifier (document title if available, otherwise URI). +- 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 the corresponding document identifiers for the + snippets you used (title if available, otherwise URI; one per snippet; same + order as context). Context must be text‑only. +- If no relevant information is found, clearly say so and return an empty + context list and sources list. + +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.""" diff --git a/src/haiku/rag/qa/deep/__init__.py b/src/haiku/rag/qa/deep/__init__.py new file mode 100644 index 00000000..0dc102f6 --- /dev/null +++ b/src/haiku/rag/qa/deep/__init__.py @@ -0,0 +1 @@ +from haiku.rag.qa.deep.models import DeepQAAnswer 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..f5701ebd --- /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 DeepQAAnswer +from haiku.rag.qa.deep.nodes import ( + DeepQADecisionNode, + DeepQAPlanNode, + DeepQASearchDispatchNode, + DeepQASynthesizeNode, +) +from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState + + +def build_deep_qa_graph() -> Graph[DeepQAState, DeepQADeps, DeepQAAnswer]: + return Graph( + nodes=[ + DeepQAPlanNode, + DeepQASearchDispatchNode, + DeepQADecisionNode, + DeepQASynthesizeNode, + ] + ) diff --git a/src/haiku/rag/qa/deep/models.py b/src/haiku/rag/qa/deep/models.py new file mode 100644 index 00000000..e33dfc95 --- /dev/null +++ b/src/haiku/rag/qa/deep/models.py @@ -0,0 +1,20 @@ +from pydantic import BaseModel, Field + + +class DeepQAEvaluation(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 DeepQAAnswer(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..9461efe2 --- /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 DeepQAAnswer, DeepQAEvaluation +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 DeepQAPlanNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]): + provider: str + model: str + + async def run( + self, ctx: GraphRunContext[DeepQAState, DeepQADeps] + ) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]: + 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 DeepQASearchDispatchNode(self.provider, self.model) + + +@dataclass +class DeepQASearchDispatchNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]): + provider: str + model: str + + async def run( + self, ctx: GraphRunContext[DeepQAState, DeepQADeps] + ) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]: + state = ctx.state + deps = ctx.deps + + if not state.context.sub_questions: + return DeepQADecisionNode(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 DeepQASearchDispatchNode(self.provider, self.model) + + +@dataclass +class DeepQADecisionNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]): + provider: str + model: str + + async def run( + self, ctx: GraphRunContext[DeepQAState, DeepQADeps] + ) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]: + 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=DeepQAEvaluation, + 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 DeepQASynthesizeNode(self.provider, self.model) + + log( + deps, + state, + f"\n[bold cyan]🔄 Starting iteration {state.iterations + 1}...[/bold cyan]", + ) + return DeepQASearchDispatchNode(self.provider, self.model) + + +@dataclass +class DeepQASynthesizeNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]): + provider: str + model: str + + async def run( + self, ctx: GraphRunContext[DeepQAState, DeepQADeps] + ) -> End[DeepQAAnswer]: + 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=DeepQAAnswer, + 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/src/haiku/rag/research/__init__.py b/src/haiku/rag/research/__init__.py index b034748a..e49dd051 100644 --- a/src/haiku/rag/research/__init__.py +++ b/src/haiku/rag/research/__init__.py @@ -1,28 +1,3 @@ +from haiku.rag.graph.models import SearchAnswer from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies -from haiku.rag.research.graph import ( - PlanNode, - ResearchDeps, - ResearchState, - build_research_graph, -) -from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer -from haiku.rag.research.stream import ( - ResearchStateSnapshot, - ResearchStreamEvent, - stream_research_graph, -) - -__all__ = [ - "ResearchDependencies", - "ResearchContext", - "SearchAnswer", - "EvaluationResult", - "ResearchReport", - "ResearchDeps", - "ResearchState", - "PlanNode", - "build_research_graph", - "stream_research_graph", - "ResearchStreamEvent", - "ResearchStateSnapshot", -] +from haiku.rag.research.models import EvaluationResult, ResearchReport diff --git a/src/haiku/rag/research/common.py b/src/haiku/rag/research/common.py index fddddb99..bd6e349d 100644 --- a/src/haiku/rag/research/common.py +++ b/src/haiku/rag/research/common.py @@ -1,39 +1,8 @@ -from typing import TYPE_CHECKING, Any - from pydantic_ai import format_as_xml -from pydantic_ai.models.openai import OpenAIChatModel -from pydantic_ai.providers.ollama import OllamaProvider -from pydantic_ai.providers.openai import OpenAIProvider -from haiku.rag.config import Config from haiku.rag.research.dependencies import ResearchContext from haiku.rag.research.models import InsightAnalysis -if TYPE_CHECKING: # pragma: no cover - from haiku.rag.research.state import ResearchDeps, ResearchState - - -def get_model(provider: str, model: str) -> Any: - if provider == "ollama": - return OpenAIChatModel( - model_name=model, - provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"), - ) - elif provider == "vllm": - return OpenAIChatModel( - model_name=model, - provider=OpenAIProvider( - base_url=f"{Config.VLLM_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1", - api_key="none", - ), - ) - else: - return f"{provider}:{model}" - - -def log(deps: "ResearchDeps", state: "ResearchState", msg: str) -> None: - deps.emit_log(msg, state) - def format_context_for_prompt(context: ResearchContext) -> str: """Format the research context as XML for inclusion in prompts.""" diff --git a/src/haiku/rag/research/dependencies.py b/src/haiku/rag/research/dependencies.py index 2c698fa9..9606c658 100644 --- a/src/haiku/rag/research/dependencies.py +++ b/src/haiku/rag/research/dependencies.py @@ -4,11 +4,11 @@ from pydantic import BaseModel, Field from rich.console import Console from haiku.rag.client import HaikuRAG +from haiku.rag.graph.models import SearchAnswer from haiku.rag.research.models import ( GapRecord, InsightAnalysis, InsightRecord, - SearchAnswer, ) from haiku.rag.research.stream import ResearchStream diff --git a/src/haiku/rag/research/graph.py b/src/haiku/rag/research/graph.py index 7f70895f..422ad1fd 100644 --- a/src/haiku/rag/research/graph.py +++ b/src/haiku/rag/research/graph.py @@ -1,23 +1,12 @@ from pydantic_graph import Graph +from haiku.rag.graph.nodes.analysis import AnalyzeInsightsNode, DecisionNode +from haiku.rag.graph.nodes.plan import PlanNode +from haiku.rag.graph.nodes.search import SearchDispatchNode +from haiku.rag.graph.nodes.synthesize import SynthesizeNode from haiku.rag.research.models import ResearchReport -from haiku.rag.research.nodes.analysis import AnalyzeInsightsNode, DecisionNode -from haiku.rag.research.nodes.plan import PlanNode -from haiku.rag.research.nodes.search import SearchDispatchNode -from haiku.rag.research.nodes.synthesize import SynthesizeNode from haiku.rag.research.state import ResearchDeps, ResearchState -__all__ = [ - "PlanNode", - "SearchDispatchNode", - "AnalyzeInsightsNode", - "DecisionNode", - "SynthesizeNode", - "ResearchState", - "ResearchDeps", - "build_research_graph", -] - def build_research_graph() -> Graph[ResearchState, ResearchDeps, ResearchReport]: return Graph( diff --git a/src/haiku/rag/research/models.py b/src/haiku/rag/research/models.py index eedcb835..aae2f32a 100644 --- a/src/haiku/rag/research/models.py +++ b/src/haiku/rag/research/models.py @@ -131,31 +131,6 @@ class InsightAnalysis(BaseModel): ) -class ResearchPlan(BaseModel): - main_question: str - sub_questions: list[str] - - -class SearchAnswer(BaseModel): - """Structured output for the SearchSpecialist agent.""" - - query: str = Field(description="The search query that was performed") - answer: str = Field(description="The answer generated based on the context") - context: list[str] = Field( - description=( - "Only the minimal set of relevant snippets (verbatim) that directly " - "support the answer" - ) - ) - sources: list[str] = Field( - description=( - "Document titles (if available) or URIs corresponding to the" - " snippets actually used in the answer (one per snippet; omit if none)" - ), - default_factory=list, - ) - - class EvaluationResult(BaseModel): """Result of analysis and evaluation.""" diff --git a/src/haiku/rag/research/prompts.py b/src/haiku/rag/research/prompts.py index 4f32e5fe..5b90696e 100644 --- a/src/haiku/rag/research/prompts.py +++ b/src/haiku/rag/research/prompts.py @@ -1,49 +1,3 @@ -PLAN_PROMPT = """You are the research orchestrator for a focused, iterative -workflow. - -Responsibilities: -1. Understand and decompose the main question -2. Propose a minimal, high‑leverage plan -3. Coordinate specialized agents to gather evidence -4. Iterate based on gaps and new findings - -Plan requirements: -- Produce at most 3 sub_questions that together cover the main question. -- Each sub_question must be a standalone, self‑contained query that can run - without extra context. Include concrete entities, scope, timeframe, and any - qualifiers. Avoid ambiguous pronouns (it/they/this/that). -- Prioritize the highest‑value aspects first; avoid redundancy and overlap. -- Prefer questions that are likely answerable from the current knowledge base; - if coverage is uncertain, make scopes narrower and specific. -- Order sub_questions by execution priority (most valuable first).""" - -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. - -Tool usage: -- Always call search_and_answer before drafting any answer. -- The tool returns snippets with verbatim `text`, a relevance `score`, and the - originating document identifier (document title if available, otherwise URI). -- 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 the corresponding document identifiers for the - snippets you used (title if available, otherwise URI; one per snippet; same - order as context). Context must be text‑only. -- If no relevant information is found, clearly say so and return an empty - context list and sources list. - -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.""" - INSIGHT_AGENT_PROMPT = """You are the insight aggregation specialist for the research workflow. diff --git a/tests/test_app.py b/tests/test_app.py index f596afa0..6cd549c9 100644 --- a/tests/test_app.py +++ b/tests/test_app.py @@ -246,3 +246,111 @@ 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_verbose(app: HaikuRAGApp, monkeypatch): + """Test asking a question with verbose (should be ignored for non-deep).""" + mock_answer = "Test answer" + mock_client = AsyncMock() + mock_client.ask.return_value = mock_answer + 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): + await app.ask("test question", verbose=True) + + mock_client.ask.assert_called_once_with("test question", cite=False) + + +@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 DeepQAAnswer + + mock_output = DeepQAAnswer(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 DeepQAAnswer + + mock_output = DeepQAAnswer( + 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 + + +@pytest.mark.asyncio +async def test_ask_with_deep_and_verbose(app: HaikuRAGApp, monkeypatch): + """Test asking a question with deep QA and verbose output.""" + from haiku.rag.qa.deep.models import DeepQAAnswer + + mock_output = DeepQAAnswer(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, verbose=True) + + mock_graph.run.assert_called_once() + call_kwargs = mock_graph.run.call_args[1] + assert call_kwargs["deps"].console is not None diff --git a/tests/test_cli.py b/tests/test_cli.py index 19ea95a1..72aa361e 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, verbose=False ) @@ -220,7 +220,51 @@ 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, verbose=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, verbose=False + ) + + +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, verbose=False + ) + + +def test_ask_with_deep_and_verbose(): + 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", "--verbose"]) + + assert result.exit_code == 0 + mock_app_instance.ask.assert_called_once_with( + question="What is Python?", cite=False, deep=True, verbose=True + ) def test_info(): diff --git a/tests/test_deep_qa.py b/tests/test_deep_qa.py new file mode 100644 index 00000000..ee9b1a1d --- /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 DeepQAAnswer +from haiku.rag.qa.deep.nodes import ( + DeepQADecisionNode, + DeepQAPlanNode, + DeepQASearchDispatchNode, + DeepQASynthesizeNode, +) +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 DeepQASearchDispatchNode(self.provider, self.model) + + async def fake_search_dispatch_run(self, ctx) -> Any: + if not ctx.state.context.sub_questions: + return DeepQADecisionNode(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 DeepQASearchDispatchNode(self.provider, self.model) + + async def fake_decision_run(self, ctx) -> Any: + ctx.state.iterations += 1 + return DeepQASynthesizeNode(self.provider, self.model) + + async def fake_synthesize_run(self, ctx) -> Any: + from pydantic_graph import End + + return End( + DeepQAAnswer( + answer="haiku.rag is a RAG system with components A, B, C.", + sources=["test.md"], + ) + ) + + monkeypatch.setattr(DeepQAPlanNode, "run", fake_plan_run) + monkeypatch.setattr(DeepQASearchDispatchNode, "run", fake_search_dispatch_run) + monkeypatch.setattr(DeepQADecisionNode, "run", fake_decision_run) + monkeypatch.setattr(DeepQASynthesizeNode, "run", fake_synthesize_run) + + start = DeepQAPlanNode(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 DeepQASearchDispatchNode(self.provider, self.model) + + async def fake_search_dispatch_run(self, ctx) -> Any: + if not ctx.state.context.sub_questions: + return DeepQADecisionNode(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 DeepQASearchDispatchNode(self.provider, self.model) + + async def fake_decision_run(self, ctx) -> Any: + ctx.state.iterations += 1 + return DeepQASynthesizeNode(self.provider, self.model) + + async def fake_synthesize_run(self, ctx) -> Any: + from pydantic_graph import End + + return End( + DeepQAAnswer( + answer="Python is a programming language [python.md].", + sources=["python.md"], + ) + ) + + monkeypatch.setattr(DeepQAPlanNode, "run", fake_plan_run) + monkeypatch.setattr(DeepQASearchDispatchNode, "run", fake_search_dispatch_run) + monkeypatch.setattr(DeepQADecisionNode, "run", fake_decision_run) + monkeypatch.setattr(DeepQASynthesizeNode, "run", fake_synthesize_run) + + start = DeepQAPlanNode(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 diff --git a/tests/test_research_graph_integration.py b/tests/test_research_graph_integration.py index 3f3aef78..e3fdd153 100644 --- a/tests/test_research_graph_integration.py +++ b/tests/test_research_graph_integration.py @@ -2,15 +2,15 @@ from typing import Any, cast import pytest +from haiku.rag.graph.models import SearchAnswer +from haiku.rag.graph.nodes.analysis import AnalyzeInsightsNode, DecisionNode +from haiku.rag.graph.nodes.plan import PlanNode +from haiku.rag.graph.nodes.search import SearchDispatchNode +from haiku.rag.graph.nodes.synthesize import SynthesizeNode from haiku.rag.research.dependencies import ResearchContext from haiku.rag.research.graph import ( - AnalyzeInsightsNode, - DecisionNode, - PlanNode, ResearchDeps, ResearchState, - SearchDispatchNode, - SynthesizeNode, build_research_graph, ) from haiku.rag.research.models import ( @@ -21,7 +21,6 @@ from haiku.rag.research.models import ( InsightRecord, InsightStatus, ResearchReport, - SearchAnswer, ) from haiku.rag.research.stream import stream_research_graph