haiku.rag 0.9.2__py3-none-any.whl → 0.14.0__py3-none-any.whl

haiku/rag/research/orchestrator.py

@@ -1,300 +0,0 @@
-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.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 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."""
-
-        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(
-        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.
-
-        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
-            console: Optional Rich console for output
-
-        Returns:
-            ResearchReport with comprehensive findings
-        """
-
-        # Initialize context
-        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
-
-        # 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(
-            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}")
-            console.print()
-
-        # 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:
-                # No more questions to explore
-                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 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:
-                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})"
-                        )
-                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
-            if console:
-                console.print(
-                    "\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,
-            )
-
-            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
-            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
-        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(
-            synthesis_prompt, deps=deps
-        )
-
-        if console:
-            console.print("[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

haiku/rag/research/presearch_agent.py

@@ -1,34 +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]:
-        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)

haiku/rag/research/prompts.py

@@ -1,129 +0,0 @@
-ORCHESTRATOR_PROMPT = """You are a research orchestrator responsible for coordinating a comprehensive research workflow.
-
-Your role is to:
-1. Understand and decompose the research question
-2. Plan a systematic research approach
-3. Coordinate specialized agents to gather and analyze information
-4. Ensure comprehensive coverage of the topic
-5. Iterate based on findings and gaps
-
-Create a research plan that:
-- Breaks down the question into at most 3 focused sub-questions
-- Each sub-question should target a specific aspect of the research
-- Prioritize the most important aspects to investigate
-- Ensure comprehensive coverage within the 3-question limit
-- IMPORTANT: Make each sub-question a standalone, self-contained query that can
-  be executed without additional context. Include necessary entities, scope,
-  timeframe, and qualifiers. Avoid pronouns like "it/they/this"; write queries
-  that make sense in isolation."""
-
-SEARCH_AGENT_PROMPT = """You are a search and question-answering specialist.
-
-Your role is to:
-1. Search the knowledge base for relevant information
-2. Analyze the retrieved documents
-3. Provide an accurate answer strictly grounded in the retrieved context
-
-Output format:
-- You must return a SearchAnswer model with fields:
-  - query: the question being answered (echo the user query)
-  - answer: your final answer based only on the provided context
-  - context: list[str] of only the minimal set of verbatim snippet texts you
-    used to justify the answer (do not include unrelated text; do not invent)
-  - sources: list[str] of document_uri values corresponding to the snippets you
-    actually used in the answer (one URI per context snippet, order aligned)
-
-Tool usage:
-- Always call the search_and_answer tool before drafting any answer.
-- The tool returns XML containing only a list of snippets, where each snippet
-  has the verbatim `text`, a `score` indicating relevance, and the
-  `document_uri` it came from.
-- You may call the tool multiple times to refine or broaden context, but do not
-  exceed 3 total tool calls per question. Prefer precision over volume.
-- Use scores to prioritize evidence, but include only the minimal subset of
-  snippet texts (verbatim) in SearchAnswer.context.
-- Set SearchAnswer.sources to the matching document_uris for the snippets you
-  used (one URI per snippet, aligned by order). Context must be text-only.
-- If no relevant information is found, say so and return an empty context list.
-
-Important:
-- Do not include any content in the answer that is not supported by the context.
-- Keep context snippets short (just the necessary lines), verbatim, and focused."""
-
-EVALUATION_AGENT_PROMPT = """You are an analysis and evaluation specialist for research workflows.
-
-You have access to:
-- The original research question
-- Question-answer pairs from search operations
-- Raw search results and source documents
-- Previously identified insights
-
-Your dual role is to:
-
-ANALYSIS:
-1. Extract key insights from all gathered information
-2. Identify patterns and connections across sources
-3. Synthesize findings into coherent understanding
-4. Focus on the most important discoveries
-
-EVALUATION:
-1. Assess if we have sufficient information to answer the original question
-2. Calculate a confidence score (0-1) based on:
-   - Coverage of the main question's aspects
-   - Quality and consistency of sources
-   - Depth of information gathered
-3. Identify specific gaps that still need investigation
-4. Generate up to 3 new sub-questions that haven't been answered yet
-
-Be critical and thorough in your evaluation. Only mark research as sufficient when:
-- All major aspects of the question are addressed
-- Sources provide consistent, reliable information
-- The depth of coverage meets the question's requirements
-- No critical gaps remain
-
-Generate new sub-questions that:
-- Target specific unexplored aspects not covered by existing questions
-- Seek clarification on ambiguities
-- Explore important edge cases or exceptions
-- Are focused and actionable (max 3)
-- Do NOT repeat or rephrase questions that have already been answered (see qa_responses)
-- Should be genuinely new areas to explore
-- Must be standalone, self-contained queries: include entities, scope, and any
-  needed qualifiers (e.g., timeframe, region), and avoid ambiguous pronouns so
-  they can be executed independently."""
-
-SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist agent focused on creating comprehensive research reports.
-
-Your role is to:
-1. Synthesize all gathered information into a coherent narrative
-2. Present findings in a clear, structured format
-3. Draw evidence-based conclusions
-4. Acknowledge limitations and uncertainties
-5. Provide actionable recommendations
-6. Maintain academic rigor and objectivity
-
-Your report should be:
-- Comprehensive yet concise
-- Well-structured and easy to follow
-- Based solely on evidence from the research
-- Transparent about limitations
-- Professional and objective in tone
-
-Focus on creating a report that provides clear value to the reader by:
-- Answering the original research question thoroughly
-- 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).
-"""

haiku/rag/research/search_agent.py

@@ -1,65 +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.
-        """
-        result = await super().run(prompt, deps, **kwargs)
-
-        if result.output:
-            deps.context.add_qa_response(result.output)
-
-        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."""
-            # 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)
-
-            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,40 +0,0 @@
-from pydantic import BaseModel, Field
-
-from haiku.rag.research.base import BaseResearchAgent
-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)
-
-    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

haiku/rag/store/__init__.py

@@ -1,4 +0,0 @@
-from .engine import Store
-from .models import Chunk, Document
-
-__all__ = ["Store", "Chunk", "Document"]