haiku.rag 0.10.2__py3-none-any.whl → 0.19.3__py3-none-any.whl

haiku/rag/research/models.py

@@ -1,70 +0,0 @@
-from pydantic import BaseModel, Field
-
-
-class ResearchPlan(BaseModel):
-    main_question: str
-    sub_questions: list[str]
-
-
-class SearchAnswer(BaseModel):
-    """Structured output for the SearchSpecialist agent."""
-
-    query: str = Field(description="The search query that was performed")
-    answer: str = Field(description="The answer generated based on the context")
-    context: list[str] = Field(
-        description=(
-            "Only the minimal set of relevant snippets (verbatim) that directly "
-            "support the answer"
-        )
-    )
-    sources: list[str] = Field(
-        description=(
-            "Document titles (if available) or URIs corresponding to the"
-            " snippets actually used in the answer (one per snippet; omit if none)"
-        ),
-        default_factory=list,
-    )
-
-
-class EvaluationResult(BaseModel):
-    """Result of analysis and evaluation."""
-
-    key_insights: list[str] = Field(
-        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,
-        default=[],
-    )
-    confidence_score: float = Field(
-        description="Confidence level in the completeness of research (0-1)",
-        ge=0.0,
-        le=1.0,
-    )
-    is_sufficient: bool = Field(
-        description="Whether the research is sufficient to answer the original question"
-    )
-    reasoning: str = Field(
-        description="Explanation of why the research is or isn't complete"
-    )
-
-
-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"
-    )

haiku/rag/research/nodes/evaluate.py

@@ -1,80 +0,0 @@
-from dataclasses import dataclass
-
-from pydantic_ai import Agent
-from pydantic_graph import BaseNode, GraphRunContext
-
-from haiku.rag.research.common import format_context_for_prompt, get_model, log
-from haiku.rag.research.dependencies import (
-    ResearchDependencies,
-)
-from haiku.rag.research.models import EvaluationResult, ResearchReport
-from haiku.rag.research.nodes.synthesize import SynthesizeNode
-from haiku.rag.research.prompts import EVALUATION_AGENT_PROMPT
-from haiku.rag.research.state import ResearchDeps, ResearchState
-
-
-@dataclass
-class EvaluateNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
-    provider: str
-    model: str
-
-    async def run(
-        self, ctx: GraphRunContext[ResearchState, ResearchDeps]
-    ) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
-        state = ctx.state
-        deps = ctx.deps
-
-        log(
-            deps.console,
-            "\n[bold cyan]📊 Analyzing and evaluating research progress...[/bold cyan]",
-        )
-
-        agent = Agent(
-            model=get_model(self.provider, self.model),
-            output_type=EvaluationResult,
-            instructions=EVALUATION_AGENT_PROMPT,
-            retries=3,
-            deps_type=ResearchDependencies,
-        )
-
-        context_xml = format_context_for_prompt(state.context)
-        prompt = (
-            "Analyze gathered information and evaluate completeness for the original question.\n\n"
-            f"{context_xml}"
-        )
-        agent_deps = ResearchDependencies(
-            client=deps.client, context=state.context, console=deps.console
-        )
-        eval_result = await agent.run(prompt, deps=agent_deps)
-        output = eval_result.output
-
-        for insight in output.key_insights:
-            state.context.add_insight(insight)
-        for new_q in output.new_questions:
-            if new_q not in state.sub_questions:
-                state.sub_questions.append(new_q)
-
-        state.last_eval = output
-        state.iterations += 1
-
-        if output.key_insights:
-            log(deps.console, "   [bold]Key insights:[/bold]")
-            for ins in output.key_insights:
-                log(deps.console, f"   • {ins}")
-        log(
-            deps.console,
-            f"   Confidence: [yellow]{output.confidence_score:.1%}[/yellow]",
-        )
-        status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
-        log(deps.console, f"   Sufficient: {status}")
-
-        from haiku.rag.research.nodes.search import SearchDispatchNode
-
-        if (
-            output.is_sufficient
-            and output.confidence_score >= state.confidence_threshold
-        ) or state.iterations >= state.max_iterations:
-            log(deps.console, "\n[bold green]✅ Stopping research.[/bold green]")
-            return SynthesizeNode(self.provider, self.model)
-
-        return SearchDispatchNode(self.provider, self.model)

haiku/rag/research/nodes/plan.py

