From cba8ad769633173928fec296e8915a83defa0145 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 19 Sep 2025 13:40:52 +0300 Subject: [PATCH 1/7] Research agent as a pydantic graph --- docs/agents.md | 71 ++++---- docs/cli.md | 18 ++ pyproject.toml | 2 + src/haiku/rag/app.py | 57 +++++-- src/haiku/rag/cli.py | 12 ++ src/haiku/rag/reranking/mxbai.py | 2 +- src/haiku/rag/research/__init__.py | 37 ++-- src/haiku/rag/research/base.py | 130 -------------- src/haiku/rag/research/common.py | 53 ++++++ src/haiku/rag/research/dependencies.py | 28 +-- src/haiku/rag/research/evaluation_agent.py | 85 --------- src/haiku/rag/research/graph.py | 29 ++++ src/haiku/rag/research/models.py | 70 ++++++++ src/haiku/rag/research/nodes/evaluate.py | 81 +++++++++ src/haiku/rag/research/nodes/plan.py | 64 +++++++ src/haiku/rag/research/nodes/search.py | 92 ++++++++++ src/haiku/rag/research/nodes/synthesize.py | 51 ++++++ src/haiku/rag/research/orchestrator.py | 170 ------------------ src/haiku/rag/research/presearch_agent.py | 39 ----- src/haiku/rag/research/search_agent.py | 69 -------- src/haiku/rag/research/state.py | 25 +++ src/haiku/rag/research/synthesis_agent.py | 60 ------- tests/research/test_evaluation_agent.py | 17 -- tests/research/test_orchestrator.py | 189 --------------------- tests/research/test_search_agent.py | 14 -- tests/research/test_synthesis_agent.py | 14 -- tests/test_research_graph.py | 26 +++ tests/test_research_graph_integration.py | 89 ++++++++++ 28 files changed, 709 insertions(+), 885 deletions(-) delete mode 100644 src/haiku/rag/research/base.py create mode 100644 src/haiku/rag/research/common.py delete mode 100644 src/haiku/rag/research/evaluation_agent.py create mode 100644 src/haiku/rag/research/graph.py create mode 100644 src/haiku/rag/research/models.py create mode 100644 src/haiku/rag/research/nodes/evaluate.py create mode 100644 src/haiku/rag/research/nodes/plan.py create mode 100644 src/haiku/rag/research/nodes/search.py create mode 100644 src/haiku/rag/research/nodes/synthesize.py delete mode 100644 src/haiku/rag/research/orchestrator.py delete mode 100644 src/haiku/rag/research/presearch_agent.py delete mode 100644 src/haiku/rag/research/search_agent.py create mode 100644 src/haiku/rag/research/state.py delete mode 100644 src/haiku/rag/research/synthesis_agent.py delete mode 100644 tests/research/test_evaluation_agent.py delete mode 100644 tests/research/test_orchestrator.py delete mode 100644 tests/research/test_search_agent.py delete mode 100644 tests/research/test_synthesis_agent.py create mode 100644 tests/test_research_graph.py create mode 100644 tests/test_research_graph_integration.py diff --git a/docs/agents.md b/docs/agents.md index 502bc3e5..2d018701 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -36,50 +36,57 @@ answer = await agent.answer("What is climate change?") print(answer) ``` -### Research Multi‑Agent +### Research Graph -The research workflow coordinates specialized agents to plan, search, analyze, and synthesize a comprehensive answer. It is designed for deeper questions that benefit from iterative investigation and structured reporting. +The research workflow is now 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. -Components: +Key nodes: -- Orchestrator: Plans, coordinates, and loops until confidence is sufficient -- Presearch Survey: Runs a quick KB scan and summarizes relevant chunk text to - ground the initial plan (plain-text summary; no URIs or scores) -- Search Specialist: Performs targeted RAG searches and answers sub‑questions -- Analysis & Evaluation: Extracts insights, identifies gaps, proposes new questions -- Synthesis: Produces a final structured research report +- Plan: builds up to 3 standalone sub‑questions (uses an internal presearch tool) +- Search (batched): answers sub‑questions using the KB with minimal, verbatim context +- Evaluate: extracts insights, proposes new questions, and checks sufficiency/confidence +- Synthesize: generates a final structured report Primary models: -- `ResearchPlan` — produced by the orchestrator when planning - - `main_question: str` - - `sub_questions: list[str]` (standalone, self‑contained queries) -- `SearchAnswer` — produced by the search specialist for each sub‑question - - `query: str` — the executed sub‑question - - `answer: str` — the agent’s answer grounded in retrieved context - - `context: list[str]` — minimal verbatim snippets used for the answer - - `sources: list[str]` — document URIs aligned with `context` -- `EvaluationResult` — insights, new standalone questions, sufficiency & confidence -- `ResearchReport` — the final synthesized report +- `SearchAnswer` — one per sub‑question (query, answer, context, sources) +- `EvaluationResult` — insights, new questions, sufficiency, confidence +- `ResearchReport` — final report (title, executive summary, findings, conclusions, …) +CLI usage: + +```bash +haiku-rag research "How does haiku.rag organize and query documents?" \ + --max-iterations 2 \ + --confidence-threshold 0.8 \ + --max-concurrency 3 \ + --verbose +``` Python usage: ```python from haiku.rag.client import HaikuRAG -from haiku.rag.research import ResearchOrchestrator - -client = HaikuRAG(path_to_db) -orchestrator = ResearchOrchestrator(provider="ollama", model="gpt-oss") - -report = await orchestrator.conduct_research( - question="What are the main drivers and recent trends of global temperature anomalies since 1990?", - client=client, - max_iterations=2, - confidence_threshold=0.8, - verbose=True, +from haiku.rag.research import ( + ResearchContext, + ResearchDeps, + ResearchState, + build_research_graph, + PlanNode, ) -print(report.title) -print(report.executive_summary) +async with HaikuRAG(path_to_db) as client: + graph = build_research_graph() + state = ResearchState( + question="What are the main drivers and trends of global temperature anomalies since 1990?", + context=ResearchContext(original_question=... ), + max_iterations=2, + confidence_threshold=0.8, + max_concurrency=3, + ) + deps = ResearchDeps(client=client) + result = await graph.run(PlanNode(provider=None, model=None), state=state, deps=deps) + report = result.output + print(report.title) + print(report.executive_summary) ``` diff --git a/docs/cli.md b/docs/cli.md index 5785580c..cc543d08 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -84,6 +84,24 @@ 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. +## Research + +Run the multi-step research graph: + +```bash +haiku-rag research "How does haiku.rag organize and query documents?" \ + --max-iterations 2 \ + --confidence-threshold 0.8 \ + --max-concurrency 3 \ + --verbose +``` + +Flags: +- `--max-iterations, -n`: maximum search/evaluate cycles (default: 3) +- `--confidence-threshold`: stop once evaluation confidence meets/exceeds this (default: 0.8) +- `--max-concurrency`: number of sub-questions searched in parallel each iteration (default: 3) +- `--verbose`: show planning, searching previews, evaluation summary, and stop reason + ## Server Start the MCP server: diff --git a/pyproject.toml b/pyproject.toml index e3cde554..f696f737 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -29,6 +29,7 @@ dependencies = [ "lancedb>=0.25.0", "pydantic>=2.11.9", "pydantic-ai>=1.0.8", + "pydantic-graph>=1.0.8", "python-dotenv>=1.1.1", "rich>=14.1.0", "tiktoken>=0.11.0", @@ -90,6 +91,7 @@ line-ending = "auto" [tool.pyright] venvPath = "." venv = ".venv" +pythonVersion = "3.12" [tool.pytest.ini_options] asyncio_default_fixture_loop_scope = "session" diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 704b9205..c4988612 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -9,7 +9,13 @@ from haiku.rag.client import HaikuRAG from haiku.rag.config import Config from haiku.rag.mcp import create_mcp_server from haiku.rag.monitor import FileWatcher -from haiku.rag.research.orchestrator import ResearchOrchestrator +from haiku.rag.research.dependencies import ResearchContext +from haiku.rag.research.graph import ( + PlanNode, + ResearchDeps, + ResearchState, + build_research_graph, +) from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.document import Document @@ -80,28 +86,53 @@ class HaikuRAGApp: self.console.print(f"[red]Error: {e}[/red]") async def research( - self, question: str, max_iterations: int = 3, verbose: bool = False + self, + question: str, + max_iterations: int = 3, + confidence_threshold: float = 0.8, + max_concurrency: int = 3, + verbose: bool = False, ): - """Run multi-agent research on a question.""" + """Run research via the pydantic-graph pipeline (default).""" async with HaikuRAG(db_path=self.db_path) as client: try: - # Create orchestrator with default config or fallback to QA - orchestrator = ResearchOrchestrator() - if verbose: - self.console.print( - f"[bold cyan]Starting research with {orchestrator.provider}:{orchestrator.model}[/bold cyan]" - ) + self.console.print("[bold cyan]Starting research[/bold cyan]") self.console.print(f"[bold blue]Question:[/bold blue] {question}") self.console.print() - # Conduct research - report = await orchestrator.conduct_research( + graph = build_research_graph() + state = ResearchState( question=question, - client=client, + context=ResearchContext(original_question=question), max_iterations=max_iterations, - verbose=verbose, + confidence_threshold=confidence_threshold, + max_concurrency=max_concurrency, ) + deps = ResearchDeps( + client=client, console=self.console if verbose else None + ) + + start = PlanNode( + provider=Config.RESEARCH_PROVIDER or Config.QA_PROVIDER, + model=Config.RESEARCH_MODEL or Config.QA_MODEL, + ) + # Prefer graph.run; fall back to iter if unavailable + report = None + try: + result = await graph.run(start, state=state, deps=deps) + report = result.output + except Exception: + from pydantic_graph import End + + async with graph.iter(start, state=state, deps=deps) as run: + node = run.next_node + while not isinstance(node, End): + node = await run.next(node) + if run.result: + report = run.result.output + if report is None: + raise RuntimeError("Graph did not produce a report") # Display the report self.console.print("[bold green]Research Report[/bold green]") diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index d02e56c7..6c4f5fef 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -250,6 +250,16 @@ def research( "-n", help="Maximum search/analyze iterations", ), + confidence_threshold: float = typer.Option( + 0.8, + "--confidence-threshold", + help="Minimum confidence (0-1) to stop", + ), + max_concurrency: int = typer.Option( + 3, + "--max-concurrency", + help="Max concurrent searches per iteration (planned)", + ), db: Path = typer.Option( Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb", "--db", @@ -266,6 +276,8 @@ def research( app.research( question=question, max_iterations=max_iterations, + confidence_threshold=confidence_threshold, + max_concurrency=max_concurrency, verbose=verbose, ) ) diff --git a/src/haiku/rag/reranking/mxbai.py b/src/haiku/rag/reranking/mxbai.py index 032edac5..df39ac4a 100644 --- a/src/haiku/rag/reranking/mxbai.py +++ b/src/haiku/rag/reranking/mxbai.py @@ -1,4 +1,4 @@ -from mxbai_rerank import MxbaiRerankV2 +from mxbai_rerank import MxbaiRerankV2 # pyright: ignore[reportMissingImports] from haiku.rag.config import Config from haiku.rag.reranking.base import RerankerBase diff --git a/src/haiku/rag/research/__init__.py b/src/haiku/rag/research/__init__.py index e1e4f1e4..e48953e4 100644 --- a/src/haiku/rag/research/__init__.py +++ b/src/haiku/rag/research/__init__.py @@ -1,37 +1,20 @@ -"""Multi-agent research workflow for advanced RAG queries.""" - -from haiku.rag.research.base import ( - BaseResearchAgent, - ResearchOutput, - SearchAnswer, - SearchResult, -) from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies -from haiku.rag.research.evaluation_agent import ( - AnalysisEvaluationAgent, - EvaluationResult, +from haiku.rag.research.graph import ( + PlanNode, + ResearchDeps, + ResearchState, + build_research_graph, ) -from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan -from haiku.rag.research.presearch_agent import PresearchSurveyAgent -from haiku.rag.research.search_agent import SearchSpecialistAgent -from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent +from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer __all__ = [ - # Base classes - "BaseResearchAgent", "ResearchDependencies", "ResearchContext", - "SearchResult", - "ResearchOutput", - # Specialized agents "SearchAnswer", - "SearchSpecialistAgent", - "PresearchSurveyAgent", - "AnalysisEvaluationAgent", "EvaluationResult", - "SynthesisAgent", "ResearchReport", - # Orchestrator - "ResearchOrchestrator", - "ResearchPlan", + "ResearchDeps", + "ResearchState", + "PlanNode", + "build_research_graph", ] diff --git a/src/haiku/rag/research/base.py b/src/haiku/rag/research/base.py deleted file mode 100644 index a2ce6ad8..00000000 --- a/src/haiku/rag/research/base.py +++ /dev/null @@ -1,130 +0,0 @@ -from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any - -from pydantic import BaseModel, Field -from pydantic_ai import Agent -from pydantic_ai.models.openai import OpenAIChatModel -from pydantic_ai.output import ToolOutput -from pydantic_ai.providers.ollama import OllamaProvider -from pydantic_ai.providers.openai import OpenAIProvider -from pydantic_ai.run import AgentRunResult - -from haiku.rag.config import Config - -if TYPE_CHECKING: - from haiku.rag.research.dependencies import ResearchDependencies - - -class BaseResearchAgent[T](ABC): - """Base class for all research agents.""" - - def __init__( - self, - provider: str, - model: str, - output_type: type[T], - ): - self.provider = provider - self.model = model - self.output_type = output_type - - model_obj = self._get_model(provider, model) - - # Import deps type lazily to avoid circular import during module load - from haiku.rag.research.dependencies import ResearchDependencies - - # If the agent is expected to return plain text, pass `str` directly. - # Otherwise, wrap the model with ToolOutput for robust tool-handling retries. - agent_output_type: Any - if self.output_type is str: # plain text output - agent_output_type = str - else: - agent_output_type = ToolOutput(self.output_type, max_retries=3) - - self._agent = Agent( - model=model_obj, - deps_type=ResearchDependencies, - output_type=agent_output_type, - instructions=self.get_system_prompt(), - retries=3, - ) - - # Register tools - self.register_tools() - - def _get_model(self, provider: str, model: str): - """Get the appropriate model object for the provider.""" - 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: - # For all other providers, use the provider:model format - return f"{provider}:{model}" - - @abstractmethod - def get_system_prompt(self) -> str: - """Return the system prompt for this agent.""" - pass - - def register_tools(self) -> None: - """Register agent-specific tools.""" - pass - - async def run( - self, prompt: str, deps: "ResearchDependencies", **kwargs - ) -> AgentRunResult[T]: - """Execute the agent.""" - return await self._agent.run(prompt, deps=deps, **kwargs) - - @property - def agent(self) -> Agent[Any, T]: - """Access the underlying Pydantic AI agent.""" - return self._agent - - -class SearchResult(BaseModel): - """Standard search result format.""" - - content: str - score: float - document_uri: str - metadata: dict[str, Any] = Field(default_factory=dict) - - -class ResearchOutput(BaseModel): - """Standard research output format.""" - - summary: str - detailed_findings: list[str] - sources: list[str] - confidence: float - - -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 URIs corresponding to the snippets actually used in the" - " answer (one URI per snippet; omit if none)" - ), - default_factory=list, - ) diff --git a/src/haiku/rag/research/common.py b/src/haiku/rag/research/common.py new file mode 100644 index 00000000..4c821d3c --- /dev/null +++ b/src/haiku/rag/research/common.py @@ -0,0 +1,53 @@ +from typing import 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 + + +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(console, msg: str) -> None: + if console: + console.print(msg) + + +def format_context_for_prompt(context: ResearchContext) -> str: + """Format the research context as XML for inclusion in prompts.""" + + context_data = { + "original_question": context.original_question, + "unanswered_questions": context.sub_questions, + "qa_responses": [ + { + "question": qa.query, + "answer": qa.answer, + "context_snippets": qa.context, + "sources": qa.sources, # pyright: ignore[reportAttributeAccessIssue] + } + for qa in context.qa_responses + ], + "insights": context.insights, + "gaps": context.gaps, + } + return format_as_xml(context_data, root_tag="research_context") diff --git a/src/haiku/rag/research/dependencies.py b/src/haiku/rag/research/dependencies.py index ad1ed935..c075f852 100644 --- a/src/haiku/rag/research/dependencies.py +++ b/src/haiku/rag/research/dependencies.py @@ -1,9 +1,8 @@ from pydantic import BaseModel, Field -from pydantic_ai import format_as_xml from rich.console import Console from haiku.rag.client import HaikuRAG -from haiku.rag.research.base import SearchAnswer +from haiku.rag.research.models import SearchAnswer class ResearchContext(BaseModel): @@ -13,7 +12,7 @@ class ResearchContext(BaseModel): sub_questions: list[str] = Field( default_factory=list, description="Decomposed sub-questions" ) - qa_responses: list["SearchAnswer"] = Field( + qa_responses: list[SearchAnswer] = Field( default_factory=list, description="Structured QA pairs used during research" ) insights: list[str] = Field( @@ -23,7 +22,7 @@ class ResearchContext(BaseModel): default_factory=list, description="Identified information gaps" ) - def add_qa_response(self, qa: "SearchAnswer") -> None: + def add_qa_response(self, qa: SearchAnswer) -> None: """Add a structured QA response (minimal context already included).""" self.qa_responses.append(qa) @@ -46,24 +45,3 @@ class ResearchDependencies(BaseModel): client: HaikuRAG = Field(description="RAG client for document operations") context: ResearchContext = Field(description="Shared research context") console: Console | None = None - - -def _format_context_for_prompt(context: ResearchContext) -> str: - """Format the research context as XML for inclusion in prompts.""" - - context_data = { - "original_question": context.original_question, - "unanswered_questions": context.sub_questions, - "qa_responses": [ - { - "question": qa.query, - "answer": qa.answer, - "context_snippets": qa.context, - "sources": qa.sources, - } - for qa in context.qa_responses - ], - "insights": context.insights, - "gaps": context.gaps, - } - return format_as_xml(context_data, root_tag="research_context") diff --git a/src/haiku/rag/research/evaluation_agent.py b/src/haiku/rag/research/evaluation_agent.py deleted file mode 100644 index 0867349a..00000000 --- a/src/haiku/rag/research/evaluation_agent.py +++ /dev/null @@ -1,85 +0,0 @@ -from pydantic import BaseModel, Field -from pydantic_ai.run import AgentRunResult - -from haiku.rag.research.base import BaseResearchAgent -from haiku.rag.research.dependencies import ( - ResearchDependencies, - _format_context_for_prompt, -) -from haiku.rag.research.prompts import EVALUATION_AGENT_PROMPT - - -class EvaluationResult(BaseModel): - """Result of analysis and evaluation.""" - - key_insights: list[str] = Field( - description="Main insights extracted from the research so far" - ) - new_questions: list[str] = Field( - description="New sub-questions to add to the research (max 3)", - max_length=3, - default=[], - ) - confidence_score: float = Field( - description="Confidence level in the completeness of research (0-1)", - ge=0.0, - le=1.0, - ) - is_sufficient: bool = Field( - description="Whether the research is sufficient to answer the original question" - ) - reasoning: str = Field( - description="Explanation of why the research is or isn't complete" - ) - - -class AnalysisEvaluationAgent(BaseResearchAgent[EvaluationResult]): - """Agent that analyzes findings and evaluates research completeness.""" - - def __init__(self, provider: str, model: str) -> None: - super().__init__(provider, model, output_type=EvaluationResult) - - async def run( - self, prompt: str, deps: ResearchDependencies, **kwargs - ) -> AgentRunResult[EvaluationResult]: - console = deps.console - if console: - console.print( - "\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]" - ) - - # Format context for the evaluation agent - context_xml = _format_context_for_prompt(deps.context) - evaluation_prompt = f"""Analyze all gathered information and evaluate the completeness of research. - -{context_xml} - -Evaluate the research progress for the original question and identify any remaining gaps.""" - - result = await super().run(evaluation_prompt, deps, **kwargs) - output = result.output - - # Store insights - for insight in output.key_insights: - deps.context.add_insight(insight) - - # Add new questions to the sub-questions list - for new_q in output.new_questions: - if new_q not in deps.context.sub_questions: - deps.context.sub_questions.append(new_q) - - if console: - if output.key_insights: - console.print(" [bold]Key insights:[/bold]") - for insight in output.key_insights: - console.print(f" • {insight}") - console.print( - f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]" - ) - status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]" - console.print(f" Sufficient: {status}") - - return result - - def get_system_prompt(self) -> str: - return EVALUATION_AGENT_PROMPT diff --git a/src/haiku/rag/research/graph.py b/src/haiku/rag/research/graph.py new file mode 100644 index 00000000..9c8282e9 --- /dev/null +++ b/src/haiku/rag/research/graph.py @@ -0,0 +1,29 @@ +from pydantic_graph import Graph + +from haiku.rag.research.models import ResearchReport +from haiku.rag.research.nodes.evaluate import EvaluateNode +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", + "EvaluateNode", + "SynthesizeNode", + "ResearchState", + "ResearchDeps", + "build_research_graph", +] + + +def build_research_graph() -> Graph[ResearchState, ResearchDeps, ResearchReport]: + return Graph( + nodes=[ + PlanNode, + SearchDispatchNode, + EvaluateNode, + SynthesizeNode, + ] + ) diff --git a/src/haiku/rag/research/models.py b/src/haiku/rag/research/models.py new file mode 100644 index 00000000..3f789842 --- /dev/null +++ b/src/haiku/rag/research/models.py @@ -0,0 +1,70 @@ +from pydantic import BaseModel, Field + + +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 URIs corresponding to the snippets actually used in the" + " answer (one URI per snippet; omit if none)" + ), + default_factory=list, + ) + + +class EvaluationResult(BaseModel): + """Result of analysis and evaluation.""" + + key_insights: list[str] = Field( + description="Main insights extracted from the research so far" + ) + new_questions: list[str] = Field( + description="New sub-questions to add to the research (max 3)", + max_length=3, + default=[], + ) + confidence_score: float = Field( + description="Confidence level in the completeness of research (0-1)", + ge=0.0, + le=1.0, + ) + is_sufficient: bool = Field( + description="Whether the research is sufficient to answer the original question" + ) + reasoning: str = Field( + description="Explanation of why the research is or isn't complete" + ) + + +class ResearchReport(BaseModel): + """Final research report structure.""" + + title: str = Field(description="Concise title for the research") + executive_summary: str = Field(description="Brief overview of key findings") + main_findings: list[str] = Field( + description="Primary research findings with supporting evidence" + ) + conclusions: list[str] = Field(description="Evidence-based conclusions") + limitations: list[str] = Field( + description="Limitations of the current research", default=[] + ) + recommendations: list[str] = Field( + description="Actionable recommendations based on findings", default=[] + ) + sources_summary: str = Field( + description="Summary of sources used and their reliability" + ) diff --git a/src/haiku/rag/research/nodes/evaluate.py b/src/haiku/rag/research/nodes/evaluate.py new file mode 100644 index 00000000..79cb80d5 --- /dev/null +++ b/src/haiku/rag/research/nodes/evaluate.py @@ -0,0 +1,81 @@ +from dataclasses import dataclass + +from pydantic_ai import Agent +from pydantic_graph import BaseNode, GraphRunContext + +from haiku.rag.research.common import format_context_for_prompt, get_model, log +from haiku.rag.research.dependencies import ( + ResearchDependencies, +) +from haiku.rag.research.models import EvaluationResult, ResearchReport +from haiku.rag.research.nodes.synthesize import SynthesizeNode +from haiku.rag.research.prompts import EVALUATION_AGENT_PROMPT +from haiku.rag.research.state import ResearchDeps, ResearchState + + +@dataclass +class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): + provider: str + model: str + + async def run( + self, ctx: GraphRunContext[ResearchState, ResearchDeps] + ) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]: + state = ctx.state + deps = ctx.deps + + log( + deps.console, + "\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]", + ) + + agent = Agent( + model=get_model(self.provider, self.model), + output_type=EvaluationResult, + instructions=EVALUATION_AGENT_PROMPT, + retries=3, + deps_type=ResearchDependencies, + ) + + context_xml = format_context_for_prompt(state.context) + prompt = ( + "Analyze gathered information and evaluate completeness for the original question.\n\n" + f"{context_xml}" + ) + agent_deps = ResearchDependencies( + client=deps.client, context=state.context, console=deps.console + ) + eval_result = await agent.run(prompt, deps=agent_deps) + output = eval_result.output + + for insight in output.key_insights: + state.context.add_insight(insight) + for new_q in output.new_questions: + if new_q not in state.sub_questions: + state.sub_questions.append(new_q) + + state.last_eval = output + state.iterations += 1 + + if deps.console: + if output.key_insights: + deps.console.print(" [bold]Key insights:[/bold]") + for ins in output.key_insights: + deps.console.print(f" • {ins}") + deps.console.print( + f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]" + ) + status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]" + deps.console.print(f" Sufficient: {status}") + + from haiku.rag.research.nodes.search import SearchDispatchNode + + if ( + output.is_sufficient + and output.confidence_score >= state.confidence_threshold + ) or state.iterations >= state.max_iterations: + if deps.console: + deps.console.print("\n[bold green]✅ Stopping research.[/bold green]") + return SynthesizeNode(self.provider, self.model) + + return SearchDispatchNode(self.provider, self.model) diff --git a/src/haiku/rag/research/nodes/plan.py b/src/haiku/rag/research/nodes/plan.py new file mode 100644 index 00000000..c5005dd0 --- /dev/null +++ b/src/haiku/rag/research/nodes/plan.py @@ -0,0 +1,64 @@ +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.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 ORCHESTRATOR_PROMPT +from haiku.rag.research.state import ResearchDeps, ResearchState + + +@dataclass +class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): + provider: str + model: str + + async def run( + self, ctx: GraphRunContext[ResearchState, ResearchDeps] + ) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]: + state = ctx.state + deps = ctx.deps + + log(deps.console, "\n[bold cyan]📋 Creating research plan...[/bold cyan]") + + plan_agent = Agent( + model=get_model(self.provider, self.model), + output_type=ResearchPlan, + instructions=( + ORCHESTRATOR_PROMPT + + "\n\nUse the gather_context tool once on the main question before planning." + ), + retries=3, + deps_type=ResearchDependencies, + ) + + @plan_agent.tool + async def gather_context( + ctx2: RunContext[ResearchDependencies], 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 research approach for the main question.\n\n" + f"Main question: {state.question}" + ) + + agent_deps = ResearchDependencies( + client=deps.client, context=state.context, console=deps.console + ) + plan_result = await plan_agent.run(prompt, deps=agent_deps) + state.sub_questions = list(plan_result.output.sub_questions) + + if deps.console: + deps.console.print("\n[bold green]✅ Research Plan Created:[/bold green]") + deps.console.print(f" [bold]Main Question:[/bold] {state.question}") + deps.console.print(" [bold]Sub-questions:[/bold]") + for i, sq in enumerate(state.sub_questions, 1): + deps.console.print(f" {i}. {sq}") + + return SearchDispatchNode(self.provider, self.model) diff --git a/src/haiku/rag/research/nodes/search.py b/src/haiku/rag/research/nodes/search.py new file mode 100644 index 00000000..4640b3c1 --- /dev/null +++ b/src/haiku/rag/research/nodes/search.py @@ -0,0 +1,92 @@ +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, GraphRunContext + +from haiku.rag.research.common import get_model, log +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.state import ResearchDeps, ResearchState + + +@dataclass +class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): + provider: str + model: str + + async def run( + self, ctx: GraphRunContext[ResearchState, ResearchDeps] + ) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]: + state = ctx.state + deps = ctx.deps + if not state.sub_questions: + from haiku.rag.research.nodes.evaluate import EvaluateNode + + return EvaluateNode(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.sub_questions and len(batch) < take: + batch.append(state.sub_questions.pop(0)) + + async def answer_one(sub_q: str) -> SearchAnswer | None: + if deps.console: + deps.console.print( + 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=ResearchDependencies, + ) + + @agent.tool + async def search_and_answer( + ctx2: RunContext[ResearchDependencies], 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_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 = ResearchDependencies( + 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.console, 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) + if deps.console: + preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "") + deps.console.log(f" [green]✓[/green] {preview}") + + return SearchDispatchNode(self.provider, self.model) diff --git a/src/haiku/rag/research/nodes/synthesize.py b/src/haiku/rag/research/nodes/synthesize.py new file mode 100644 index 00000000..e2ec6be4 --- /dev/null +++ b/src/haiku/rag/research/nodes/synthesize.py @@ -0,0 +1,51 @@ +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.research.models import ResearchReport +from haiku.rag.research.prompts import SYNTHESIS_AGENT_PROMPT +from haiku.rag.research.state import ResearchDeps, ResearchState + + +@dataclass +class SynthesizeNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): + provider: str + model: str + + async def run( + self, ctx: GraphRunContext[ResearchState, ResearchDeps] + ) -> End[ResearchReport]: + state = ctx.state + deps = ctx.deps + + log( + deps.console, + "\n[bold cyan]📝 Generating final research report...[/bold cyan]", + ) + + agent = Agent( + model=get_model(self.provider, self.model), + output_type=ResearchReport, + instructions=SYNTHESIS_AGENT_PROMPT, + retries=3, + deps_type=ResearchDependencies, + ) + + context_xml = format_context_for_prompt(state.context) + prompt = ( + "Generate a comprehensive research report based on all gathered information.\n\n" + f"{context_xml}\n\n" + "Create a detailed report that synthesizes all findings into a coherent response." + ) + agent_deps = ResearchDependencies( + client=deps.client, context=state.context, console=deps.console + ) + result = await agent.run(prompt, deps=agent_deps) + + log(deps.console, "[bold green]✅ Research complete![/bold green]") + return End(result.output) diff --git a/src/haiku/rag/research/orchestrator.py b/src/haiku/rag/research/orchestrator.py deleted file mode 100644 index af8ed766..00000000 --- a/src/haiku/rag/research/orchestrator.py +++ /dev/null @@ -1,170 +0,0 @@ -from typing import Any - -from pydantic import BaseModel, Field -from pydantic_ai.run import AgentRunResult -from rich.console import Console - -from haiku.rag.config import Config -from haiku.rag.research.base import BaseResearchAgent -from haiku.rag.research.dependencies import ( - ResearchContext, - ResearchDependencies, -) -from haiku.rag.research.evaluation_agent import ( - AnalysisEvaluationAgent, - EvaluationResult, -) -from haiku.rag.research.presearch_agent import PresearchSurveyAgent -from haiku.rag.research.prompts import ORCHESTRATOR_PROMPT -from haiku.rag.research.search_agent import SearchSpecialistAgent -from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent - - -class ResearchPlan(BaseModel): - """Research execution plan.""" - - main_question: str = Field(description="The main research question") - sub_questions: list[str] = Field( - description="Decomposed sub-questions to investigate (max 3)", max_length=3 - ) - - -class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): - """Orchestrator agent that coordinates the research workflow.""" - - def __init__( - self, - provider: str | None = Config.RESEARCH_PROVIDER, - model: str | None = None, - ): - # Use provided values or fall back to config defaults - provider = provider or Config.RESEARCH_PROVIDER or Config.QA_PROVIDER - model = model or Config.RESEARCH_MODEL or Config.QA_MODEL - - super().__init__(provider, model, output_type=ResearchPlan) - - self.search_agent: SearchSpecialistAgent = SearchSpecialistAgent( - provider, model - ) - self.presearch_agent: PresearchSurveyAgent = PresearchSurveyAgent( - provider, model - ) - self.evaluation_agent: AnalysisEvaluationAgent = AnalysisEvaluationAgent( - provider, model - ) - self.synthesis_agent: SynthesisAgent = SynthesisAgent(provider, model) - - def get_system_prompt(self) -> str: - return ORCHESTRATOR_PROMPT - - def _should_stop_research( - self, - evaluation_result: AgentRunResult[EvaluationResult], - confidence_threshold: float, - ) -> bool: - """Determine if research should stop based on evaluation.""" - - result = evaluation_result.output - return result.is_sufficient and result.confidence_score >= confidence_threshold - - async def conduct_research( - self, - question: str, - client: Any, - max_iterations: int = 3, - confidence_threshold: float = 0.8, - verbose: bool = False, - ) -> ResearchReport: - """Conduct comprehensive research on a question. - - Args: - question: The research question to investigate - client: HaikuRAG client for document operations - max_iterations: Maximum number of search-analyze-clarify cycles - confidence_threshold: Minimum confidence level to stop research (0-1) - verbose: If True, print progress and intermediate results - - Returns: - ResearchReport with comprehensive findings - """ - - # Initialize context - context = ResearchContext(original_question=question) - deps = ResearchDependencies(client=client, context=context) - if verbose: - deps.console = Console() - - console = deps.console - # Create initial research plan - if console: - console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]") - - # Run a simple presearch survey to summarize KB context - presearch_result = await self.presearch_agent.run(question, deps=deps) - plan_prompt = ( - "Create a research plan for the main question below.\n\n" - f"Main question: {question}\n\n" - "Use this brief presearch summary to inform the plan. Focus the 3 sub-questions " - "on the most important aspects not already obvious from the current KB context.\n\n" - f"{presearch_result.output}" - ) - - plan_result: AgentRunResult[ResearchPlan] = await self.run( - plan_prompt, deps=deps - ) - context.sub_questions = plan_result.output.sub_questions - - if console: - console.print("\n[bold green]✅ Research Plan Created:[/bold green]") - console.print( - f" [bold]Main Question:[/bold] {plan_result.output.main_question}" - ) - console.print(" [bold]Sub-questions:[/bold]") - for i, sq in enumerate(plan_result.output.sub_questions, 1): - console.print(f" {i}. {sq}") - - # Execute research iterations - for iteration in range(max_iterations): - if console: - console.rule( - f"[bold yellow]🔄 Iteration {iteration + 1}/{max_iterations}[/bold yellow]" - ) - - # Check if we have questions to search - if not context.sub_questions: - if console: - console.print( - "[yellow]No more questions to explore. Concluding research.[/yellow]" - ) - break - - # Use current sub-questions for this iteration - questions_to_search = context.sub_questions[:] - - # Search phase - answer all questions in this iteration - if console: - console.print( - f"\n[bold cyan]🔍 Searching & Answering {len(questions_to_search)} questions:[/bold cyan]" - ) - - for search_question in questions_to_search: - await self.search_agent.run(search_question, deps=deps) - - # Analysis and Evaluation phase - - evaluation_result = await self.evaluation_agent.run("", deps=deps) - - # Check if research is sufficient - if self._should_stop_research(evaluation_result, confidence_threshold): - if console: - console.print( - f"\n[bold green]✅ Stopping research:[/bold green] {evaluation_result.output.reasoning}" - ) - break - - # Generate final report - report_result: AgentRunResult[ResearchReport] = await self.synthesis_agent.run( - "", deps=deps - ) - - return report_result.output diff --git a/src/haiku/rag/research/presearch_agent.py b/src/haiku/rag/research/presearch_agent.py deleted file mode 100644 index 5482dfc8..00000000 --- a/src/haiku/rag/research/presearch_agent.py +++ /dev/null @@ -1,39 +0,0 @@ -from pydantic_ai import RunContext -from pydantic_ai.run import AgentRunResult - -from haiku.rag.research.base import BaseResearchAgent -from haiku.rag.research.dependencies import ResearchDependencies -from haiku.rag.research.prompts import PRESEARCH_AGENT_PROMPT - - -class PresearchSurveyAgent(BaseResearchAgent[str]): - """Presearch agent that gathers verbatim context and summarizes it.""" - - def __init__(self, provider: str, model: str) -> None: - super().__init__(provider, model, str) - - async def run( - self, prompt: str, deps: ResearchDependencies, **kwargs - ) -> AgentRunResult[str]: - console = deps.console - if console: - console.print( - "\n[bold cyan]🔎 Presearch: summarizing KB context...[/bold cyan]" - ) - - return await super().run(prompt, deps, **kwargs) - - def get_system_prompt(self) -> str: - return PRESEARCH_AGENT_PROMPT - - def register_tools(self) -> None: - @self.agent.tool - async def gather_context( - ctx: RunContext[ResearchDependencies], - query: str, - limit: int = 6, - ) -> str: - """Return verbatim concatenation of relevant chunk texts.""" - results = await ctx.deps.client.search(query, limit=limit) - expanded = await ctx.deps.client.expand_context(results) - return "\n\n".join(chunk.content for chunk, _ in expanded) diff --git a/src/haiku/rag/research/search_agent.py b/src/haiku/rag/research/search_agent.py deleted file mode 100644 index 321b6fee..00000000 --- a/src/haiku/rag/research/search_agent.py +++ /dev/null @@ -1,69 +0,0 @@ -from pydantic_ai import RunContext -from pydantic_ai.format_prompt import format_as_xml -from pydantic_ai.run import AgentRunResult - -from haiku.rag.research.base import BaseResearchAgent, SearchAnswer -from haiku.rag.research.dependencies import ResearchDependencies -from haiku.rag.research.prompts import SEARCH_AGENT_PROMPT - - -class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]): - """Agent specialized in answering questions using RAG search.""" - - def __init__(self, provider: str, model: str) -> None: - super().__init__(provider, model, output_type=SearchAnswer) - - async def run( - self, prompt: str, deps: ResearchDependencies, **kwargs - ) -> AgentRunResult[SearchAnswer]: - """Execute the agent and persist the QA pair in shared context. - - Pydantic AI enforces `SearchAnswer` as the output model; we just store - the QA response with the last search results as sources. - """ - console = deps.console - if console: - console.print(f"\t{prompt}") - - result = await super().run(prompt, deps, **kwargs) - deps.context.add_qa_response(result.output) - deps.context.sub_questions.remove(prompt) - if console: - answer = result.output.answer - answer_preview = answer[:150] + "…" if len(answer) > 150 else answer - console.log(f"\n [green]✓[/green] {answer_preview}") - - return result - - def get_system_prompt(self) -> str: - return SEARCH_AGENT_PROMPT - - def register_tools(self) -> None: - """Register search-specific tools.""" - - @self.agent.tool - async def search_and_answer( - ctx: RunContext[ResearchDependencies], - query: str, - limit: int = 5, - ) -> str: - """Search the KB and return a concise context pack.""" - search_results = await ctx.deps.client.search(query, limit=limit) - expanded = await ctx.deps.client.expand_context(search_results) - - snippet_entries = [ - { - "text": chunk.content, - "score": score, - "document_uri": (chunk.document_uri or ""), - } - for chunk, score in expanded - ] - - # Return an XML-formatted payload with the question and snippets. - if snippet_entries: - return format_as_xml(snippet_entries, root_tag="snippets") - else: - return ( - f"No relevant information found in the knowledge base for: {query}" - ) diff --git a/src/haiku/rag/research/state.py b/src/haiku/rag/research/state.py new file mode 100644 index 00000000..6085e871 --- /dev/null +++ b/src/haiku/rag/research/state.py @@ -0,0 +1,25 @@ +from dataclasses import dataclass, field + +from rich.console import Console + +from haiku.rag.client import HaikuRAG +from haiku.rag.research.dependencies import ResearchContext +from haiku.rag.research.models import EvaluationResult + + +@dataclass +class ResearchDeps: + client: HaikuRAG + console: Console | None = None + + +@dataclass +class ResearchState: + question: str + context: ResearchContext + sub_questions: list[str] = field(default_factory=list) + iterations: int = 0 + max_iterations: int = 3 + max_concurrency: int = 3 + confidence_threshold: float = 0.8 + last_eval: EvaluationResult | None = None diff --git a/src/haiku/rag/research/synthesis_agent.py b/src/haiku/rag/research/synthesis_agent.py deleted file mode 100644 index e3499e95..00000000 --- a/src/haiku/rag/research/synthesis_agent.py +++ /dev/null @@ -1,60 +0,0 @@ -from pydantic import BaseModel, Field -from pydantic_ai.run import AgentRunResult - -from haiku.rag.research.base import BaseResearchAgent -from haiku.rag.research.dependencies import ( - ResearchDependencies, - _format_context_for_prompt, -) -from haiku.rag.research.prompts import SYNTHESIS_AGENT_PROMPT - - -class ResearchReport(BaseModel): - """Final research report structure.""" - - title: str = Field(description="Concise title for the research") - executive_summary: str = Field(description="Brief overview of key findings") - main_findings: list[str] = Field( - description="Primary research findings with supporting evidence" - ) - conclusions: list[str] = Field(description="Evidence-based conclusions") - limitations: list[str] = Field( - description="Limitations of the current research", default=[] - ) - recommendations: list[str] = Field( - description="Actionable recommendations based on findings", default=[] - ) - sources_summary: str = Field( - description="Summary of sources used and their reliability" - ) - - -class SynthesisAgent(BaseResearchAgent[ResearchReport]): - """Agent specialized in synthesizing research into comprehensive reports.""" - - def __init__(self, provider: str, model: str) -> None: - super().__init__(provider, model, output_type=ResearchReport) - - async def run( - self, prompt: str, deps: ResearchDependencies, **kwargs - ) -> AgentRunResult[ResearchReport]: - console = deps.console - if console: - console.print( - "\n[bold cyan]📝 Generating final research report...[/bold cyan]" - ) - - context_xml = _format_context_for_prompt(deps.context) - synthesis_prompt = f"""Generate a comprehensive research report based on all gathered information. - -{context_xml} - -Create a detailed report that synthesizes all findings into a coherent response.""" - result = await super().run(synthesis_prompt, deps, **kwargs) - if console: - console.print("[bold green]✅ Research complete![/bold green]") - - return result - - def get_system_prompt(self) -> str: - return SYNTHESIS_AGENT_PROMPT diff --git a/tests/research/test_evaluation_agent.py b/tests/research/test_evaluation_agent.py deleted file mode 100644 index 0c3a50d6..00000000 --- a/tests/research/test_evaluation_agent.py +++ /dev/null @@ -1,17 +0,0 @@ -from haiku.rag.config import Config -from haiku.rag.research.evaluation_agent import ( - AnalysisEvaluationAgent, - EvaluationResult, -) - - -class TestAnalysisEvaluationAgent: - """Lean tests for AnalysisEvaluationAgent without LLM mocking.""" - - def test_agent_initialization(self): - agent = AnalysisEvaluationAgent( - provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL - ) - assert agent.provider == Config.RESEARCH_PROVIDER - assert agent.model == Config.RESEARCH_MODEL - assert agent.output_type == EvaluationResult diff --git a/tests/research/test_orchestrator.py b/tests/research/test_orchestrator.py deleted file mode 100644 index 304eb788..00000000 --- a/tests/research/test_orchestrator.py +++ /dev/null @@ -1,189 +0,0 @@ -from unittest.mock import AsyncMock, create_autospec - -import pytest -from pydantic_ai.models.test import TestModel - -from haiku.rag.client import HaikuRAG -from haiku.rag.config import Config -from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies -from haiku.rag.research.evaluation_agent import EvaluationResult -from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan -from haiku.rag.research.synthesis_agent import ResearchReport -from haiku.rag.store.models.chunk import Chunk - - -@pytest.fixture -def test_model(): - """Create a test model for orchestrator testing.""" - return TestModel() - - -@pytest.fixture -def mock_client(): - """Create a mock HaikuRAG client.""" - client = create_autospec(HaikuRAG, instance=True) - client.search = AsyncMock() - client.expand_context = AsyncMock() - return client - - -@pytest.fixture -def research_context(): - """Create a research context.""" - return ResearchContext(original_question="What is climate change?") - - -@pytest.fixture -def research_deps(mock_client, research_context): - """Create research dependencies.""" - return ResearchDependencies(client=mock_client, context=research_context) - - -def create_mock_chunk(chunk_id: str, content: str, score: float = 0.8): - """Helper to create mock chunk objects.""" - return Chunk( - id=chunk_id, - document_id=f"doc_{chunk_id}", - content=content, - document_uri=f"doc_{chunk_id}.md", - metadata={}, - ), score - - -class TestResearchOrchestrator: - """Test suite for ResearchOrchestrator.""" - - def test_orchestrator_uses_config_defaults(self): - """Test that orchestrator uses config defaults when no args provided.""" - orchestrator = ResearchOrchestrator() - - # Should use RESEARCH_PROVIDER/MODEL if set, else QA_PROVIDER/MODEL - assert orchestrator.provider is not None - assert orchestrator.model is not None - - # All agents should use the same provider/model - assert orchestrator.search_agent.provider == orchestrator.provider - assert orchestrator.search_agent.model == orchestrator.model - assert orchestrator.evaluation_agent.provider == orchestrator.provider - assert orchestrator.evaluation_agent.model == orchestrator.model - assert orchestrator.synthesis_agent.provider == orchestrator.provider - assert orchestrator.synthesis_agent.model == orchestrator.model - - def test_orchestrator_initialization(self): - """Test that orchestrator initializes all agents correctly.""" - orchestrator = ResearchOrchestrator( - provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL - ) - - # Check all agents are initialized - assert orchestrator.search_agent is not None - assert orchestrator.evaluation_agent is not None - assert orchestrator.synthesis_agent is not None - - # Check they all use the same provider and model - assert orchestrator.search_agent.provider == Config.RESEARCH_PROVIDER - assert orchestrator.search_agent.model == Config.RESEARCH_MODEL - assert orchestrator.evaluation_agent.provider == Config.RESEARCH_PROVIDER - assert orchestrator.evaluation_agent.model == Config.RESEARCH_MODEL - assert orchestrator.synthesis_agent.provider == Config.RESEARCH_PROVIDER - assert orchestrator.synthesis_agent.model == Config.RESEARCH_MODEL - - def test_orchestrator_has_correct_output_type(self): - """Test that orchestrator's output type is ResearchPlan.""" - orchestrator = ResearchOrchestrator( - provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL - ) - assert orchestrator.output_type == ResearchPlan - - def test_orchestrator_has_no_tools(self): - """Test that orchestrator no longer registers tools (direct agent calls now).""" - orchestrator = ResearchOrchestrator( - provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL - ) - - # Get the tools from the agent - tools = orchestrator.agent._function_toolset.tools - tool_names = list(tools.keys()) - - # Should have no tools since we call agents directly now - assert len(tool_names) == 0 - - def test_should_stop_research_logic(self): - """Test the stopping logic based on EvaluationResult.""" - orchestrator = ResearchOrchestrator( - provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL - ) - - # Create mock evaluation results - from unittest.mock import MagicMock - - # Sufficient research result - sufficient_result = MagicMock() - sufficient_result.output = EvaluationResult( - key_insights=["Climate is changing", "Human activity is the cause"], - new_questions=[], - confidence_score=0.9, - is_sufficient=True, - reasoning="All aspects covered comprehensively", - ) - - # Insufficient research result - insufficient_result = MagicMock() - insufficient_result.output = EvaluationResult( - key_insights=["Some data found"], - new_questions=[ - "What about economic impacts?", - "Regional variations?", - ], - confidence_score=0.4, - is_sufficient=False, - reasoning="Major gaps remain in understanding", - ) - - # Test with sufficient research (threshold 0.8) - assert orchestrator._should_stop_research(sufficient_result, 0.8) - - # Test with insufficient research - assert not orchestrator._should_stop_research(insufficient_result, 0.8) - - # Test with high confidence but below threshold - sufficient_result.output.confidence_score = 0.75 - assert not orchestrator._should_stop_research(sufficient_result, 0.8) - - # Test with is_sufficient=False even with high confidence - insufficient_result.output.confidence_score = 0.95 - assert not orchestrator._should_stop_research(insufficient_result, 0.8) - - @pytest.mark.asyncio - async def test_conduct_research_workflow(self, test_model, mock_client): - """Test the basic research workflow using TestModel.""" - orchestrator = ResearchOrchestrator( - provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL - ) - - # Setup mock client returns - mock_chunks = [ - create_mock_chunk("1", "Climate change information"), - ] - mock_client.search.return_value = mock_chunks - mock_client.expand_context.return_value = mock_chunks - - # Use TestModel for all agents - with orchestrator.agent.override(model=test_model): - with orchestrator.search_agent.agent.override(model=test_model): - with orchestrator.evaluation_agent.agent.override(model=test_model): - with orchestrator.synthesis_agent.agent.override(model=test_model): - # Run the research - report = await orchestrator.conduct_research( - "What is climate change?", mock_client, max_iterations=1 - ) - - # Verify we got a valid report structure - assert isinstance(report, ResearchReport) - assert report.title - assert report.executive_summary - assert isinstance(report.main_findings, list) - assert isinstance(report.conclusions, list) - assert isinstance(report.limitations, list) - assert isinstance(report.recommendations, list) - assert report.sources_summary diff --git a/tests/research/test_search_agent.py b/tests/research/test_search_agent.py deleted file mode 100644 index 71196f27..00000000 --- a/tests/research/test_search_agent.py +++ /dev/null @@ -1,14 +0,0 @@ -from haiku.rag.config import Config -from haiku.rag.research import SearchAnswer, SearchSpecialistAgent - - -class TestSearchSpecialistAgent: - """Lean tests for SearchSpecialistAgent without LLM mocking.""" - - def test_agent_initialization(self): - agent = SearchSpecialistAgent( - provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL - ) - assert agent.provider == Config.RESEARCH_PROVIDER - assert agent.model == Config.RESEARCH_MODEL - assert agent.output_type is SearchAnswer diff --git a/tests/research/test_synthesis_agent.py b/tests/research/test_synthesis_agent.py deleted file mode 100644 index 65146782..00000000 --- a/tests/research/test_synthesis_agent.py +++ /dev/null @@ -1,14 +0,0 @@ -from haiku.rag.config import Config -from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent - - -class TestSynthesisAgent: - """Lean tests for SynthesisAgent without LLM mocking.""" - - def test_agent_initialization(self): - agent = SynthesisAgent( - provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL - ) - assert agent.provider == Config.RESEARCH_PROVIDER - assert agent.model == Config.RESEARCH_MODEL - assert agent.output_type == ResearchReport diff --git a/tests/test_research_graph.py b/tests/test_research_graph.py new file mode 100644 index 00000000..8b50a6f3 --- /dev/null +++ b/tests/test_research_graph.py @@ -0,0 +1,26 @@ +import asyncio + +from haiku.rag.research.dependencies import ResearchContext +from haiku.rag.research.graph import ResearchState, build_research_graph + + +def test_build_graph_and_state(): + graph = build_research_graph() + assert graph is not None + + state = ResearchState( + question="What are the key features of haiku.rag?", + context=ResearchContext( + original_question="What are the key features of haiku.rag?" + ), + max_iterations=1, + confidence_threshold=0.8, + ) + assert state.iterations == 0 + assert state.sub_questions == [] + + +def test_async_loop_available(): + # Ensure an event loop can be created in test env + loop = asyncio.new_event_loop() + loop.close() diff --git a/tests/test_research_graph_integration.py b/tests/test_research_graph_integration.py new file mode 100644 index 00000000..ea719155 --- /dev/null +++ b/tests/test_research_graph_integration.py @@ -0,0 +1,89 @@ +from typing import Any, cast + +import pytest + +from haiku.rag.research.dependencies import ResearchContext +from haiku.rag.research.graph import ( + EvaluateNode, + PlanNode, + ResearchDeps, + ResearchState, + SearchDispatchNode, + SynthesizeNode, + build_research_graph, +) +from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer + + +@pytest.mark.asyncio +async def test_graph_end_to_end_with_patched_nodes(monkeypatch): + graph = build_research_graph() + + state = ResearchState( + question="What is haiku.rag?", + context=ResearchContext(original_question="What is haiku.rag?"), + max_iterations=1, + confidence_threshold=0.5, + max_concurrency=2, + ) + deps = ResearchDeps( + client=cast(Any, None), console=None + ) # client unused in patched nodes + + async def fake_plan_run(self, ctx) -> Any: + ctx.state.sub_questions = [ + "Describe haiku.rag in one sentence", + "List core components of haiku.rag", + ] + return SearchDispatchNode(self.provider, self.model) + + async def fake_search_dispatch_run(self, ctx) -> Any: + # Answer all pending questions deterministically, then move to evaluation + while ctx.state.sub_questions: + q = ctx.state.sub_questions.pop(0) + # pydantic BaseModel kwargs not fully typed for pyright + ctx.state.context.add_qa_response( + SearchAnswer(query=q, answer="A", context=["x"], sources=["s"]) # pyright: ignore[reportCallIssue] + ) + return EvaluateNode(self.provider, self.model) + + async def fake_evaluate_run(self, ctx) -> Any: + ctx.state.last_eval = EvaluationResult( + key_insights=["ok"], + new_questions=[], + confidence_score=1.0, + is_sufficient=True, + reasoning="done", + ) + ctx.state.iterations += 1 + return SynthesizeNode(self.provider, self.model) + + async def fake_synthesize_run(self, ctx) -> Any: + report = ResearchReport( + title="Haiku RAG", + executive_summary="...", + main_findings=["f1"], + conclusions=["c1"], + limitations=[], + recommendations=[], + sources_summary="s", + ) + from pydantic_graph import End + + return End(report) + + monkeypatch.setattr(PlanNode, "run", fake_plan_run, raising=False) + monkeypatch.setattr( + SearchDispatchNode, "run", fake_search_dispatch_run, raising=False + ) + monkeypatch.setattr(EvaluateNode, "run", fake_evaluate_run, raising=False) + monkeypatch.setattr(SynthesizeNode, "run", fake_synthesize_run, raising=False) + + start = PlanNode(provider="test", model="test") + + result = await graph.run(start, state=state, deps=deps) + report = result.output + + assert isinstance(report, ResearchReport) + assert report.title == "Haiku RAG" + assert len(state.context.qa_responses) == 2 From db47eb45cb7abf38ddc2d4fdea70b7fbd468c3b7 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 19 Sep 2025 15:16:46 +0300 Subject: [PATCH 2/7] Update prompts --- src/haiku/rag/research/nodes/plan.py | 4 +- src/haiku/rag/research/prompts.py | 184 ++++++++++++--------------- 2 files changed, 86 insertions(+), 102 deletions(-) diff --git a/src/haiku/rag/research/nodes/plan.py b/src/haiku/rag/research/nodes/plan.py index c5005dd0..a412639f 100644 --- a/src/haiku/rag/research/nodes/plan.py +++ b/src/haiku/rag/research/nodes/plan.py @@ -7,7 +7,7 @@ from haiku.rag.research.common import get_model, log 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 ORCHESTRATOR_PROMPT +from haiku.rag.research.prompts import PLAN_PROMPT from haiku.rag.research.state import ResearchDeps, ResearchState @@ -28,7 +28,7 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): model=get_model(self.provider, self.model), output_type=ResearchPlan, instructions=( - ORCHESTRATOR_PROMPT + PLAN_PROMPT + "\n\nUse the gather_context tool once on the main question before planning." ), retries=3, diff --git a/src/haiku/rag/research/prompts.py b/src/haiku/rag/research/prompts.py index af5038f7..6d9aad48 100644 --- a/src/haiku/rag/research/prompts.py +++ b/src/haiku/rag/research/prompts.py @@ -1,129 +1,113 @@ -ORCHESTRATOR_PROMPT = """You are a research orchestrator responsible for coordinating a comprehensive research workflow. +PLAN_PROMPT = """You are the research orchestrator for a focused, iterative +workflow. -Your role is to: -1. Understand and decompose the research question -2. Plan a systematic research approach -3. Coordinate specialized agents to gather and analyze information -4. Ensure comprehensive coverage of the topic -5. Iterate based on findings and gaps +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 -Create a research plan that: -- Breaks down the question into at most 3 focused sub-questions -- Each sub-question should target a specific aspect of the research -- Prioritize the most important aspects to investigate -- Ensure comprehensive coverage within the 3-question limit -- IMPORTANT: Make each sub-question a standalone, self-contained query that can - be executed without additional context. Include necessary entities, scope, - timeframe, and qualifiers. Avoid pronouns like "it/they/this"; write queries - that make sense in isolation.""" +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. +SEARCH_AGENT_PROMPT = """You are a search and question‑answering specialist. -Your role is to: -1. Search the knowledge base for relevant information -2. Analyze the retrieved documents -3. Provide an accurate answer strictly grounded in the retrieved context - -Output format: -- You must return a SearchAnswer model with fields: - - query: the question being answered (echo the user query) - - answer: your final answer based only on the provided context - - context: list[str] of only the minimal set of verbatim snippet texts you - used to justify the answer (do not include unrelated text; do not invent) - - sources: list[str] of document_uri values corresponding to the snippets you - actually used in the answer (one URI per context snippet, order aligned) +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 the search_and_answer tool before drafting any answer. -- The tool returns XML containing only a list of snippets, where each snippet - has the verbatim `text`, a `score` indicating relevance, and the - `document_uri` it came from. +- Always call search_and_answer before drafting any answer. +- The tool returns snippets with verbatim `text`, a relevance `score`, and the + originating `document_uri`. - You may call the tool multiple times to refine or broaden context, but do not - exceed 3 total tool calls per question. Prefer precision over volume. + 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. -- Set SearchAnswer.sources to the matching document_uris for the snippets you - used (one URI per snippet, aligned by order). Context must be text-only. -- If no relevant information is found, say so and return an empty context list. + snippet texts (verbatim) in SearchAnswer.context (typically 1‑4). +- Set SearchAnswer.sources to the corresponding document_uris for the snippets + you used (one URI 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. -Important: -- Do not include any content in the answer that is not supported by the context. -- Keep context snippets short (just the necessary lines), verbatim, and focused.""" +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.""" -EVALUATION_AGENT_PROMPT = """You are an analysis and evaluation specialist for research workflows. +EVALUATION_AGENT_PROMPT = """You are an analysis and evaluation specialist for +the research workflow. -You have access to: -- The original research question -- Question-answer pairs from search operations -- Raw search results and source documents +Inputs available: +- Original research question +- Question–answer pairs produced by search +- Raw search results and source metadata - Previously identified insights -Your dual role is to: - ANALYSIS: -1. Extract key insights from all gathered information -2. Identify patterns and connections across sources -3. Synthesize findings into coherent understanding -4. Focus on the most important discoveries +1. Extract the most important, non‑obvious insights from the collected evidence. +2. Identify patterns, agreements, and disagreements across sources. +3. Note material uncertainties and assumptions. EVALUATION: -1. Assess if we have sufficient information to answer the original question -2. Calculate a confidence score (0-1) based on: - - Coverage of the main question's aspects - - Quality and consistency of sources - - Depth of information gathered -3. Identify specific gaps that still need investigation -4. Generate up to 3 new sub-questions that haven't been answered yet +1. Decide if we have sufficient information to answer the original question. +2. Provide a confidence_score in [0,1] considering: + - Coverage of the main question’s aspects + - Quality, consistency, and diversity of sources + - Depth and specificity of evidence +3. List concrete gaps that still need investigation. +4. Propose up to 3 new sub_questions that would close the highest‑value gaps. -Be critical and thorough in your evaluation. Only mark research as sufficient when: -- All major aspects of the question are addressed -- Sources provide consistent, reliable information -- The depth of coverage meets the question's requirements -- No critical gaps remain +Strictness: +- Only mark research as sufficient when all major aspects are addressed with + consistent, reliable evidence and no critical gaps remain. -Generate new sub-questions that: -- Target specific unexplored aspects not covered by existing questions -- Seek clarification on ambiguities -- Explore important edge cases or exceptions -- Are focused and actionable (max 3) -- Do NOT repeat or rephrase questions that have already been answered (see qa_responses) -- Should be genuinely new areas to explore -- Must be standalone, self-contained queries: include entities, scope, and any - needed qualifiers (e.g., timeframe, region), and avoid ambiguous pronouns so - they can be executed independently.""" +New sub_questions must: +- Be genuinely new (not answered or duplicative; check qa_responses). +- Be standalone and specific (entities, scope, timeframe/region if relevant). +- Be actionable and scoped to the knowledge base (narrow if necessary). +- Be ordered by expected impact (most valuable first).""" -SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist agent focused on creating comprehensive research reports. +SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist producing the final +research report. -Your role is to: -1. Synthesize all gathered information into a coherent narrative -2. Present findings in a clear, structured format -3. Draw evidence-based conclusions -4. Acknowledge limitations and uncertainties -5. Provide actionable recommendations -6. Maintain academic rigor and objectivity +Goals: +1. Synthesize all gathered information into a coherent narrative. +2. Present findings clearly and concisely. +3. Draw evidence‑based conclusions and recommendations. +4. State limitations and uncertainties transparently. -Your report should be: -- Comprehensive yet concise -- Well-structured and easy to follow -- Based solely on evidence from the research -- Transparent about limitations -- Professional and objective in tone +Report guidelines (map to output fields): +- title: concise (5–12 words), informative. +- executive_summary: 3–5 sentences summarizing the overall answer. +- main_findings: 4–8 one‑sentence bullets; each reflects evidence from the + research (do not include inline citations or snippet text). +- conclusions: 2–4 bullets that follow logically from findings. +- recommendations: 2–5 actionable bullets tied to findings. +- limitations: 1–3 bullets describing key constraints or uncertainties. +- sources_summary: 2–4 sentences summarizing sources used and their reliability. -Focus on creating a report that provides clear value to the reader by: -- Answering the original research question thoroughly -- Highlighting the most important findings -- Explaining the implications of the research -- Suggesting concrete next steps""" +Style: +- Base all content solely on the collected evidence. +- Be professional, objective, and specific. +- Avoid meta commentary and refrain from speculation beyond the evidence.""" PRESEARCH_AGENT_PROMPT = """You are a rapid research surveyor. Task: -- Call the gather_context tool once with the main question to obtain a - relevant texts from the Knowledge Base (KB). -- Read that context and produce a brief natural-language summary describing - what the KB appears to contain relative to the question. +- Call gather_context once on the main question to obtain relevant text from + the knowledge base (KB). +- Read that context and produce a short natural‑language summary of what the + KB appears to contain relative to the question. Rules: - Base the summary strictly on the provided text; do not invent. -- Output only the summary as plain text (one short paragraph). -""" +- Output only the summary as plain text (one short paragraph).""" From 42fe2201af75b94152ffcbb5be85651fb49021ed Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 19 Sep 2025 15:46:53 +0300 Subject: [PATCH 3/7] Only configure logfire in cli when not in development --- src/haiku/rag/cli.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index 6c4f5fef..c1073e9b 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -13,10 +13,10 @@ from haiku.rag.logging import configure_cli_logging from haiku.rag.migration import migrate_sqlite_to_lancedb from haiku.rag.utils import is_up_to_date -logfire.configure(send_to_logfire="if-token-present") -logfire.instrument_pydantic_ai() - -if not Config.ENV == "development": +if Config.ENV == "development": + logfire.configure(send_to_logfire="if-token-present") + logfire.instrument_pydantic_ai() +else: warnings.filterwarnings("ignore") cli = typer.Typer( From 9daf14fd42190f700547e3152a3ad008346a788f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 19 Sep 2025 16:28:03 +0300 Subject: [PATCH 4/7] Research graph diagram --- .pre-commit-config.yaml | 10 ---------- docs/agents.md | 14 +++++++++++++- mkdocs.yml | 6 +++++- 3 files changed, 18 insertions(+), 12 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 57298115..19c80c7c 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -20,13 +20,3 @@ repos: rev: v1.1.399 hooks: - id: pyright - - - repo: https://github.com/RodrigoGonzalez/check-mkdocs - rev: v1.2.0 - hooks: - - id: check-mkdocs - name: check-mkdocs - args: ["--config", "mkdocs.yml"] # Optional, mkdocs.yml is the default - # If you have additional plugins or libraries that are not included in - # check-mkdocs, add them here - additional_dependencies: ["mkdocs-material"] diff --git a/docs/agents.md b/docs/agents.md index 2d018701..4d2d82c7 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -38,7 +38,19 @@ print(answer) ### Research Graph -The research workflow is now 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. +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. + +```mermaid +--- +title: Research graph +--- +stateDiagram-v2 + PlanNode --> SearchDispatchNode + SearchDispatchNode --> EvaluateNode + EvaluateNode --> SearchDispatchNode + EvaluateNode --> SynthesizeNode + SynthesizeNode --> [*] +``` Key nodes: diff --git a/mkdocs.yml b/mkdocs.yml index 943f2576..2edd9719 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -76,4 +76,8 @@ markdown_extensions: use_pygments: true - pymdownx.inlinehilite - pymdownx.snippets - - pymdownx.superfences + - pymdownx.superfences: + custom_fences: + - name: mermaid + class: mermaid + format: !!python/name:pymdownx.superfences.fence_code_format From 459004d3ad3e221b807ba6d7687ae4af3ea3f67c Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 19 Sep 2025 16:33:39 +0300 Subject: [PATCH 5/7] Update README --- README.md | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/README.md b/README.md index d6aab4d9..e6851b5f 100644 --- a/README.md +++ b/README.md @@ -11,6 +11,7 @@ Retrieval-Augmented Generation (RAG) library built on LanceDB. - **Local LanceDB**: No external servers required, supports also LanceDB cloud storage, S3, Google Cloud & Azure - **Multiple embedding providers**: Ollama, VoyageAI, OpenAI, vLLM - **Multiple QA providers**: Any provider/model supported by Pydantic AI +- **Research graph (multi‑agent)**: Plan → Search → Evaluate → Synthesize with agentic AI - **Native hybrid search**: Vector + full-text search with native LanceDB RRF reranking - **Reranking**: Default search result reranking with MixedBread AI, Cohere, or vLLM - **Question answering**: Built-in QA agents on your documents @@ -38,6 +39,14 @@ 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 +# Multi‑agent research (iterative plan/search/evaluate) +haiku-rag research \ + "What are the main drivers and trends of global temperature anomalies since 1990?" \ + --max-iterations 2 \ + --confidence-threshold 0.8 \ + --max-concurrency 3 \ + --verbose + # Rebuild database (re-chunk and re-embed all documents) haiku-rag rebuild @@ -53,6 +62,13 @@ haiku-rag serve ```python from haiku.rag.client import HaikuRAG +from haiku.rag.research import ( + ResearchContext, + ResearchDeps, + ResearchState, + build_research_graph, + PlanNode, +) async with HaikuRAG("database.lancedb") as client: # Add document @@ -70,6 +86,25 @@ async with HaikuRAG("database.lancedb") as client: # Ask questions with citations answer = await client.ask("Who is the author of haiku.rag?", cite=True) print(answer) + + # Multi‑agent research pipeline (Plan → Search → Evaluate → Synthesize) + graph = build_research_graph() + state = ResearchState( + question=( + "What are the main drivers and trends of global temperature " + "anomalies since 1990?" + ), + context=ResearchContext(original_question="…"), + max_iterations=2, + confidence_threshold=0.8, + max_concurrency=3, + ) + deps = ResearchDeps(client=client) + start = PlanNode(provider=None, model=None) + result = await graph.run(start, state=state, deps=deps) + report = result.output + print(report.title) + print(report.executive_summary) ``` ## MCP Server From 3fc461e1b12c86832d4278ef384ec300fffb5f79 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 19 Sep 2025 16:54:06 +0300 Subject: [PATCH 6/7] Set default concurrency to 1 --- src/haiku/rag/app.py | 2 +- src/haiku/rag/cli.py | 2 +- src/haiku/rag/research/state.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index c4988612..73e46392 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -90,7 +90,7 @@ class HaikuRAGApp: question: str, max_iterations: int = 3, confidence_threshold: float = 0.8, - max_concurrency: int = 3, + max_concurrency: int = 1, verbose: bool = False, ): """Run research via the pydantic-graph pipeline (default).""" diff --git a/src/haiku/rag/cli.py b/src/haiku/rag/cli.py index c1073e9b..45814204 100644 --- a/src/haiku/rag/cli.py +++ b/src/haiku/rag/cli.py @@ -256,7 +256,7 @@ def research( help="Minimum confidence (0-1) to stop", ), max_concurrency: int = typer.Option( - 3, + 1, "--max-concurrency", help="Max concurrent searches per iteration (planned)", ), diff --git a/src/haiku/rag/research/state.py b/src/haiku/rag/research/state.py index 6085e871..c153b4ca 100644 --- a/src/haiku/rag/research/state.py +++ b/src/haiku/rag/research/state.py @@ -20,6 +20,6 @@ class ResearchState: sub_questions: list[str] = field(default_factory=list) iterations: int = 0 max_iterations: int = 3 - max_concurrency: int = 3 + max_concurrency: int = 1 confidence_threshold: float = 0.8 last_eval: EvaluationResult | None = None From fbb43c06c73d5bc1cdef9fdab1289fe366806af7 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 19 Sep 2025 17:25:04 +0300 Subject: [PATCH 7/7] Cleanup console logging --- src/haiku/rag/app.py | 6 ++++++ src/haiku/rag/research/nodes/evaluate.py | 23 +++++++++++------------ src/haiku/rag/research/nodes/plan.py | 11 +++++------ src/haiku/rag/research/nodes/search.py | 11 +++++------ 4 files changed, 27 insertions(+), 24 deletions(-) diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 73e46392..d0948ea4 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -145,6 +145,12 @@ class HaikuRAGApp: self.console.print(report.executive_summary) self.console.print() + # Confidence (from last evaluation) + if state.last_eval: + conf = state.last_eval.confidence_score # type: ignore[attr-defined] + self.console.print(f"[bold cyan]Confidence:[/bold cyan] {conf:.1%}") + self.console.print() + # Main Findings if report.main_findings: self.console.print("[bold cyan]Main Findings:[/bold cyan]") diff --git a/src/haiku/rag/research/nodes/evaluate.py b/src/haiku/rag/research/nodes/evaluate.py index 79cb80d5..7270d0a0 100644 --- a/src/haiku/rag/research/nodes/evaluate.py +++ b/src/haiku/rag/research/nodes/evaluate.py @@ -57,16 +57,16 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): state.last_eval = output state.iterations += 1 - if deps.console: - if output.key_insights: - deps.console.print(" [bold]Key insights:[/bold]") - for ins in output.key_insights: - deps.console.print(f" • {ins}") - deps.console.print( - f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]" - ) - status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]" - deps.console.print(f" Sufficient: {status}") + if output.key_insights: + log(deps.console, " [bold]Key insights:[/bold]") + for ins in output.key_insights: + log(deps.console, f" • {ins}") + log( + deps.console, + f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]", + ) + status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]" + log(deps.console, f" Sufficient: {status}") from haiku.rag.research.nodes.search import SearchDispatchNode @@ -74,8 +74,7 @@ class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): output.is_sufficient and output.confidence_score >= state.confidence_threshold ) or state.iterations >= state.max_iterations: - if deps.console: - deps.console.print("\n[bold green]✅ Stopping research.[/bold green]") + log(deps.console, "\n[bold green]✅ Stopping research.[/bold green]") return SynthesizeNode(self.provider, self.model) return SearchDispatchNode(self.provider, self.model) diff --git a/src/haiku/rag/research/nodes/plan.py b/src/haiku/rag/research/nodes/plan.py index a412639f..653c12c5 100644 --- a/src/haiku/rag/research/nodes/plan.py +++ b/src/haiku/rag/research/nodes/plan.py @@ -54,11 +54,10 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): plan_result = await plan_agent.run(prompt, deps=agent_deps) state.sub_questions = list(plan_result.output.sub_questions) - if deps.console: - deps.console.print("\n[bold green]✅ Research Plan Created:[/bold green]") - deps.console.print(f" [bold]Main Question:[/bold] {state.question}") - deps.console.print(" [bold]Sub-questions:[/bold]") - for i, sq in enumerate(state.sub_questions, 1): - deps.console.print(f" {i}. {sq}") + log(deps.console, "\n[bold green]✅ Research Plan Created:[/bold green]") + log(deps.console, f" [bold]Main Question:[/bold] {state.question}") + log(deps.console, " [bold]Sub-questions:[/bold]") + for i, sq in enumerate(state.sub_questions, 1): + log(deps.console, f" {i}. {sq}") return SearchDispatchNode(self.provider, self.model) diff --git a/src/haiku/rag/research/nodes/search.py b/src/haiku/rag/research/nodes/search.py index 4640b3c1..0943338a 100644 --- a/src/haiku/rag/research/nodes/search.py +++ b/src/haiku/rag/research/nodes/search.py @@ -36,11 +36,10 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): batch.append(state.sub_questions.pop(0)) async def answer_one(sub_q: str) -> SearchAnswer | None: - if deps.console: - deps.console.print( - f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}" - ) - + log( + deps.console, + 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), @@ -87,6 +86,6 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]): state.context.add_qa_response(ans) if deps.console: preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "") - deps.console.log(f" [green]✓[/green] {preview}") + log(deps.console, f" [green]✓[/green] {preview}") return SearchDispatchNode(self.provider, self.model)