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

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

View file

@ -70,14 +70,14 @@ from haiku.rag.client import HaikuRAG
from haiku.rag.research import ResearchOrchestrator from haiku.rag.research import ResearchOrchestrator
client = HaikuRAG(path_to_db) 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( report = await orchestrator.conduct_research(
question="What are the main drivers and recent trends of global temperature anomalies since 1990?", question="What are the main drivers and recent trends of global temperature anomalies since 1990?",
client=client, client=client,
max_iterations=2, max_iterations=2,
confidence_threshold=0.8, confidence_threshold=0.8,
verbose=False, verbose=True,
) )
print(report.title) print(report.title)

View file

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

View file

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

View file

@ -1,4 +1,6 @@
from pydantic import BaseModel, Field 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.client import HaikuRAG
from haiku.rag.research.base import SearchAnswer from haiku.rag.research.base import SearchAnswer
@ -43,3 +45,25 @@ class ResearchDependencies(BaseModel):
client: HaikuRAG = Field(description="RAG client for document operations") client: HaikuRAG = Field(description="RAG client for document operations")
context: ResearchContext = Field(description="Shared research context") 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,6 +1,11 @@
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from pydantic_ai.run import AgentRunResult
from haiku.rag.research.base import BaseResearchAgent 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 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: def __init__(self, provider: str, model: str) -> None:
super().__init__(provider, model, output_type=EvaluationResult) 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: def get_system_prompt(self) -> str:
return EVALUATION_AGENT_PROMPT return EVALUATION_AGENT_PROMPT
def register_tools(self) -> None:
"""No additional tools needed - uses LLM capabilities directly."""
pass

View file

