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