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

Refactor to give its research subagent full responsibility

Browse source
This commit is contained in:
Yiorgis Gozadinos 2025-09-19 11:18:13 +03:00
parent 9c128c1d36
commit a215a64686
No known key found for this signature in database
10 changed files with 155 additions and 130 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
docs/agents.md
View file

@ -70,17 +70,14 @@ from haiku.rag.client import HaikuRAG
from haiku.rag.research import ResearchOrchestrator
client = HaikuRAG(path_to_db)
orchestrator = ResearchOrchestrator(
provider="ollama",
model="gpt-oss",
verbose=True
)
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,
)
print(report.title)

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

@ -86,7 +86,7 @@ class HaikuRAGApp:
async with HaikuRAG(db_path=self.db_path) as client:
try:
# Create orchestrator with default config or fallback to QA
orchestrator = ResearchOrchestrator(verbose=verbose)
orchestrator = ResearchOrchestrator()
if verbose:
self.console.print(
@ -100,6 +100,7 @@ class HaikuRAGApp:
question=question,
client=client,
max_iterations=max_iterations,
verbose=verbose,
)
# Display the report

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

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

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

47
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,5 +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

155
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,
@ -34,7 +36,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
self,
provider: str | None = Config.RESEARCH_PROVIDER,
model: str | None = None,
verbose: bool = False,
):
# Use provided values or fall back to config defaults
provider = provider or Config.RESEARCH_PROVIDER or Config.QA_PROVIDER
@ -52,34 +53,10 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
provider, model
)
self.synthesis_agent: SynthesisAgent = SynthesisAgent(provider, model)
if verbose:
self._console = Console()
self.verbose = verbose
def get_system_prompt(self) -> str:
return ORCHESTRATOR_PROMPT
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")
def _should_stop_research(
self,
evaluation_result: AgentRunResult[EvaluationResult],
@ -90,20 +67,13 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
result = evaluation_result.output
return result.is_sufficient and result.confidence_score >= confidence_threshold
def _log(self, line="", rule=False):
if not self._console:
return
if rule:
self._console.rule(line)
else:
self._console.print(line)
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.
@ -113,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
@ -122,14 +91,16 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
# 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
self._log("\n[bold cyan]📋 Creating research plan...[/bold cyan]")
self._log("\n[bold cyan]🔎 Presearch: summarizing KB context...[/bold cyan]")
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"
@ -143,109 +114,57 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
)
context.sub_questions = plan_result.output.sub_questions
self._log("\n[bold green]✅ Research Plan Created:[/bold green]")
self._log(f" [bold]Main Question:[/bold] {plan_result.output.main_question}")
self._log(" [bold]Sub-questions:[/bold]")
for i, sq in enumerate(plan_result.output.sub_questions, 1):
self._log(f" {i}. {sq}")
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):
self._log(
f"[bold yellow]🔄 Iteration {iteration + 1}/{max_iterations}[/bold yellow]",
rule=True,
)
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
self._log(
"[yellow]No more questions to explore. Concluding research.[/yellow]"
)
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
self._log(
f"\n[bold cyan]🔍 Searching & Answering {len(questions_to_search)} questions:[/bold cyan]"
)
for i, q in enumerate(questions_to_search, 1):
self._log(f" {i}. {q}")
if console:
console.print(
f"\n[bold cyan]🔍 Searching & Answering {len(questions_to_search)} questions:[/bold cyan]"
)
# Run searches for all questions and remove answered ones
for search_question in questions_to_search:
await self.search_agent.run(search_question, deps=deps)
if self._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
)
self._log(f"\n [green]✓[/green] {search_question}")
self._log(f" {answer_preview}")
# Analysis and Evaluation phase
self._log(
"\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,
)
output = evaluation_result.output
if output.key_insights:
self._log(" [bold]Key insights:[/bold]")
for insight in output.key_insights:
self._log(f" • {insight}")
self._log(f" Confidence: [yellow]{output.confidence_score:.1%}[/yellow]")
status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
self._log(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):
self._log(
f"\n[bold green]✅ Stopping research:[/bold green] {evaluation_result.output.reasoning}"
)
if console:
console.print(
f"\n[bold green]✅ Stopping research:[/bold green] {evaluation_result.output.reasoning}"
)
break
# Generate final report
self._log("\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
)
self._log("[bold green]✅ Research complete![/bold green]")
return report_result.output

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)

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

@ -21,9 +21,18 @@ 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.
"""
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:
@ -39,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)

26
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,5 +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

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