diff --git a/docs/agents.md b/docs/agents.md index 41a52c65..33579666 100644 --- a/docs/agents.md +++ b/docs/agents.md @@ -43,6 +43,8 @@ The research workflow coordinates specialized agents to plan, search, analyze, a Components: - Orchestrator: Plans, coordinates, and loops until confidence is sufficient +- Presearch Survey: Runs a quick KB scan and summarizes relevant chunk text to + ground the initial plan (plain-text summary; no URIs or scores) - Search Specialist: Performs targeted RAG searches and answers sub‑questions - Analysis & Evaluation: Extracts insights, identifies gaps, proposes new questions - Synthesis: Produces a final structured research report diff --git a/src/haiku/rag/app.py b/src/haiku/rag/app.py index 01311b86..9347d255 100644 --- a/src/haiku/rag/app.py +++ b/src/haiku/rag/app.py @@ -122,12 +122,7 @@ class HaikuRAGApp: self.console.print(f"• {finding}") self.console.print() - # Themes - if report.themes: - self.console.print("[bold cyan]Key Themes:[/bold cyan]") - for theme, explanation in report.themes.items(): - self.console.print(f"• [bold]{theme}[/bold]: {explanation}") - self.console.print() + # (Themes section removed) # Conclusions if report.conclusions: @@ -261,7 +256,7 @@ class HaikuRAGApp: elif transport == "sse": await server.run_sse_async() else: - await server.run_http_async("streamable-http") + await server.run_http_async(transport="streamable-http") except KeyboardInterrupt: pass finally: diff --git a/src/haiku/rag/qa/agent.py b/src/haiku/rag/qa/agent.py index b7edb6fa..3493b0ce 100644 --- a/src/haiku/rag/qa/agent.py +++ b/src/haiku/rag/qa/agent.py @@ -49,6 +49,9 @@ class QuestionAnswerAgent: limit: int = 3, ) -> list[SearchResult]: """Search the knowledge base for relevant documents.""" + + # Remove quotes from queries as this requires positional indexing in lancedb + query = query.replace('"', "") search_results = await ctx.deps.client.search(query, limit=limit) expanded_results = await ctx.deps.client.expand_context(search_results) diff --git a/src/haiku/rag/research/__init__.py b/src/haiku/rag/research/__init__.py index e0716ee8..e1e4f1e4 100644 --- a/src/haiku/rag/research/__init__.py +++ b/src/haiku/rag/research/__init__.py @@ -12,6 +12,7 @@ from haiku.rag.research.evaluation_agent import ( EvaluationResult, ) from haiku.rag.research.orchestrator import ResearchOrchestrator, ResearchPlan +from haiku.rag.research.presearch_agent import PresearchSurveyAgent from haiku.rag.research.search_agent import SearchSpecialistAgent from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent @@ -25,6 +26,7 @@ __all__ = [ # Specialized agents "SearchAnswer", "SearchSpecialistAgent", + "PresearchSurveyAgent", "AnalysisEvaluationAgent", "EvaluationResult", "SynthesisAgent", diff --git a/src/haiku/rag/research/base.py b/src/haiku/rag/research/base.py index e61b6ab1..1a0796c9 100644 --- a/src/haiku/rag/research/base.py +++ b/src/haiku/rag/research/base.py @@ -33,10 +33,18 @@ class BaseResearchAgent[T](ABC): # Import deps type lazily to avoid circular import during module load from haiku.rag.research.dependencies import ResearchDependencies + # If the agent is expected to return plain text, pass `str` directly. + # Otherwise, wrap the model with ToolOutput for robust tool-handling retries. + agent_output_type: Any + if self.output_type is str: # plain text output + agent_output_type = str + else: + agent_output_type = ToolOutput(self.output_type, max_retries=3) + self._agent = Agent( model=model_obj, deps_type=ResearchDependencies, - output_type=ToolOutput(self.output_type, max_retries=3), + output_type=agent_output_type, system_prompt=self.get_system_prompt(), ) diff --git a/src/haiku/rag/research/evaluation_agent.py b/src/haiku/rag/research/evaluation_agent.py index 5b924c81..8d3f5541 100644 --- a/src/haiku/rag/research/evaluation_agent.py +++ b/src/haiku/rag/research/evaluation_agent.py @@ -11,7 +11,9 @@ class EvaluationResult(BaseModel): description="Main insights extracted from the research so far" ) new_questions: list[str] = Field( - description="New sub-questions to add to the research (max 3)", max_length=3 + description="New sub-questions to add to the research (max 3)", + max_length=3, + default=[], ) confidence_score: float = Field( description="Confidence level in the completeness of research (0-1)", diff --git a/src/haiku/rag/research/orchestrator.py b/src/haiku/rag/research/orchestrator.py index ffde5736..666ed3cf 100644 --- a/src/haiku/rag/research/orchestrator.py +++ b/src/haiku/rag/research/orchestrator.py @@ -12,6 +12,7 @@ from haiku.rag.research.evaluation_agent import ( AnalysisEvaluationAgent, EvaluationResult, ) +from haiku.rag.research.presearch_agent import PresearchSurveyAgent from haiku.rag.research.prompts import ORCHESTRATOR_PROMPT from haiku.rag.research.search_agent import SearchSpecialistAgent from haiku.rag.research.synthesis_agent import ResearchReport, SynthesisAgent @@ -41,6 +42,9 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): self.search_agent: SearchSpecialistAgent = SearchSpecialistAgent( provider, model ) + self.presearch_agent: PresearchSurveyAgent = PresearchSurveyAgent( + provider, model + ) self.evaluation_agent: AnalysisEvaluationAgent = AnalysisEvaluationAgent( provider, model ) @@ -61,7 +65,12 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): "original_question": context.original_question, "unanswered_questions": context.sub_questions, "qa_responses": [ - {"question": qa.query, "answer": qa.answer} + { + "question": qa.query, + "answer": qa.answer, + "context_snippets": qa.context, + "sources": qa.sources, + } for qa in context.qa_responses ], "insights": context.insights, @@ -99,12 +108,38 @@ class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]): # Use provided console or create a new one console = console or Console() if verbose else None + # Run a simple presearch survey to summarize KB context + if console: + console.print( + "\n[bold cyan]🔎 Presearch: summarizing KB context...[/bold cyan]" + ) + + presearch_result = await self.presearch_agent.run(question, deps=deps) + # 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}" + ) + plan_result: AgentRunResult[ResearchPlan] = await self.run( - f"Create a research plan for: {question}", deps=deps + plan_prompt, deps=deps ) context.sub_questions = plan_result.output.sub_questions diff --git a/src/haiku/rag/research/presearch_agent.py b/src/haiku/rag/research/presearch_agent.py new file mode 100644 index 00000000..8c6e7c33 --- /dev/null +++ b/src/haiku/rag/research/presearch_agent.py @@ -0,0 +1,34 @@ +from pydantic_ai import RunContext +from pydantic_ai.run import AgentRunResult + +from haiku.rag.research.base import BaseResearchAgent +from haiku.rag.research.dependencies import ResearchDependencies +from haiku.rag.research.prompts import PRESEARCH_AGENT_PROMPT + + +class PresearchSurveyAgent(BaseResearchAgent[str]): + """Presearch agent that gathers verbatim context and summarizes it.""" + + def __init__(self, provider: str, model: str) -> None: + super().__init__(provider, model, str) + + async def run( + self, prompt: str, deps: ResearchDependencies, **kwargs + ) -> AgentRunResult[str]: + return await super().run(prompt, deps, **kwargs) + + def get_system_prompt(self) -> str: + return PRESEARCH_AGENT_PROMPT + + def register_tools(self) -> None: + @self.agent.tool + async def gather_context( + ctx: RunContext[ResearchDependencies], + query: 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/prompts.py b/src/haiku/rag/research/prompts.py index 4bcec6a6..af5038f7 100644 --- a/src/haiku/rag/research/prompts.py +++ b/src/haiku/rag/research/prompts.py @@ -114,3 +114,16 @@ Focus on creating a report that provides clear value to the reader by: - Highlighting the most important findings - Explaining the implications of the research - Suggesting concrete next steps""" + +PRESEARCH_AGENT_PROMPT = """You are a rapid research surveyor. + +Task: +- Call the gather_context tool once with the main question to obtain a + relevant texts from the Knowledge Base (KB). +- Read that context and produce a brief natural-language summary describing + what the KB appears to contain relative to the question. + +Rules: +- Base the summary strictly on the provided text; do not invent. +- Output only the summary as plain text (one short paragraph). +""" diff --git a/src/haiku/rag/research/search_agent.py b/src/haiku/rag/research/search_agent.py index 200f619a..b2cdc17f 100644 --- a/src/haiku/rag/research/search_agent.py +++ b/src/haiku/rag/research/search_agent.py @@ -42,6 +42,7 @@ class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]): ) -> 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 45349789..e0dbd0a3 100644 --- a/src/haiku/rag/research/synthesis_agent.py +++ b/src/haiku/rag/research/synthesis_agent.py @@ -12,11 +12,12 @@ class ResearchReport(BaseModel): main_findings: list[str] = Field( description="Primary research findings with supporting evidence" ) - themes: dict[str, str] = Field(description="Major themes and their explanations") conclusions: list[str] = Field(description="Evidence-based conclusions") - limitations: list[str] = Field(description="Limitations of the current research") + limitations: list[str] = Field( + description="Limitations of the current research", default=[] + ) recommendations: list[str] = Field( - description="Actionable recommendations based on findings" + description="Actionable recommendations based on findings", default=[] ) sources_summary: str = Field( description="Summary of sources used and their reliability" diff --git a/tests/research/test_orchestrator.py b/tests/research/test_orchestrator.py index 09d44c6e..3c298c82 100644 --- a/tests/research/test_orchestrator.py +++ b/tests/research/test_orchestrator.py @@ -172,7 +172,6 @@ class TestResearchOrchestrator: assert report.title assert report.executive_summary assert isinstance(report.main_findings, list) - assert isinstance(report.themes, dict) assert isinstance(report.conclusions, list) assert isinstance(report.limitations, list) assert isinstance(report.recommendations, list)