@@ -1,63 +0,0 @@
-from dataclasses import dataclass
-
-from pydantic_ai import Agent, RunContext
-from pydantic_graph import BaseNode, GraphRunContext
-
-from haiku.rag.research.common import get_model, log
-from haiku.rag.research.dependencies import ResearchDependencies
-from haiku.rag.research.models import ResearchPlan, ResearchReport
-from haiku.rag.research.nodes.search import SearchDispatchNode
-from haiku.rag.research.prompts import PLAN_PROMPT
-from haiku.rag.research.state import ResearchDeps, ResearchState
-
-
-@dataclass
-class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
-    provider: str
-    model: str
-
-    async def run(
-        self, ctx: GraphRunContext[ResearchState, ResearchDeps]
-    ) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
-        state = ctx.state
-        deps = ctx.deps
-
-        log(deps.console, "\n[bold cyan]📋 Creating research plan...[/bold cyan]")
-
-        plan_agent = Agent(
-            model=get_model(self.provider, self.model),
-            output_type=ResearchPlan,
-            instructions=(
-                PLAN_PROMPT
-                + "\n\nUse the gather_context tool once on the main question before planning."
-            ),
-            retries=3,
-            deps_type=ResearchDependencies,
-        )
-
-        @plan_agent.tool
-        async def gather_context(
-            ctx2: RunContext[ResearchDependencies], query: str, limit: int = 6
-        ) -> str:
-            results = await ctx2.deps.client.search(query, limit=limit)
-            expanded = await ctx2.deps.client.expand_context(results)
-            return "\n\n".join(chunk.content for chunk, _ in expanded)
-
-        prompt = (
-            "Plan a focused research approach for the main question.\n\n"
-            f"Main question: {state.question}"
-        )
-
-        agent_deps = ResearchDependencies(
-            client=deps.client, context=state.context, console=deps.console
-        )
-        plan_result = await plan_agent.run(prompt, deps=agent_deps)
-        state.sub_questions = list(plan_result.output.sub_questions)
-
-        log(deps.console, "\n[bold green]✅ Research Plan Created:[/bold green]")
-        log(deps.console, f"   [bold]Main Question:[/bold] {state.question}")
-        log(deps.console, "   [bold]Sub-questions:[/bold]")
-        for i, sq in enumerate(state.sub_questions, 1):
-            log(deps.console, f"      {i}. {sq}")
-
-        return SearchDispatchNode(self.provider, self.model)

haiku/rag/research/nodes/search.py

@@ -1,93 +0,0 @@
-import asyncio
-from dataclasses import dataclass
-from typing import Any
-
-from pydantic_ai import Agent, RunContext
-from pydantic_ai.format_prompt import format_as_xml
-from pydantic_ai.output import ToolOutput
-from pydantic_graph import BaseNode, GraphRunContext
-
-from haiku.rag.research.common import get_model, log
-from haiku.rag.research.dependencies import ResearchDependencies
-from haiku.rag.research.models import ResearchReport, SearchAnswer
-from haiku.rag.research.prompts import SEARCH_AGENT_PROMPT
-from haiku.rag.research.state import ResearchDeps, ResearchState
-
-
-@dataclass
-class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
-    provider: str
-    model: str
-
-    async def run(
-        self, ctx: GraphRunContext[ResearchState, ResearchDeps]
-    ) -> BaseNode[ResearchState, ResearchDeps, ResearchReport]:
-        state = ctx.state
-        deps = ctx.deps
-        if not state.sub_questions:
-            from haiku.rag.research.nodes.evaluate import EvaluateNode
-
-            return EvaluateNode(self.provider, self.model)
-
-        # Take up to max_concurrency questions and answer them concurrently
-        take = max(1, state.max_concurrency)
-        batch: list[str] = []
-        while state.sub_questions and len(batch) < take:
-            batch.append(state.sub_questions.pop(0))
-
-        async def answer_one(sub_q: str) -> SearchAnswer | None:
-            log(
-                deps.console,
-                f"\n[bold cyan]🔍 Searching & Answering:[/bold cyan] {sub_q}",
-            )
-            agent = Agent(
-                model=get_model(self.provider, self.model),
-                output_type=ToolOutput(SearchAnswer, max_retries=3),
-                instructions=SEARCH_AGENT_PROMPT,
-                retries=3,
-                deps_type=ResearchDependencies,
-            )
-
-            @agent.tool
-            async def search_and_answer(
-                ctx2: RunContext[ResearchDependencies], query: str, limit: int = 5
-            ) -> str:
-                search_results = await ctx2.deps.client.search(query, limit=limit)
-                expanded = await ctx2.deps.client.expand_context(search_results)
-
-                entries: list[dict[str, Any]] = [
-                    {
-                        "text": chunk.content,
-                        "score": score,
-                        "document_uri": (
-                            chunk.document_title or chunk.document_uri or ""
-                        ),
-                    }
-                    for chunk, score in expanded
-                ]
-                if not entries:
-                    return f"No relevant information found in the knowledge base for: {query}"
-
-                return format_as_xml(entries, root_tag="snippets")
-
-            agent_deps = ResearchDependencies(
-                client=deps.client, context=state.context, console=deps.console
-            )
-            try:
-                result = await agent.run(sub_q, deps=agent_deps)
-            except Exception as e:
-                log(deps.console, f"[red]Search failed:[/red] {e}")
-                return None
-
-            return result.output
-
-        answers = await asyncio.gather(*(answer_one(q) for q in batch))
-        for ans in answers:
-            if ans is None:
-                continue
-            state.context.add_qa_response(ans)
-            if deps.console:
-                preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "")
-                log(deps.console, f"   [green]✓[/green] {preview}")
-
-        return SearchDispatchNode(self.provider, self.model)

