Merge pull request #64 from ggozad/feat/research-as-graph

Research agents as a graph
This commit is contained in:
Yiorgis Gozadinos 2025-09-19 17:26:36 +03:00 committed by GitHub
commit b648e08673
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
32 changed files with 852 additions and 1000 deletions

View file

@ -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"]

View file

@ -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 (multiagent)**: 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
# Multiagent 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)
# Multiagent 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

View file

@ -36,50 +36,69 @@ answer = await agent.answer("What is climate change?")
print(answer)
```
### Research MultiAgent
### 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 implemented as a typed pydanticgraph. It plans, searches (in parallel batches), evaluates, and synthesizes into a final report — with clear stop conditions and shared state.
Components:
```mermaid
---
title: Research graph
---
stateDiagram-v2
PlanNode --> SearchDispatchNode
SearchDispatchNode --> EvaluateNode
EvaluateNode --> SearchDispatchNode
EvaluateNode --> SynthesizeNode
SynthesizeNode --> [*]
```
- 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 subquestions
- Analysis & Evaluation: Extracts insights, identifies gaps, proposes new questions
- Synthesis: Produces a final structured research report
Key nodes:
- Plan: builds up to 3 standalone subquestions (uses an internal presearch tool)
- Search (batched): answers subquestions 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, selfcontained queries)
- `SearchAnswer` — produced by the search specialist for each subquestion
- `query: str` — the executed subquestion
- `answer: str` — the agents 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 subquestion (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)
```

View file

@ -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:

View file

@ -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

View file

@ -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"

View file

@ -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 = 1,
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]")
@ -114,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]")

View file

@ -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(
@ -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(
1,
"--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,
)
)

View file

@ -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

View file

@ -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",
]

View file

@ -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,
)

View file

@ -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")

View file

@ -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")

View file

@ -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

View file

@ -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,
]
)

View file

@ -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"
)

View file

@ -0,0 +1,80 @@
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 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
if (
output.is_sufficient
and output.confidence_score >= state.confidence_threshold
) or state.iterations >= state.max_iterations:
log(deps.console, "\n[bold green]✅ Stopping research.[/bold green]")
return SynthesizeNode(self.provider, self.model)
return SearchDispatchNode(self.provider, self.model)

View file

@ -0,0 +1,63 @@
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 PLAN_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=(
PLAN_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)
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)

View file

@ -0,0 +1,91 @@
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:
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),
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 "")
log(deps.console, f" [green]✓[/green] {preview}")
return SearchDispatchNode(self.provider, self.model)

View file

@ -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)

View file

@ -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

View file

@ -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)

View file

@ -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, highleverage 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, selfcontained query that can run
without extra context. Include concrete entities, scope, timeframe, and any
qualifiers. Avoid ambiguous pronouns (it/they/this/that).
- Prioritize the highestvalue 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 questionanswering 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 14).
- Set SearchAnswer.sources to the corresponding document_uris for the snippets
you used (one URI per snippet; same order as context). Context must be textonly.
- 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
- Questionanswer 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, nonobvious 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 questions 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 highestvalue 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 evidencebased 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 (512 words), informative.
- executive_summary: 35 sentences summarizing the overall answer.
- main_findings: 48 onesentence bullets; each reflects evidence from the
research (do not include inline citations or snippet text).
- conclusions: 24 bullets that follow logically from findings.
- recommendations: 25 actionable bullets tied to findings.
- limitations: 13 bullets describing key constraints or uncertainties.
- sources_summary: 24 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 naturallanguage 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)."""

View file

@ -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}"
)

View file

@ -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 = 1
confidence_threshold: float = 0.8
last_eval: EvaluationResult | None = None

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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

View file

@ -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()

View file

@ -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