From 5bfa720864842013b14bbea6626bae0abde46a3c Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 19 Sep 2025 09:05:04 +0300 Subject: [PATCH 1/7] Remove openai models from research tests --- tests/research/test_evaluation_agent.py | 9 ++++--- tests/research/test_orchestrator.py | 33 ++++++++++++++++--------- tests/research/test_search_agent.py | 9 ++++--- tests/research/test_synthesis_agent.py | 9 ++++--- 4 files changed, 40 insertions(+), 20 deletions(-) diff --git a/tests/research/test_evaluation_agent.py b/tests/research/test_evaluation_agent.py index fd63c5a4..0c3a50d6 100644 --- a/tests/research/test_evaluation_agent.py +++ b/tests/research/test_evaluation_agent.py @@ -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 diff --git a/tests/research/test_orchestrator.py b/tests/research/test_orchestrator.py index 3c298c82..304eb788 100644 --- a/tests/research/test_orchestrator.py +++ b/tests/research/test_orchestrator.py @@ -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 = [ diff --git a/tests/research/test_search_agent.py b/tests/research/test_search_agent.py index a7a9062c..71196f27 100644 --- a/tests/research/test_search_agent.py +++ b/tests/research/test_search_agent.py @@ -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 diff --git a/tests/research/test_synthesis_agent.py b/tests/research/test_synthesis_agent.py index 666662c7..65146782 100644 --- a/tests/research/test_synthesis_agent.py +++ b/tests/research/test_synthesis_agent.py @@ -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 From 3abe669681df38fe1e43c2d8c0fad8165318d983 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 19 Sep 2025 09:13:27 +0300 Subject: [PATCH 2/7] Simplify presearch prompt --- src/haiku/rag/research/orchestrator.py | 19 +++---------------- 1 file changed, 3 insertions(+), 16 deletions(-) diff --git a/src/haiku/rag/research/orchestrator.py b/src/haiku/rag/research/orchestrator.py index 666ed3cf..1274f384 100644 --- a/src/haiku/rag/research/orchestrator.py +++ b/src/haiku/rag/research/orchestrator.py @@ -108,40 +108,27 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): # Use provided console or create a new one console = console or Console() if verbose else None + # Create initial research plan # Run a simple presearch survey to summarize KB context if console: + console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]") 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 - 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", - ) - 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: From ad7a9ed3249f62e7389541fb62997ee230d3d7cb Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 19 Sep 2025 09:17:14 +0300 Subject: [PATCH 3/7] Provide register_tools to the base research agent --- src/haiku/rag/research/base.py | 1 - src/haiku/rag/research/evaluation_agent.py | 4 ---- src/haiku/rag/research/orchestrator.py | 5 ----- src/haiku/rag/research/synthesis_agent.py | 6 ------ 4 files changed, 16 deletions(-) diff --git a/src/haiku/rag/research/base.py b/src/haiku/rag/research/base.py index 1a0796c9..6f612b7a 100644 --- a/src/haiku/rag/research/base.py +++ b/src/haiku/rag/research/base.py @@ -75,7 +75,6 @@ class BaseResearchAgent[T](ABC): """Return the system prompt for this agent.""" pass - @abstractmethod def register_tools(self) -> None: """Register agent-specific tools.""" pass diff --git a/src/haiku/rag/research/evaluation_agent.py b/src/haiku/rag/research/evaluation_agent.py index 8d3f5541..519f36f7 100644 --- a/src/haiku/rag/research/evaluation_agent.py +++ b/src/haiku/rag/research/evaluation_agent.py @@ -36,7 +36,3 @@ class AnalysisEvaluationAgent(BaseResearchAgent[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 diff --git a/src/haiku/rag/research/orchestrator.py b/src/haiku/rag/research/orchestrator.py index 1274f384..5fe788c3 100644 --- a/src/haiku/rag/research/orchestrator.py +++ b/src/haiku/rag/research/orchestrator.py @@ -53,11 +53,6 @@ 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 _format_context_for_prompt(self, context: ResearchContext) -> str: """Format the research context as XML for inclusion in prompts.""" diff --git a/src/haiku/rag/research/synthesis_agent.py b/src/haiku/rag/research/synthesis_agent.py index e0dbd0a3..e99ee719 100644 --- a/src/haiku/rag/research/synthesis_agent.py +++ b/src/haiku/rag/research/synthesis_agent.py @@ -32,9 +32,3 @@ class SynthesisAgent(BaseResearchAgent[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 From 5f102c8475daccb40d66afc1a2e658e7ca30a66f Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 19 Sep 2025 09:43:57 +0300 Subject: [PATCH 4/7] Set research agents to retry 3 times --- src/haiku/rag/research/base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/haiku/rag/research/base.py b/src/haiku/rag/research/base.py index 6f612b7a..2f05f94c 100644 --- a/src/haiku/rag/research/base.py +++ b/src/haiku/rag/research/base.py @@ -46,6 +46,7 @@ class BaseResearchAgent[T](ABC): deps_type=ResearchDependencies, output_type=agent_output_type, system_prompt=self.get_system_prompt(), + retries=3, ) # Register tools From ccaae7cca1f0fad153eb9c202b63a57160354293 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 19 Sep 2025 09:44:46 +0300 Subject: [PATCH 5/7] Simplify research verbose logging --- docs/agents.md | 7 +- src/haiku/rag/app.py | 4 +- src/haiku/rag/research/orchestrator.py | 153 +++++++++++-------------- 3 files changed, 74 insertions(+), 90 deletions(-) diff --git a/docs/agents.md b/docs/agents.md index 33579666..991436ce 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -70,14 +70,17 @@ 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", + verbose=True +) 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) diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 9347d255..e2f2f898 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -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() + orchestrator = ResearchOrchestrator(verbose=verbose) if verbose: self.console.print( @@ -100,8 +100,6 @@ class HaikuRAGApp: question=question, client=client, max_iterations=max_iterations, - verbose=verbose, - console=self.console if verbose else None, ) # Display the report diff --git a/src/haiku/rag/research/orchestrator.py b/src/haiku/rag/research/orchestrator.py index 5fe788c3..2de52f60 100644 --- a/src/haiku/rag/research/orchestrator.py +++ b/src/haiku/rag/research/orchestrator.py @@ -31,7 +31,10 @@ 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, + verbose: bool = False, ): # Use provided values or fall back to config defaults provider = provider or Config.RESEARCH_PROVIDER or Config.QA_PROVIDER @@ -49,6 +52,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 @@ -73,14 +80,30 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): } return format_as_xml(context_data, root_tag="research_context") + def _should_stop_research( + self, + evaluation_result: AgentRunResult[EvaluationResult], + confidence_threshold: float, + ) -> bool: + """Determine if research should stop based on evaluation.""" + + result = evaluation_result.output + return result.is_sufficient and result.confidence_score >= confidence_threshold + + 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, - console: Console | None = None, ) -> ResearchReport: """Conduct comprehensive research on a question. @@ -100,16 +123,10 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): 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 # Run a simple presearch survey to summarize KB context - if console: - console.print("\n[bold cyan]📋 Creating research plan...[/bold cyan]") - console.print( - "\n[bold cyan]🔎 Presearch: summarizing KB context...[/bold cyan]" - ) + 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) @@ -126,42 +143,36 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): ) 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() + 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}") # Execute research iterations for iteration in range(max_iterations): - if console: - console.rule( - f"[bold yellow]🔄 Iteration {iteration + 1}/{max_iterations}[/bold yellow]" - ) + self._log( + f"[bold yellow]🔄 Iteration {iteration + 1}/{max_iterations}[/bold yellow]", + rule=True, + ) # 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]" - ) + self._log( + "[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}") + 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}") # Run searches for all questions and remove answered ones answered_questions = [] @@ -169,27 +180,22 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): 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})" - ) + self._log( + f"\n [red]×[/red] Omitting failed question: {search_question} ({e})" + ) finally: answered_questions.append(search_question) - if console and context.qa_responses: + 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] + "..." + 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}") + self._log(f"\n [green]✓[/green] {search_question}") + self._log(f" {answer_preview}") # Remove answered questions from the list for question in answered_questions: @@ -197,10 +203,9 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): context.sub_questions.remove(question) # Analysis and Evaluation phase - if console: - console.print( - "\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]" - ) + 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) @@ -215,19 +220,14 @@ Evaluate the research progress for the original question and identify any remain 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}") + 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: @@ -240,17 +240,13 @@ Evaluate the research progress for the original question and identify any remain # 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}" - ) + self._log( + 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]" - ) + 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) @@ -264,19 +260,6 @@ Create a detailed report that synthesizes all findings into a coherent response. synthesis_prompt, deps=deps ) - if console: - console.print("[bold green]✅ Research complete![/bold green]") + self._log("[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 From 9c128c1d360b4fd4c0a8e42a46396df9b15bc35d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 19 Sep 2025 10:29:59 +0300 Subject: [PATCH 6/7] Make answered question handling the responsibility of the search agent --- src/haiku/rag/research/orchestrator.py | 18 ++---------------- src/haiku/rag/research/search_agent.py | 6 ++---- 2 files changed, 4 insertions(+), 20 deletions(-) diff --git a/src/haiku/rag/research/orchestrator.py b/src/haiku/rag/research/orchestrator.py index 2de52f60..1481579c 100644 --- a/src/haiku/rag/research/orchestrator.py +++ b/src/haiku/rag/research/orchestrator.py @@ -165,7 +165,7 @@ 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 self._log( @@ -175,17 +175,8 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): self._log(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 - self._log( - f"\n [red]×[/red] Omitting failed question: {search_question} ({e})" - ) - finally: - answered_questions.append(search_question) - + 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] @@ -197,11 +188,6 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): self._log(f"\n [green]✓[/green] {search_question}") self._log(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 self._log( "\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]" diff --git a/src/haiku/rag/research/search_agent.py b/src/haiku/rag/research/search_agent.py index b2cdc17f..8eb465fe 100644 --- a/src/haiku/rag/research/search_agent.py +++ b/src/haiku/rag/research/search_agent.py @@ -22,10 +22,8 @@ class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]): 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) - + deps.context.add_qa_response(result.output) + deps.context.sub_questions.remove(prompt) return result def get_system_prompt(self) -> str: From a215a64686ce545c3dd4287c7089eb9ba28a0b97 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 19 Sep 2025 11:18:13 +0300 Subject: [PATCH 7/7] Refactor to give its research subagent full responsibility --- docs/agents.md | 7 +- src/haiku/rag/app.py | 3 +- src/haiku/rag/research/base.py | 2 +- src/haiku/rag/research/dependencies.py | 24 ++++ src/haiku/rag/research/evaluation_agent.py | 47 +++++++ src/haiku/rag/research/orchestrator.py | 155 +++++---------------- src/haiku/rag/research/presearch_agent.py | 7 +- src/haiku/rag/research/search_agent.py | 12 +- src/haiku/rag/research/synthesis_agent.py | 26 ++++ tests/generate_benchmark_db.py | 2 +- 10 files changed, 155 insertions(+), 130 deletions(-) diff --git a/docs/agents.md b/docs/agents.md index 991436ce..502bc3e5 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -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) diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index e2f2f898..704b9205 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -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 diff --git a/src/haiku/rag/research/base.py b/src/haiku/rag/research/base.py index 2f05f94c..a2ce6ad8 100644 --- a/src/haiku/rag/research/base.py +++ b/src/haiku/rag/research/base.py @@ -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, ) diff --git a/src/haiku/rag/research/dependencies.py b/src/haiku/rag/research/dependencies.py index 3438a796..ad1ed935 100644 --- a/src/haiku/rag/research/dependencies.py +++ b/src/haiku/rag/research/dependencies.py @@ -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") diff --git a/src/haiku/rag/research/evaluation_agent.py b/src/haiku/rag/research/evaluation_agent.py index 519f36f7..0867349a 100644 --- a/src/haiku/rag/research/evaluation_agent.py +++ b/src/haiku/rag/research/evaluation_agent.py @@ -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 diff --git a/src/haiku/rag/research/orchestrator.py b/src/haiku/rag/research/orchestrator.py index 1481579c..af8ed766 100644 --- a/src/haiku/rag/research/orchestrator.py +++ b/src/haiku/rag/research/orchestrator.py @@ -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 diff --git a/src/haiku/rag/research/presearch_agent.py b/src/haiku/rag/research/presearch_agent.py index 8c6e7c33..5482dfc8 100644 --- a/src/haiku/rag/research/presearch_agent.py +++ b/src/haiku/rag/research/presearch_agent.py @@ -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) diff --git a/src/haiku/rag/research/search_agent.py b/src/haiku/rag/research/search_agent.py index 8eb465fe..321b6fee 100644 --- a/src/haiku/rag/research/search_agent.py +++ b/src/haiku/rag/research/search_agent.py @@ -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) diff --git a/src/haiku/rag/research/synthesis_agent.py b/src/haiku/rag/research/synthesis_agent.py index e99ee719..e3499e95 100644 --- a/src/haiku/rag/research/synthesis_agent.py +++ b/src/haiku/rag/research/synthesis_agent.py @@ -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 diff --git a/tests/generate_benchmark_db.py b/tests/generate_benchmark_db.py index 3dd1e3e1..756884ef 100644 --- a/tests/generate_benchmark_db.py +++ b/tests/generate_benchmark_db.py @@ -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()