haiku/rag/research/nodes/synthesize.py

@@ -1,51 +0,0 @@
-from dataclasses import dataclass
-
-from pydantic_ai import Agent
-from pydantic_graph import BaseNode, End, GraphRunContext
-
-from haiku.rag.research.common import format_context_for_prompt, get_model, log
-from haiku.rag.research.dependencies import (
-    ResearchDependencies,
-)
-from haiku.rag.research.models import ResearchReport
-from haiku.rag.research.prompts import SYNTHESIS_AGENT_PROMPT
-from haiku.rag.research.state import ResearchDeps, ResearchState
-
-
-@dataclass
-class SynthesizeNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
-    provider: str
-    model: str
-
-    async def run(
-        self, ctx: GraphRunContext[ResearchState, ResearchDeps]
-    ) -> End[ResearchReport]:
-        state = ctx.state
-        deps = ctx.deps
-
-        log(
-            deps.console,
-            "\n[bold cyan]📝 Generating final research report...[/bold cyan]",
-        )
-
-        agent = Agent(
-            model=get_model(self.provider, self.model),
-            output_type=ResearchReport,
-            instructions=SYNTHESIS_AGENT_PROMPT,
-            retries=3,
-            deps_type=ResearchDependencies,
-        )
-
-        context_xml = format_context_for_prompt(state.context)
-        prompt = (
-            "Generate a comprehensive research report based on all gathered information.\n\n"
-            f"{context_xml}\n\n"
-            "Create a detailed report that synthesizes all findings into a coherent response."
-        )
-        agent_deps = ResearchDependencies(
-            client=deps.client, context=state.context, console=deps.console
-        )
-        result = await agent.run(prompt, deps=agent_deps)
-
-        log(deps.console, "[bold green]✅ Research complete![/bold green]")
-        return End(result.output)

haiku/rag/research/prompts.py

