haiku.rag 0.9.2__py3-none-any.whl → 0.10.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/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-0.9.2.dist-info/RECORD

@@ -1,50 +0,0 @@
-haiku/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-haiku/rag/app.py,sha256=o64L7aj5V8lYHxWhAKBNj1tGfXiN6xr0_Cc1dEYd3As,11483
-haiku/rag/chunker.py,sha256=PVe6ysv8UlacUd4Zb3_8RFWIaWDXnzBAy2VDJ4TaUsE,1555
-haiku/rag/cli.py,sha256=3nlzrT5FPCyfnu51KHchLG4Cj2eVv9YsuGHMShBnVb0,9845
-haiku/rag/client.py,sha256=NJVGXzVzpoVy1sttz_xEU7mXWtObKT8pGpvo5pZyzwc,21288
-haiku/rag/config.py,sha256=SPEIv2IElZmZh4Wsp8gk7ViRW5ZzD-UGmIqRAXscDdI,2134
-haiku/rag/logging.py,sha256=dm65AwADpcQsH5OAPtRA-4hsw0w5DK-sGOvzYkj6jzw,1720
-haiku/rag/mcp.py,sha256=bR9Y-Nz-hvjiql20Y0KE0hwNGwyjmPGX8K9d-qmXptY,4683
-haiku/rag/migration.py,sha256=M--KnSF3lxgKjxmokb4vuzGH-pV8eg0C_8e7jvPqW8Y,11058
-haiku/rag/monitor.py,sha256=r386nkhdlsU8UECwIuVwnrSlgMk3vNIuUZGNIzkZuec,2770
-haiku/rag/reader.py,sha256=qkPTMJuQ_o4sK-8zpDl9WFYe_MJ7aL_gUw6rczIpW-g,3274
-haiku/rag/utils.py,sha256=aiuPu_rrfpyIvJJq0o5boUIIvCdNzdpKwAIPYYn3iG8,4965
-haiku/rag/embeddings/__init__.py,sha256=44IfDITGIFTflGT6UEmiYOwpWFVbYv5smLY59D0YeCs,1419
-haiku/rag/embeddings/base.py,sha256=BnSviKrlzjv3L0sZJs_T-pxfawd-bcTak-rsX-D2f3A,497
-haiku/rag/embeddings/ollama.py,sha256=LuLlHH6RGoO9_gFCIlbmesuXOj017gTw6z-p8Ez0CfE,595
-haiku/rag/embeddings/openai.py,sha256=fIFCk-jpUtaW0xsnrQnJ824O0UCjaGG2sgvBzREhilc,503
-haiku/rag/embeddings/vllm.py,sha256=vhaUnCn6VMkfSluLhWKtSV-sekFaPsp4pKo2N7-SBCY,626
-haiku/rag/embeddings/voyageai.py,sha256=UW-MW4tJKnPB6Fs2P7A3yt-ZeRm46H9npckchSriPX8,661
-haiku/rag/qa/__init__.py,sha256=Sl7Kzrg9CuBOcMF01wc1NtQhUNWjJI0MhIHfCWrb8V4,434
-haiku/rag/qa/agent.py,sha256=f4Keh-ESgctNbTg96QL95HYjINVLOcxa8t8crx92MMk,3081
-haiku/rag/qa/prompts.py,sha256=LhRfDtO8Pb06lpr4PpwEaKUYItZ5OiIkeqcCogcssHY,3347
-haiku/rag/reranking/__init__.py,sha256=IRXHs4qPu6VbGJQpzSwhgtVWWumURH_vEoVFE-extlo,894
-haiku/rag/reranking/base.py,sha256=LM9yUSSJ414UgBZhFTgxGprlRqzfTe4I1vgjricz2JY,405
-haiku/rag/reranking/cohere.py,sha256=1iTdiaa8vvb6oHVB2qpWzUOVkyfUcimVSZp6Qr4aq4c,1049
-haiku/rag/reranking/mxbai.py,sha256=46sVTsTIkzIX9THgM3u8HaEmgY7evvEyB-N54JTHvK8,867
-haiku/rag/reranking/vllm.py,sha256=xVGH9ss-ISWdJ5SKUUHUbTqBo7PIEmA_SQv0ScdJ6XA,1479
-haiku/rag/research/__init__.py,sha256=qLF41YayAxW_VeHhuTceVuz9hw1FNbuRV9VMhonUMW0,1078
-haiku/rag/research/base.py,sha256=ZUvqh-IxU8r4mOPRKjwZcHciKcIfrTnP6Q_9jVElelQ,4041
-haiku/rag/research/dependencies.py,sha256=vZctKC5donqhm8LFO6hQdXZZXzjdW1__4eUlhyZn058,1573
-haiku/rag/research/evaluation_agent.py,sha256=yyBobKr8MRwiox59I2Jqycp02ju9EGVaI9FceRGL188,1386
-haiku/rag/research/orchestrator.py,sha256=LrxRG74BWun0T4uOxhc9AuitxbPioS_eG_nE098ftyY,11765
-haiku/rag/research/presearch_agent.py,sha256=vf-WlY46g5tuuLKMBuPXXYYffynsBw7KVLr8LoTNHnU,1292
-haiku/rag/research/prompts.py,sha256=pVRB7_b_p3JaLF1bC3ANTbSFY78ypSjDhoq6peoU6jo,5685
-haiku/rag/research/search_agent.py,sha256=0iK7vCd9w7h8pWJgB6VUSPOdjlzB8peboNSXxuEGBK0,2464
-haiku/rag/research/synthesis_agent.py,sha256=jo5rg7aL4zGXLQP105cANqRPIiwJLqYe2unO5BQkNvE,1511
-haiku/rag/store/__init__.py,sha256=hq0W0DAC7ysqhWSP2M2uHX8cbG6kbr-sWHxhq6qQcY0,103
-haiku/rag/store/engine.py,sha256=fNrykqMX7PRSCt4LSRfuJ66OLrb8BVYq2bpbfI2iaWU,8455
-haiku/rag/store/models/__init__.py,sha256=s0E72zneGlowvZrFWaNxHYjOAUjgWdLxzdYsnvNRVlY,88
-haiku/rag/store/models/chunk.py,sha256=ZNyTfO6lh3rXWLVYO3TZcitbL4LSUGr42fR6jQQ5iQc,364
-haiku/rag/store/models/document.py,sha256=zSSpt6pyrMJAIXGQvIcqojcqUzwZnhp3WxVokaWxNRc,396
-haiku/rag/store/repositories/__init__.py,sha256=Olv5dLfBQINRV3HrsfUpjzkZ7Qm7goEYyMNykgo_DaY,291
-haiku/rag/store/repositories/chunk.py,sha256=1RmPyEYRYOFbrALbmLOo62t3f-xO2KgxUjcvPdrRZlc,14467
-haiku/rag/store/repositories/document.py,sha256=XoLCrMrZqs0iCZoHlDOfRDaVUux77Vdu5iZczduF1rY,7812
-haiku/rag/store/repositories/settings.py,sha256=wx3fuP_5CpPflZHRrIkeoer6ml-iD0qXERh5k6MQRzI,5291
-haiku/rag/store/upgrades/__init__.py,sha256=wUiEoSiHTahvuagx93E4FB07v123AhdbOjwUkPusiIg,14
-haiku_rag-0.9.2.dist-info/METADATA,sha256=IM9tGvye83CRTj2wOFtPP7oD9KtJvp3RXh4QdCFknD4,4681
-haiku_rag-0.9.2.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-haiku_rag-0.9.2.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
-haiku_rag-0.9.2.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
-haiku_rag-0.9.2.dist-info/RECORD,,