haiku.rag 0.9.3__py3-none-any.whl → 0.10.0__py3-none-any.whl

haiku/rag/research/orchestrator.py

@@ -1,170 +0,0 @@
-from typing import Any
-
-from pydantic import BaseModel, Field
-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.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
-
-
-class ResearchPlan(BaseModel):
-    """Research execution plan."""
-
-    main_question: str = Field(description="The main research question")
-    sub_questions: list[str] = Field(
-        description="Decomposed sub-questions to investigate (max 3)", max_length=3
-    )
-
-
-class ResearchOrchestrator(BaseResearchAgent[ResearchPlan]):
-    """Orchestrator agent that coordinates the research workflow."""
-
-    def __init__(
-        self,
-        provider: str | None = Config.RESEARCH_PROVIDER,
-        model: str | None = None,
-    ):
-        # Use provided values or fall back to config defaults
-        provider = provider or Config.RESEARCH_PROVIDER or Config.QA_PROVIDER
-        model = model or Config.RESEARCH_MODEL or Config.QA_MODEL
-
-        super().__init__(provider, model, output_type=ResearchPlan)
-
-        self.search_agent: SearchSpecialistAgent = SearchSpecialistAgent(
-            provider, model
-        )
-        self.presearch_agent: PresearchSurveyAgent = PresearchSurveyAgent(
-            provider, model
-        )
-        self.evaluation_agent: AnalysisEvaluationAgent = AnalysisEvaluationAgent(
-            provider, model
-        )
-        self.synthesis_agent: SynthesisAgent = SynthesisAgent(provider, model)
-
-    def get_system_prompt(self) -> str:
-        return ORCHESTRATOR_PROMPT
-
-    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
-
-    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.
-
-        Args:
-            question: The research question to investigate
-            client: HaikuRAG client for document operations
-            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
-
-        Returns:
-            ResearchReport with comprehensive findings
-        """
-
-        # 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
-        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"
-            "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"{presearch_result.output}"
-        )
-
-        plan_result: AgentRunResult[ResearchPlan] = await self.run(
-            plan_prompt, deps=deps
-        )
-        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}")
-
-        # Execute research iterations
-        for iteration in range(max_iterations):
-            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:
-                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
-            if console:
-                console.print(
-                    f"\n[bold cyan]🔍 Searching & Answering {len(questions_to_search)} questions:[/bold cyan]"
-                )
-
-            for search_question in questions_to_search:
-                await self.search_agent.run(search_question, deps=deps)
-
-            # Analysis and Evaluation phase
-
-            evaluation_result = await self.evaluation_agent.run("", deps=deps)
-
-            # 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}"
-                    )
-                break
-
-        # Generate final report
-        report_result: AgentRunResult[ResearchReport] = await self.synthesis_agent.run(
-            "", deps=deps
-        )
-
-        return report_result.output

haiku/rag/research/presearch_agent.py

@@ -1,39 +0,0 @@
-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]:
-        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:
-        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."""
-            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)

haiku/rag/research/search_agent.py

@@ -1,69 +0,0 @@
-from pydantic_ai import RunContext
-from pydantic_ai.format_prompt import format_as_xml
-from pydantic_ai.run import AgentRunResult
-
-from haiku.rag.research.base import BaseResearchAgent, SearchAnswer
-from haiku.rag.research.dependencies import ResearchDependencies
-from haiku.rag.research.prompts import SEARCH_AGENT_PROMPT
-
-
-class SearchSpecialistAgent(BaseResearchAgent[SearchAnswer]):
-    """Agent specialized in answering questions using RAG search."""
-
-    def __init__(self, provider: str, model: str) -> None:
-        super().__init__(provider, model, output_type=SearchAnswer)
-
-    async def run(
-        self, prompt: str, deps: ResearchDependencies, **kwargs
-    ) -> AgentRunResult[SearchAnswer]:
-        """Execute the agent and persist the QA pair in shared context.
-
-        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:
-        return SEARCH_AGENT_PROMPT
-
-    def register_tools(self) -> None:
-        """Register search-specific tools."""
-
-        @self.agent.tool
-        async def search_and_answer(
-            ctx: RunContext[ResearchDependencies],
-            query: str,
-            limit: int = 5,
-        ) -> str:
-            """Search the KB and return a concise context pack."""
-            search_results = await ctx.deps.client.search(query, limit=limit)
-            expanded = await ctx.deps.client.expand_context(search_results)
-
-            snippet_entries = [
-                {
-                    "text": chunk.content,
-                    "score": score,
-                    "document_uri": (chunk.document_uri or ""),
-                }
-                for chunk, score in expanded
-            ]
-
-            # Return an XML-formatted payload with the question and snippets.
-            if snippet_entries:
-                return format_as_xml(snippet_entries, root_tag="snippets")
-            else:
-                return (
-                    f"No relevant information found in the knowledge base for: {query}"
-                )

haiku/rag/research/synthesis_agent.py

@@ -1,60 +0,0 @@
-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
-
-
-class ResearchReport(BaseModel):
-    """Final research report structure."""
-
-    title: str = Field(description="Concise title for the research")
-    executive_summary: str = Field(description="Brief overview of key findings")
-    main_findings: list[str] = Field(
-        description="Primary research findings with supporting evidence"
-    )
-    conclusions: list[str] = Field(description="Evidence-based conclusions")
-    limitations: list[str] = Field(
-        description="Limitations of the current research", default=[]
-    )
-    recommendations: list[str] = Field(
-        description="Actionable recommendations based on findings", default=[]
-    )
-    sources_summary: str = Field(
-        description="Summary of sources used and their reliability"
-    )
-
-
-class SynthesisAgent(BaseResearchAgent[ResearchReport]):
-    """Agent specialized in synthesizing research into comprehensive reports."""
-
-    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