@ -1,13 +1,15 @@
from typing import Any from typing import Any
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from pydantic_ai.format_prompt import format_as_xml
from pydantic_ai.run import AgentRunResult from pydantic_ai.run import AgentRunResult
from rich.console import Console from rich.console import Console
from haiku.rag.config import Config from haiku.rag.config import Config
from haiku.rag.research.base import BaseResearchAgent 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 ( from haiku.rag.research.evaluation_agent import (
AnalysisEvaluationAgent, AnalysisEvaluationAgent,
EvaluationResult, EvaluationResult,
@ -31,7 +33,9 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
"""Orchestrator agent that coordinates the research workflow.""" """Orchestrator agent that coordinates the research workflow."""
def __init__( 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 # Use provided values or fall back to config defaults
provider = provider or Config.RESEARCH_PROVIDER or Config.QA_PROVIDER provider = provider or Config.RESEARCH_PROVIDER or Config.QA_PROVIDER
@ -53,30 +57,15 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
def get_system_prompt(self) -> str: def get_system_prompt(self) -> str:
return ORCHESTRATOR_PROMPT return ORCHESTRATOR_PROMPT
def register_tools(self) -> None: def _should_stop_research(
"""Register orchestration tools.""" self,
# Tools are no longer needed - orchestrator directly calls agents evaluation_result: AgentRunResult[EvaluationResult],
pass confidence_threshold: float,
) -> bool:
"""Determine if research should stop based on evaluation."""
def _format_context_for_prompt(self, context: ResearchContext) -> str: result = evaluation_result.output
"""Format the research context as XML for inclusion in prompts.""" return result.is_sufficient and result.confidence_score >= confidence_threshold
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")
async def conduct_research( async def conduct_research(
self, self,
@ -85,7 +74,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
max_iterations: int = 3, max_iterations: int = 3,
confidence_threshold: float = 0.8, confidence_threshold: float = 0.8,
verbose: bool = False, verbose: bool = False,
console: Console | None = None,
) -> ResearchReport: ) -> ResearchReport:
"""Conduct comprehensive research on a question. """Conduct comprehensive research on a question.
@ -95,7 +83,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
max_iterations: Maximum number of search-analyze-clarify cycles max_iterations: Maximum number of search-analyze-clarify cycles
confidence_threshold: Minimum confidence level to stop research (0-1) confidence_threshold: Minimum confidence level to stop research (0-1)
verbose: If True, print progress and intermediate results verbose: If True, print progress and intermediate results
console: Optional Rich console for output
Returns: Returns:
ResearchReport with comprehensive findings ResearchReport with comprehensive findings
@ -104,44 +91,27 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
# Initialize context # Initialize context
context = ResearchContext(original_question=question) context = ResearchContext(original_question=question)
deps = ResearchDependencies(client=client, context=context) deps = ResearchDependencies(client=client, context=context)
if verbose:
deps.console = Console()
# Use provided console or create a new one console = deps.console
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)
# Create initial research plan # Create initial research plan
if console: if console:
console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]") console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]")
# Include the presearch summary to ground the planning step. # Run a simple presearch survey to summarize KB context
presearch_result = await self.presearch_agent.run(question, deps=deps)
planning_context_xml = format_as_xml(
{
"original_question": question,
"presearch_summary": presearch_result.output or "",
},
root_tag="planning_context",
)
plan_prompt = ( plan_prompt = (
"Create a research plan for the main question below.\n\n" "Create a research plan for the main question below.\n\n"
f"Main question: {question}\n\n" f"Main question: {question}\n\n"
"Use this brief presearch summary to inform the plan. Focus the 3 sub-questions " "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" "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_result: AgentRunResult[ResearchPlan] = await self.run(
plan_prompt, deps=deps plan_prompt, deps=deps
) )
context.sub_questions = plan_result.output.sub_questions context.sub_questions = plan_result.output.sub_questions
if console: if console:
@ -152,7 +122,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
console.print(" [bold]Sub-questions:[/bold]") console.print(" [bold]Sub-questions:[/bold]")
for i, sq in enumerate(plan_result.output.sub_questions, 1): for i, sq in enumerate(plan_result.output.sub_questions, 1):
console.print(f" {i}. {sq}") console.print(f" {i}. {sq}")
console.print()
# Execute research iterations # Execute research iterations
for iteration in range(max_iterations): for iteration in range(max_iterations):
@ -163,7 +132,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
# Check if we have questions to search # Check if we have questions to search
if not context.sub_questions: if not context.sub_questions:
# No more questions to explore
if console: if console:
console.print( console.print(
"[yellow]No more questions to explore. Concluding research.[/yellow]" "[yellow]No more questions to explore. Concluding research.[/yellow]"
@ -171,90 +139,20 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
break break
# Use current sub-questions for this iteration # 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 # Search phase - answer all questions in this iteration
if console: if console:
console.print( console.print(
f"\n[bold cyan]🔍 Searching & Answering {len(questions_to_search)} questions:[/bold cyan]" 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: for search_question in questions_to_search:
try: await self.search_agent.run(search_question, deps=deps)
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 # Analysis and Evaluation phase
if console:
console.print(
"\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]"
)
# Format context for the evaluation agent evaluation_result = await self.evaluation_agent.run("", deps=deps)
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 # Check if research is sufficient
if self._should_stop_research(evaluation_result, confidence_threshold): 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 break
# Generate final report # 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( 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 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

@ -15,6 +15,12 @@ class PresearchSurveyAgent(BaseResearchAgent[str]):
async def run( async def run(
self, prompt: str, deps: ResearchDependencies, **kwargs self, prompt: str, deps: ResearchDependencies, **kwargs
) -> AgentRunResult[str]: ) -> 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) return await super().run(prompt, deps, **kwargs)
def get_system_prompt(self) -> str: def get_system_prompt(self) -> str:
@ -28,7 +34,6 @@ class PresearchSurveyAgent(BaseResearchAgent[str]):
limit: int = 6, limit: int = 6,
) -> str: ) -> str:
"""Return verbatim concatenation of relevant chunk texts.""" """Return verbatim concatenation of relevant chunk texts."""
query = query.replace('"', "")
results = await ctx.deps.client.search(query, limit=limit) results = await ctx.deps.client.search(query, limit=limit)
expanded = await ctx.deps.client.expand_context(results) expanded = await ctx.deps.client.expand_context(results)
return "\n\n".join(chunk.content for chunk, _ in expanded) return "\n\n".join(chunk.content for chunk, _ in expanded)

View file

@ -21,10 +21,17 @@ class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]):
Pydantic AI enforces `SearchAnswer` as the output model; we just store Pydantic AI enforces `SearchAnswer` as the output model; we just store
the QA response with the last search results as sources. 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: result = await super().run(prompt, deps, **kwargs)
deps.context.add_qa_response(result.output) 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 return result
@ -41,9 +48,6 @@ class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]):
limit: int = 5, limit: int = 5,
) -> str: ) -> str:
"""Search the KB and return a concise context pack.""" """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) search_results = await ctx.deps.client.search(query, limit=limit)
expanded = await ctx.deps.client.expand_context(search_results) expanded = await ctx.deps.client.expand_context(search_results)

View file

@ -1,6 +1,11 @@
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
from pydantic_ai.run import AgentRunResult
from haiku.rag.research.base import BaseResearchAgent 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 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: def __init__(self, provider: str, model: str) -> None:
super().__init__(provider, model, output_type=ResearchReport) 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: def get_system_prompt(self) -> str:
return SYNTHESIS_AGENT_PROMPT 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

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

View file

@ -1,3 +1,4 @@
from haiku.rag.config import Config
from haiku.rag.research.evaluation_agent import ( from haiku.rag.research.evaluation_agent import (
AnalysisEvaluationAgent, AnalysisEvaluationAgent,
EvaluationResult, EvaluationResult,
@ -8,7 +9,9 @@ class TestAnalysisEvaluationAgent:
"""Lean tests for AnalysisEvaluationAgent without LLM mocking.""" """Lean tests for AnalysisEvaluationAgent without LLM mocking."""
def test_agent_initialization(self): def test_agent_initialization(self):
agent = AnalysisEvaluationAgent(provider="openai", model="gpt-4") agent = AnalysisEvaluationAgent(
assert agent.provider == "openai" provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
assert agent.model == "gpt-4" )
assert agent.provider == Config.RESEARCH_PROVIDER
assert agent.model == Config.RESEARCH_MODEL
assert agent.output_type == EvaluationResult assert agent.output_type == EvaluationResult

View file

@ -4,6 +4,7 @@ import pytest
from pydantic_ai.models.test import TestModel from pydantic_ai.models.test import TestModel
from haiku.rag.client import HaikuRAG from haiku.rag.client import HaikuRAG
from haiku.rag.config import Config
from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
from haiku.rag.research.evaluation_agent import EvaluationResult from haiku.rag.research.evaluation_agent import EvaluationResult
from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan
@ -70,7 +71,9 @@ class TestResearchOrchestrator:
def test_orchestrator_initialization(self): def test_orchestrator_initialization(self):
"""Test that orchestrator initializes all agents correctly.""" """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 # Check all agents are initialized
assert orchestrator.search_agent is not None assert orchestrator.search_agent is not None
@ -78,21 +81,25 @@ class TestResearchOrchestrator:
assert orchestrator.synthesis_agent is not None assert orchestrator.synthesis_agent is not None
# Check they all use the same provider and model # Check they all use the same provider and model
assert orchestrator.search_agent.provider == "openai" assert orchestrator.search_agent.provider == Config.RESEARCH_PROVIDER
assert orchestrator.search_agent.model == "gpt-4" assert orchestrator.search_agent.model == Config.RESEARCH_MODEL
assert orchestrator.evaluation_agent.provider == "openai" assert orchestrator.evaluation_agent.provider == Config.RESEARCH_PROVIDER
assert orchestrator.evaluation_agent.model == "gpt-4" assert orchestrator.evaluation_agent.model == Config.RESEARCH_MODEL
assert orchestrator.synthesis_agent.provider == "openai" assert orchestrator.synthesis_agent.provider == Config.RESEARCH_PROVIDER
assert orchestrator.synthesis_agent.model == "gpt-4" assert orchestrator.synthesis_agent.model == Config.RESEARCH_MODEL
def test_orchestrator_has_correct_output_type(self): def test_orchestrator_has_correct_output_type(self):
"""Test that orchestrator's output type is ResearchPlan.""" """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 assert orchestrator.output_type == ResearchPlan
def test_orchestrator_has_no_tools(self): def test_orchestrator_has_no_tools(self):
"""Test that orchestrator no longer registers tools (direct agent calls now).""" """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 # Get the tools from the agent
tools = orchestrator.agent._function_toolset.tools tools = orchestrator.agent._function_toolset.tools
@ -103,7 +110,9 @@ class TestResearchOrchestrator:
def test_should_stop_research_logic(self): def test_should_stop_research_logic(self):
"""Test the stopping logic based on EvaluationResult.""" """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 # Create mock evaluation results
from unittest.mock import MagicMock from unittest.mock import MagicMock
@ -148,7 +157,9 @@ class TestResearchOrchestrator:
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_conduct_research_workflow(self, test_model, mock_client): async def test_conduct_research_workflow(self, test_model, mock_client):
"""Test the basic research workflow using TestModel.""" """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 # Setup mock client returns
mock_chunks = [ mock_chunks = [

View file

@ -1,3 +1,4 @@
from haiku.rag.config import Config
from haiku.rag.research import SearchAnswer, SearchSpecialistAgent from haiku.rag.research import SearchAnswer, SearchSpecialistAgent
@ -5,7 +6,9 @@ class TestSearchSpecialistAgent:
"""Lean tests for SearchSpecialistAgent without LLM mocking.""" """Lean tests for SearchSpecialistAgent without LLM mocking."""
def test_agent_initialization(self): def test_agent_initialization(self):
agent = SearchSpecialistAgent(provider="openai", model="gpt-4") agent = SearchSpecialistAgent(
assert agent.provider == "openai" provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
assert agent.model == "gpt-4" )
assert agent.provider == Config.RESEARCH_PROVIDER
assert agent.model == Config.RESEARCH_MODEL
assert agent.output_type is SearchAnswer assert agent.output_type is SearchAnswer

View file

@ -1,3 +1,4 @@
from haiku.rag.config import Config
from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent
@ -5,7 +6,9 @@ class TestSynthesisAgent:
"""Lean tests for SynthesisAgent without LLM mocking.""" """Lean tests for SynthesisAgent without LLM mocking."""
def test_agent_initialization(self): def test_agent_initialization(self):
agent = SynthesisAgent(provider="openai", model="gpt-4") agent = SynthesisAgent(
assert agent.provider == "openai" provider=Config.RESEARCH_PROVIDER, model=Config.RESEARCH_MODEL
assert agent.model == "gpt-4" )
assert agent.provider == Config.RESEARCH_PROVIDER
assert agent.model == Config.RESEARCH_MODEL
assert agent.output_type == ResearchReport assert agent.output_type == ResearchReport