@@ -1,114 +0,0 @@
-PLAN_PROMPT = """You are the research orchestrator for a focused, iterative
-workflow.
-
-Responsibilities:
-1. Understand and decompose the main question
-2. Propose a minimal, high‑leverage plan
-3. Coordinate specialized agents to gather evidence
-4. Iterate based on gaps and new findings
-
-Plan requirements:
-- Produce at most 3 sub_questions that together cover the main question.
-- Each sub_question must be a standalone, self‑contained query that can run
-  without extra context. Include concrete entities, scope, timeframe, and any
-  qualifiers. Avoid ambiguous pronouns (it/they/this/that).
-- Prioritize the highest‑value aspects first; avoid redundancy and overlap.
-- Prefer questions that are likely answerable from the current knowledge base;
-  if coverage is uncertain, make scopes narrower and specific.
-- Order sub_questions by execution priority (most valuable first)."""
-
-SEARCH_AGENT_PROMPT = """You are a search and question‑answering specialist.
-
-Tasks:
-1. Search the knowledge base for relevant evidence.
-2. Analyze retrieved snippets.
-3. Provide an answer strictly grounded in that evidence.
-
-Tool usage:
-- Always call search_and_answer before drafting any answer.
-- The tool returns snippets with verbatim `text`, a relevance `score`, and the
-  originating document identifier (document title if available, otherwise URI).
-- You may call the tool multiple times to refine or broaden context, but do not
-  exceed 3 total calls. Favor precision over volume.
-- Use scores to prioritize evidence, but include only the minimal subset of
-  snippet texts (verbatim) in SearchAnswer.context (typically 1‑4).
-- Set SearchAnswer.sources to the corresponding document identifiers for the
-  snippets you used (title if available, otherwise URI; one per snippet; same
-  order as context). Context must be text‑only.
-- If no relevant information is found, clearly say so and return an empty
-  context list and sources list.
-
-Answering rules:
-- Be direct and specific; avoid meta commentary about the process.
-- Do not include any claims not supported by the provided snippets.
-- Prefer concise phrasing; avoid copying long passages.
-- When evidence is partial, state the limits explicitly in the answer."""
-
-EVALUATION_AGENT_PROMPT = """You are an analysis and evaluation specialist for
-the research workflow.
-
-Inputs available:
-- Original research question
-- Question–answer pairs produced by search
-- Raw search results and source metadata
-- Previously identified insights
-
-ANALYSIS:
-1. Extract the most important, non‑obvious insights from the collected evidence.
-2. Identify patterns, agreements, and disagreements across sources.
-3. Note material uncertainties and assumptions.
-
-EVALUATION:
-1. Decide if we have sufficient information to answer the original question.
-2. Provide a confidence_score in [0,1] considering:
-   - Coverage of the main question’s aspects
-   - Quality, consistency, and diversity of sources
-   - Depth and specificity of evidence
-3. List concrete gaps that still need investigation.
-4. Propose up to 3 new sub_questions that would close the highest‑value gaps.
-
-Strictness:
-- Only mark research as sufficient when all major aspects are addressed with
-  consistent, reliable evidence and no critical gaps remain.
-
-New sub_questions must:
-- Be genuinely new (not answered or duplicative; check qa_responses).
-- Be standalone and specific (entities, scope, timeframe/region if relevant).
-- Be actionable and scoped to the knowledge base (narrow if necessary).
-- Be ordered by expected impact (most valuable first)."""
-
-SYNTHESIS_AGENT_PROMPT = """You are a synthesis specialist producing the final
-research report.
-
-Goals:
-1. Synthesize all gathered information into a coherent narrative.
-2. Present findings clearly and concisely.
-3. Draw evidence‑based conclusions and recommendations.
-4. State limitations and uncertainties transparently.
-
-Report guidelines (map to output fields):
-- title: concise (5–12 words), informative.
-- executive_summary: 3–5 sentences summarizing the overall answer.
-- main_findings: 4–8 one‑sentence bullets; each reflects evidence from the
-  research (do not include inline citations or snippet text).
-- conclusions: 2–4 bullets that follow logically from findings.
-- recommendations: 2–5 actionable bullets tied to findings.
-- limitations: 1–3 bullets describing key constraints or uncertainties.
-- sources_summary: 2–4 sentences summarizing sources used and their reliability.
-
-Style:
-- Base all content solely on the collected evidence.
-- Be professional, objective, and specific.
-- Avoid meta commentary and refrain from speculation beyond the evidence."""
-
-PRESEARCH_AGENT_PROMPT = """You are a rapid research surveyor.
-
-Task:
-- Call gather_context once on the main question to obtain relevant text from
-  the knowledge base (KB).
-- Read that context and produce a short natural‑language summary of 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/state.py

@@ -1,25 +0,0 @@
-from dataclasses import dataclass, field
-
-from rich.console import Console
-
-from haiku.rag.client import HaikuRAG
-from haiku.rag.research.dependencies import ResearchContext
-from haiku.rag.research.models import EvaluationResult
-
-
-@dataclass
-class ResearchDeps:
-    client: HaikuRAG
-    console: Console | None = None
-
-
-@dataclass
-class ResearchState:
-    question: str
-    context: ResearchContext
-    sub_questions: list[str] = field(default_factory=list)
-    iterations: int = 0
-    max_iterations: int = 3
-    max_concurrency: int = 1
-    confidence_threshold: float = 0.8
-    last_eval: EvaluationResult | None = None

haiku/rag/store/__init__.py

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