danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #63 from ggozad/feat/improve-research

Browse source 
Improve and clean research multi-agent implementation.
This commit is contained in:
Yiorgis Gozadinos 2025-09-19 12:10:20 +03:00 committed by GitHub
commit eba39e862b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 183 additions and 198 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
docs/agents.md
View file

@ -70,14 +70,14 @@ 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")
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=False,
verbose=True,
)
print(report.title)

1
src/haiku/rag/app.py
View file

@ -101,7 +101,6 @@ class HaikuRAGApp:
client=client,
max_iterations=max_iterations,
verbose=verbose,
console=self.console if verbose else None,
)
# Display the report

4
src/haiku/rag/research/base.py
View file

@ -45,7 +45,8 @@ class BaseResearchAgent[T](ABC):
model=model_obj,
deps_type=ResearchDependencies,
output_type=agent_output_type,
system_prompt=self.get_system_prompt(),
instructions=self.get_system_prompt(),
retries=3,
)
# Register tools
@ -75,7 +76,6 @@ class BaseResearchAgent[T](ABC):
"""Return the system prompt for this agent."""
pass
@abstractmethod
def register_tools(self) -> None:
"""Register agent-specific tools."""
pass

24
src/haiku/rag/research/dependencies.py
View file

@ -1,4 +1,6 @@
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
@ -43,3 +45,25 @@ 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")

51
src/haiku/rag/research/evaluation_agent.py
View file

@ -1,6 +1,11 @@
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
@ -34,9 +39,47 @@ class AnalysisEvaluationAgent(BaseResearchAgent[EvaluationResult]):
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
def register_tools(self) -> None:
"""No additional tools needed - uses LLM capabilities directly."""
pass

180
src/haiku/rag/research/orchestrator.py
View file

