Merge pull request #58 from ggozad/feat/agentic-research

Multi-agent (agentic) research
This commit is contained in:
Yiorgis Gozadinos 2025-09-17 13:12:20 +03:00 committed by GitHub
commit 0e6f04b19e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
25 changed files with 2208 additions and 914 deletions

2
.gitignore vendored
View file

@ -18,3 +18,5 @@ tests/data/
# environment variables # environment variables
.env .env
TODO.md TODO.md
PLAN.md
DEVNOTES.md

View file

@ -90,4 +90,5 @@ Full documentation at: https://ggozad.github.io/haiku.rag/
- [Configuration](https://ggozad.github.io/haiku.rag/configuration/) - Environment variables - [Configuration](https://ggozad.github.io/haiku.rag/configuration/) - Environment variables
- [CLI](https://ggozad.github.io/haiku.rag/cli/) - Command reference - [CLI](https://ggozad.github.io/haiku.rag/cli/) - Command reference
- [Python API](https://ggozad.github.io/haiku.rag/python/) - Complete API docs - [Python API](https://ggozad.github.io/haiku.rag/python/) - Complete API docs
- [Agents](https://ggozad.github.io/haiku.rag/agents/) - QA agent and multi-agent research
- [Benchmarks](https://ggozad.github.io/haiku.rag/benchmarks/) - Performance Benchmarks - [Benchmarks](https://ggozad.github.io/haiku.rag/benchmarks/) - Performance Benchmarks

83
docs/agents.md Normal file
View file

@ -0,0 +1,83 @@
## Agents
Two agentic flows are provided by haiku.rag:
- Simple QA Agent — a focused question answering agent
- Research MultiAgent — a multistep, analyzable research workflow
### Simple QA Agent
The simple QA agent answers a single question using the knowledge base. It retrieves relevant chunks, optionally expands context around them, and asks the model to answer strictly based on that context.
Key points:
- Uses a single `search_documents` tool to fetch relevant chunks
- Can be run with or without inline citations in the prompt
- Returns a plain string answer
Python usage:
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.qa.agent import QuestionAnswerAgent
client = HaikuRAG(path_to_db)
# Choose a provider and model (see Configuration for env defaults)
agent = QuestionAnswerAgent(
client=client,
provider="openai", # or "ollama", "vllm", etc.
model="gpt-4o-mini",
use_citations=False, # set True to bias prompt towards citing sources
)
answer = await agent.answer("What is climate change?")
print(answer)
```
### Research MultiAgent
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.
Components:
- Orchestrator: Plans, coordinates, and loops until confidence is sufficient
- 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
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
Python usage:
```python
from haiku.rag.client import HaikuRAG
from haiku.rag.research import ResearchOrchestrator
client = HaikuRAG(path_to_db)
orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4o-mini")
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=False,
)
print(report.title)
print(report.executive_summary)
```

View file

@ -55,6 +55,7 @@ haiku-rag migrate old_database.sqlite # Migrate from SQLite
- [Server](server.md) - File monitoring and server mode - [Server](server.md) - File monitoring and server mode
- [MCP](mcp.md) - Model Context Protocol integration - [MCP](mcp.md) - Model Context Protocol integration
- [Python](python.md) - Python API reference - [Python](python.md) - Python API reference
- [Agents](agents.md) - QA agent and multi-agent research
## License ## License

View file

@ -204,3 +204,5 @@ print(answer)
The QA agent will search your documents for relevant information and use the configured LLM to generate a comprehensive answer. With `cite=True`, responses include citations showing which documents were used as sources. The QA agent will search your documents for relevant information and use the configured LLM to generate a comprehensive answer. With `cite=True`, responses include citations showing which documents were used as sources.
The QA provider and model can be configured via environment variables (see [Configuration](configuration.md)). The QA provider and model can be configured via environment variables (see [Configuration](configuration.md)).
See also: [Agents](agents.md) for details on the QA agent and the multiagent research workflow.

View file

@ -61,8 +61,9 @@ nav:
- Configuration: configuration.md - Configuration: configuration.md
- CLI: cli.md - CLI: cli.md
- Server: server.md - Server: server.md
- MCP: mcp.md - Agents: agents.md
- Python: python.md - Python: python.md
- MCP: mcp.md
- Benchmarks: benchmarks.md - Benchmarks: benchmarks.md
markdown_extensions: markdown_extensions:
- admonition - admonition

View file

@ -51,16 +51,16 @@ packages = ["src/haiku"]
[dependency-groups] [dependency-groups]
dev = [ dev = [
"datasets>=3.6.0", "datasets>=4.1.0",
"logfire>=4.6.0", "logfire>=4.7.0",
"mkdocs>=1.6.1", "mkdocs>=1.6.1",
"mkdocs-material>=9.6.14", "mkdocs-material>=9.6.14",
"pre-commit>=4.2.0", "pre-commit>=4.2.0",
"pyright>=1.1.404", "pyright>=1.1.405",
"pytest>=8.4.0", "pytest>=8.4.2",
"pytest-asyncio>=1.0.0", "pytest-asyncio>=1.2.0",
"pytest-cov>=6.2.1", "pytest-cov>=7.0.0",
"ruff>=0.11.13", "ruff>=0.13.0",
] ]
[tool.ruff] [tool.ruff]

View file

@ -9,6 +9,7 @@ from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config from haiku.rag.config import Config
from haiku.rag.mcp import create_mcp_server from haiku.rag.mcp import create_mcp_server
from haiku.rag.monitor import FileWatcher from haiku.rag.monitor import FileWatcher
from haiku.rag.research.orchestrator import ResearchOrchestrator
from haiku.rag.store.models.chunk import Chunk from haiku.rag.store.models.chunk import Chunk
from haiku.rag.store.models.document import Document from haiku.rag.store.models.document import Document
@ -78,6 +79,85 @@ class HaikuRAGApp:
except Exception as e: except Exception as e:
self.console.print(f"[red]Error: {e}[/red]") self.console.print(f"[red]Error: {e}[/red]")
async def research(
self, question: str, max_iterations: int = 3, verbose: bool = False
):
"""Run multi-agent research on a question."""
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(f"[bold blue]Question:[/bold blue] {question}")
self.console.print()
# Conduct research
report = await orchestrator.conduct_research(
question=question,
client=client,
max_iterations=max_iterations,
verbose=verbose,
console=self.console if verbose else None,
)
# Display the report
self.console.print("[bold green]Research Report[/bold green]")
self.console.rule()
# Title and Executive Summary
self.console.print(f"[bold]{report.title}[/bold]")
self.console.print()
self.console.print("[bold cyan]Executive Summary:[/bold cyan]")
self.console.print(report.executive_summary)
self.console.print()
# Main Findings
if report.main_findings:
self.console.print("[bold cyan]Main Findings:[/bold cyan]")
for finding in report.main_findings:
self.console.print(f"{finding}")
self.console.print()
# Themes
if report.themes:
self.console.print("[bold cyan]Key Themes:[/bold cyan]")
for theme, explanation in report.themes.items():
self.console.print(f"• [bold]{theme}[/bold]: {explanation}")
self.console.print()
# Conclusions
if report.conclusions:
self.console.print("[bold cyan]Conclusions:[/bold cyan]")
for conclusion in report.conclusions:
self.console.print(f"{conclusion}")
self.console.print()
# Recommendations
if report.recommendations:
self.console.print("[bold cyan]Recommendations:[/bold cyan]")
for rec in report.recommendations:
self.console.print(f"{rec}")
self.console.print()
# Limitations
if report.limitations:
self.console.print("[bold yellow]Limitations:[/bold yellow]")
for limitation in report.limitations:
self.console.print(f"{limitation}")
self.console.print()
# Sources Summary
if report.sources_summary:
self.console.print("[bold cyan]Sources:[/bold cyan]")
self.console.print(report.sources_summary)
except Exception as e:
self.console.print(f"[red]Error during research: {e}[/red]")
async def rebuild(self): async def rebuild(self):
async with HaikuRAG(db_path=self.db_path, skip_validation=True) as client: async with HaikuRAG(db_path=self.db_path, skip_validation=True) as client:
try: try:

View file

@ -3,6 +3,7 @@ import warnings
from importlib.metadata import version from importlib.metadata import version
from pathlib import Path from pathlib import Path
import logfire
import typer import typer
from rich.console import Console from rich.console import Console
@ -12,6 +13,9 @@ from haiku.rag.logging import configure_cli_logging
from haiku.rag.migration import migrate_sqlite_to_lancedb from haiku.rag.migration import migrate_sqlite_to_lancedb
from haiku.rag.utils import is_up_to_date 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 not Config.ENV == "development":
warnings.filterwarnings("ignore") warnings.filterwarnings("ignore")
@ -235,6 +239,38 @@ def ask(
asyncio.run(app.ask(question=question, cite=cite)) asyncio.run(app.ask(question=question, cite=cite))
@cli.command("research", help="Run multi-agent research and output a concise report")
def research(
question: str = typer.Argument(
help="The research question to investigate",
),
max_iterations: int = typer.Option(
3,
"--max-iterations",
"-n",
help="Maximum search/analyze iterations",
),
db: Path = typer.Option(
Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
"--db",
help="Path to the LanceDB database file",
),
verbose: bool = typer.Option(
False,
"--verbose",
help="Show verbose progress output",
),
):
app = HaikuRAGApp(db_path=db)
asyncio.run(
app.research(
question=question,
max_iterations=max_iterations,
verbose=verbose,
)
)
@cli.command("settings", help="Display current configuration settings") @cli.command("settings", help="Display current configuration settings")
def settings(): def settings():
app = HaikuRAGApp(db_path=Path()) # Don't need actual DB for settings app = HaikuRAGApp(db_path=Path()) # Don't need actual DB for settings

View file

@ -27,7 +27,11 @@ class AppConfig(BaseModel):
RERANK_MODEL: str = "" RERANK_MODEL: str = ""
QA_PROVIDER: str = "ollama" QA_PROVIDER: str = "ollama"
QA_MODEL: str = "qwen3" QA_MODEL: str = "gpt-oss"
# Research defaults (fallback to QA if not provided via env)
RESEARCH_PROVIDER: str = "ollama"
RESEARCH_MODEL: str = "gpt-oss"
CHUNK_SIZE: int = 256 CHUNK_SIZE: int = 256
CONTEXT_CHUNK_RADIUS: int = 0 CONTEXT_CHUNK_RADIUS: int = 0
@ -37,9 +41,11 @@ class AppConfig(BaseModel):
MARKDOWN_PREPROCESSOR: str = "" MARKDOWN_PREPROCESSOR: str = ""
OLLAMA_BASE_URL: str = "http://localhost:11434" OLLAMA_BASE_URL: str = "http://localhost:11434"
VLLM_EMBEDDINGS_BASE_URL: str = "" VLLM_EMBEDDINGS_BASE_URL: str = ""
VLLM_RERANK_BASE_URL: str = "" VLLM_RERANK_BASE_URL: str = ""
VLLM_QA_BASE_URL: str = "" VLLM_QA_BASE_URL: str = ""
VLLM_RESEARCH_BASE_URL: str = ""
# Provider keys # Provider keys
VOYAGE_API_KEY: str = "" VOYAGE_API_KEY: str = ""

View file

@ -6,7 +6,7 @@ from pydantic_ai.providers.openai import OpenAIProvider
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config from haiku.rag.config import Config
from haiku.rag.qa.prompts import SYSTEM_PROMPT, SYSTEM_PROMPT_WITH_CITATIONS from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIONS
class SearchResult(BaseModel): class SearchResult(BaseModel):
@ -31,7 +31,9 @@ class QuestionAnswerAgent:
): ):
self._client = client self._client = client
system_prompt = SYSTEM_PROMPT_WITH_CITATIONS if use_citations else SYSTEM_PROMPT system_prompt = (
QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT
)
model_obj = self._get_model(provider, model) model_obj = self._get_model(provider, model)
self._agent = Agent( self._agent = Agent(

View file

@ -1,4 +1,4 @@
SYSTEM_PROMPT = """ QA_SYSTEM_PROMPT = """
You are a knowledgeable assistant that helps users find information from a document knowledge base. You are a knowledgeable assistant that helps users find information from a document knowledge base.
Your process: Your process:
@ -21,7 +21,7 @@ Be concise, and always maintain accuracy over completeness. Prefer short, direct
/no_think /no_think
""" """
SYSTEM_PROMPT_WITH_CITATIONS = """ QA_SYSTEM_PROMPT_WITH_CITATIONS = """
You are a knowledgeable assistant that helps users find information from a document knowledge base. You are a knowledgeable assistant that helps users find information from a document knowledge base.
IMPORTANT: You MUST use the search_documents tool for every question. Do not answer any question without first searching the knowledge base. IMPORTANT: You MUST use the search_documents tool for every question. Do not answer any question without first searching the knowledge base.

View file

@ -0,0 +1,35 @@
"""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.orchestrator import ResearchOrchestrator, ResearchPlan
from haiku.rag.research.search_agent import SearchSpecialistAgent
from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent
__all__ = [
# Base classes
"BaseResearchAgent",
"ResearchDependencies",
"ResearchContext",
"SearchResult",
"ResearchOutput",
# Specialized agents
"SearchAnswer",
"SearchSpecialistAgent",
"AnalysisEvaluationAgent",
"EvaluationResult",
"SynthesisAgent",
"ResearchReport",
# Orchestrator
"ResearchOrchestrator",
"ResearchPlan",
]

View file

@ -0,0 +1,122 @@
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
self._agent = Agent(
model=model_obj,
deps_type=ResearchDependencies,
output_type=ToolOutput(self.output_type, max_retries=3),
system_prompt=self.get_system_prompt(),
)
# 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
@abstractmethod
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,45 @@
from pydantic import BaseModel, Field
from haiku.rag.client import HaikuRAG
from haiku.rag.research.base import SearchAnswer
class ResearchContext(BaseModel):
"""Context shared across research agents."""
original_question: str = Field(description="The original research question")
sub_questions: list[str] = Field(
default_factory=list, description="Decomposed sub-questions"
)
qa_responses: list["SearchAnswer"] = Field(
default_factory=list, description="Structured QA pairs used during research"
)
insights: list[str] = Field(
default_factory=list, description="Key insights discovered"
)
gaps: list[str] = Field(
default_factory=list, description="Identified information gaps"
)
def add_qa_response(self, qa: "SearchAnswer") -> None:
"""Add a structured QA response (minimal context already included)."""
self.qa_responses.append(qa)
def add_insight(self, insight: str) -> None:
"""Add a key insight."""
if insight not in self.insights:
self.insights.append(insight)
def add_gap(self, gap: str) -> None:
"""Identify an information gap."""
if gap not in self.gaps:
self.gaps.append(gap)
class ResearchDependencies(BaseModel):
"""Dependencies for research agents with multi-agent context."""
model_config = {"arbitrary_types_allowed": True}
client: HaikuRAG = Field(description="RAG client for document operations")
context: ResearchContext = Field(description="Shared research context")

View file

@ -0,0 +1,40 @@
from pydantic import BaseModel, Field
from haiku.rag.research.base import BaseResearchAgent
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
)
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)
def get_system_prompt(self) -> str:
return EVALUATION_AGENT_PROMPT
def register_tools(self) -> None:
"""No additional tools needed - uses LLM capabilities directly."""
pass

View file

@ -0,0 +1,265 @@
from typing import Any
from pydantic import BaseModel, Field
from pydantic_ai.format_prompt import format_as_xml
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.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.evaluation_agent: AnalysisEvaluationAgent = AnalysisEvaluationAgent(
provider, model
)
self.synthesis_agent: SynthesisAgent = SynthesisAgent(provider, model)
def get_system_prompt(self) -> str:
return ORCHESTRATOR_PROMPT
def register_tools(self) -> None:
"""Register orchestration tools."""
# Tools are no longer needed - orchestrator directly calls agents
pass
def _format_context_for_prompt(self, 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}
for qa in context.qa_responses
],
"insights": context.insights,
"gaps": context.gaps,
}
return format_as_xml(context_data, root_tag="research_context")
async def conduct_research(
self,
question: str,
client: Any,
max_iterations: int = 3,
confidence_threshold: float = 0.8,
verbose: bool = False,
console: Console | None = None,
) -> 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
console: Optional Rich console for output
Returns:
ResearchReport with comprehensive findings
"""
# Initialize context
context = ResearchContext(original_question=question)
deps = ResearchDependencies(client=client, context=context)
# Use provided console or create a new one
console = console or Console() if verbose else None
# Create initial research plan
if console:
console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]")
plan_result: AgentRunResult[ResearchPlan] = await self.run(
f"Create a research plan for: {question}", 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}")
console.print()
# 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:
# No more questions to explore
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 i, q in enumerate(questions_to_search, 1):
console.print(f" {i}. {q}")
# Run searches for all questions and remove answered ones
answered_questions = []
for search_question in questions_to_search:
try:
await self.search_agent.run(search_question, deps=deps)
except Exception as e: # pragma: no cover - defensive
if console:
console.print(
f"\n [red]×[/red] Omitting failed question: {search_question} ({e})"
)
finally:
answered_questions.append(search_question)
if console and context.qa_responses:
# Show the last QA response (which should be for this question)
latest_qa = context.qa_responses[-1]
answer_preview = (
latest_qa.answer[:150] + "..."
if len(latest_qa.answer) > 150
else latest_qa.answer
)
console.print(
f"\n [green]✓[/green] {search_question[:50]}..."
if len(search_question) > 50
else f"\n [green]✓[/green] {search_question}"
)
console.print(f" {answer_preview}")
# Remove answered questions from the list
for question in answered_questions:
if question in context.sub_questions:
context.sub_questions.remove(question)
# Analysis and Evaluation phase
if console:
console.print(
"\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]"
)
# Format context for the evaluation agent
context_xml = self._format_context_for_prompt(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."""
evaluation_result = await self.evaluation_agent.run(
evaluation_prompt,
deps=deps,
)
if console and evaluation_result.output:
output = evaluation_result.output
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}")
# Store insights
for insight in evaluation_result.output.key_insights:
context.add_insight(insight)
# Add new questions to the sub-questions list
for new_q in evaluation_result.output.new_questions:
if new_q not in context.sub_questions:
context.sub_questions.append(new_q)
# 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
if console:
console.print(
"\n[bold cyan]📝 Generating final research report...[/bold cyan]"
)
# Format context for the synthesis agent
final_context_xml = self._format_context_for_prompt(context)
synthesis_prompt = f"""Generate a comprehensive research report based on all gathered information.
{final_context_xml}
Create a detailed report that synthesizes all findings into a coherent response."""
report_result: AgentRunResult[ResearchReport] = await self.synthesis_agent.run(
synthesis_prompt, deps=deps
)
if console:
console.print("[bold green]✅ Research complete![/bold green]")
return report_result.output
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
# Stop if the agent indicates sufficient information AND confidence exceeds threshold
return result.is_sufficient and result.confidence_score >= confidence_threshold

View file

@ -0,0 +1,116 @@
ORCHESTRATOR_PROMPT = """You are a research orchestrator responsible for coordinating a comprehensive research 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
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."""
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)
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.
- 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.
- 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.
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."""
EVALUATION_AGENT_PROMPT = """You are an analysis and evaluation specialist for research workflows.
You have access to:
- The original research question
- Question-answer pairs from search operations
- Raw search results and source documents
- 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
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
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
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."""
SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist agent focused on creating comprehensive research reports.
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
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
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"""

View file

@ -0,0 +1,64 @@
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.
"""
result = await super().run(prompt, deps, **kwargs)
if result.output:
deps.context.add_qa_response(result.output)
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."""
# Remove quotes from queries as this requires positional indexing in lancedb
query = query.replace('"', "")
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,39 @@
from pydantic import BaseModel, Field
from haiku.rag.research.base import BaseResearchAgent
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"
)
themes: dict[str, str] = Field(description="Major themes and their explanations")
conclusions: list[str] = Field(description="Evidence-based conclusions")
limitations: list[str] = Field(description="Limitations of the current research")
recommendations: list[str] = Field(
description="Actionable recommendations based on findings"
)
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)
def get_system_prompt(self) -> str:
return SYNTHESIS_AGENT_PROMPT
def register_tools(self) -> None:
"""Register synthesis-specific tools."""
# The agent will use its LLM capabilities directly for synthesis
# The structured output will guide the report generation
pass

View file

@ -0,0 +1,14 @@
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="openai", model="gpt-4")
assert agent.provider == "openai"
assert agent.model == "gpt-4"
assert agent.output_type == EvaluationResult

View file

@ -0,0 +1,179 @@
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.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="openai", model="gpt-4")
# 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 == "openai"
assert orchestrator.search_agent.model == "gpt-4"
assert orchestrator.evaluation_agent.provider == "openai"
assert orchestrator.evaluation_agent.model == "gpt-4"
assert orchestrator.synthesis_agent.provider == "openai"
assert orchestrator.synthesis_agent.model == "gpt-4"
def test_orchestrator_has_correct_output_type(self):
"""Test that orchestrator's output type is ResearchPlan."""
orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4")
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="openai", model="gpt-4")
# 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="openai", model="gpt-4")
# 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="openai", model="gpt-4")
# 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.themes, dict)
assert isinstance(report.conclusions, list)
assert isinstance(report.limitations, list)
assert isinstance(report.recommendations, list)
assert report.sources_summary

View file

@ -0,0 +1,11 @@
from haiku.rag.research import SearchAnswer, SearchSpecialistAgent
class TestSearchSpecialistAgent:
"""Lean tests for SearchSpecialistAgent without LLM mocking."""
def test_agent_initialization(self):
agent = SearchSpecialistAgent(provider="openai", model="gpt-4")
assert agent.provider == "openai"
assert agent.model == "gpt-4"
assert agent.output_type is SearchAnswer

View file

@ -0,0 +1,11 @@
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="openai", model="gpt-4")
assert agent.provider == "openai"
assert agent.model == "gpt-4"
assert agent.output_type == ResearchReport

1941
uv.lock

File diff suppressed because it is too large Load diff