@ -1,13 +1,15 @@
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.dependencies import (
ResearchContext,
ResearchDependencies,
)
from haiku.rag.research.evaluation_agent import (
AnalysisEvaluationAgent,
EvaluationResult,
@ -31,7 +33,9 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
"""Orchestrator agent that coordinates the research workflow."""
def __init__(
self, provider: str | None = Config.RESEARCH_PROVIDER, model: str | None = None
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
@ -53,30 +57,15 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
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 _should_stop_research(
self,
evaluation_result: AgentRunResult[EvaluationResult],
confidence_threshold: float,
) -> bool:
"""Determine if research should stop based on evaluation."""
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,
"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")
result = evaluation_result.output
return result.is_sufficient and result.confidence_score >= confidence_threshold
async def conduct_research(
self,
@ -85,7 +74,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
max_iterations: int = 3,
confidence_threshold: float = 0.8,
verbose: bool = False,
console: Console | None = None,
) -> ResearchReport:
"""Conduct comprehensive research on a question.
@ -95,7 +83,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
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
@ -104,44 +91,27 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
# Initialize context
context = ResearchContext(original_question=question)
deps = ResearchDependencies(client=client, context=context)
if verbose:
deps.console = Console()
# Use provided console or create a new one
console = console or Console() if verbose else None
# Run a simple presearch survey to summarize KB context
if console:
console.print(
"\n[bold cyan]🔎 Presearch: summarizing KB context...[/bold cyan]"
)
presearch_result = await self.presearch_agent.run(question, deps=deps)
console = deps.console
# Create initial research plan
if console:
console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]")
# Include the presearch summary to ground the planning step.
planning_context_xml = format_as_xml(
{
"original_question": question,
"presearch_summary": presearch_result.output or "",
},
root_tag="planning_context",
)
# 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"{planning_context_xml}"
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:
@ -152,7 +122,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
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):
@ -163,7 +132,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
# 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]"
@ -171,90 +139,20 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
break
# Use current sub-questions for this iteration
questions_to_search = context.sub_questions
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)
await self.search_agent.run(search_question, deps=deps)
# 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)
evaluation_result = await self.evaluation_agent.run("", deps=deps)
# Check if research is sufficient
if self._should_stop_research(evaluation_result, confidence_threshold):
@ -265,36 +163,8 @@ Evaluate the research progress for the original question and identify any remain
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
"", 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

7
src/haiku/rag/research/presearch_agent.py
View file

@ -15,6 +15,12 @@ class PresearchSurveyAgent(BaseResearchAgent[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:
@ -28,7 +34,6 @@ class PresearchSurveyAgent(BaseResearchAgent[str]):
limit: int = 6,
) -> str:
"""Return verbatim concatenation of relevant chunk texts."""
query = query.replace('"', "")
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)

16
src/haiku/rag/research/search_agent.py
View file

@ -21,10 +21,17 @@ class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]):
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)
console = deps.console
if console:
console.print(f"\t{prompt}")
if result.output:
deps.context.add_qa_response(result.output)
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
@ -41,9 +48,6 @@ class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]):
limit: int = 5,
) -> str:
"""Search the KB and return a concise context pack."""
# Remove quotes from queries as this requires positional indexing in lancedb
# XXX: Investigate how to do that with lancedb
query = query.replace('"', "")
search_results = await ctx.deps.client.search(query, limit=limit)
expanded = await ctx.deps.client.expand_context(search_results)

32
src/haiku/rag/research/synthesis_agent.py
View file

@ -1,6 +1,11 @@
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
@ -30,11 +35,26 @@ class SynthesisAgent(BaseResearchAgent[ResearchReport]):
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
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

2
tests/generate_benchmark_db.py
View file

@ -12,7 +12,7 @@ from haiku.rag.client import HaikuRAG
from haiku.rag.logging import configure_cli_logging
from haiku.rag.qa import get_qa_agent
logfire.configure()
logfire.configure(send_to_logfire="if-token-present")
logfire.instrument_pydantic_ai()
configure_cli_logging()
console = Console()

9
tests/research/test_evaluation_agent.py
View file

@ -1,3 +1,4 @@
from haiku.rag.config import Config
from haiku.rag.research.evaluation_agent import (
AnalysisEvaluationAgent,
EvaluationResult,
@ -8,7 +9,9 @@ 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"
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

33
tests/research/test_orchestrator.py
View file

@ -4,6 +4,7 @@ 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
@ -70,7 +71,9 @@ class TestResearchOrchestrator:
def test_orchestrator_initialization(self):
"""Test that orchestrator initializes all agents correctly."""
orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4")
orchestrator = ResearchOrchestrator(
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
)
# Check all agents are initialized
assert orchestrator.search_agent is not None
@ -78,21 +81,25 @@ class TestResearchOrchestrator:
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"
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="openai", model="gpt-4")
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="openai", model="gpt-4")
orchestrator = ResearchOrchestrator(
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
)
# Get the tools from the agent
tools = orchestrator.agent._function_toolset.tools
@ -103,7 +110,9 @@ class TestResearchOrchestrator:
def test_should_stop_research_logic(self):
"""Test the stopping logic based on EvaluationResult."""
orchestrator = ResearchOrchestrator(provider="openai", model="gpt-4")
orchestrator = ResearchOrchestrator(
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
)
# Create mock evaluation results
from unittest.mock import MagicMock
@ -148,7 +157,9 @@ class TestResearchOrchestrator:
@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")
orchestrator = ResearchOrchestrator(
provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
)
# Setup mock client returns
mock_chunks = [

9
tests/research/test_search_agent.py
View file

@ -1,3 +1,4 @@
from haiku.rag.config import Config
from haiku.rag.research import SearchAnswer, SearchSpecialistAgent
@ -5,7 +6,9 @@ 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"
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

9
tests/research/test_synthesis_agent.py
View file

@ -1,3 +1,4 @@
from haiku.rag.config import Config
from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent
@ -5,7 +6,9 @@ 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"
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