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

haiku/rag/app.py

@@ -9,7 +9,13 @@ from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
 from haiku.rag.mcp import create_mcp_server
 from haiku.rag.monitor import FileWatcher
-from haiku.rag.research.orchestrator import ResearchOrchestrator
+from haiku.rag.research.dependencies import ResearchContext
+from haiku.rag.research.graph import (
+    PlanNode,
+    ResearchDeps,
+    ResearchState,
+    build_research_graph,
+)
 from haiku.rag.store.models.chunk import Chunk
 from haiku.rag.store.models.document import Document
 
@@ -80,30 +86,54 @@ class HaikuRAGApp:
                 self.console.print(f"[red]Error: {e}[/red]")
 
     async def research(
-        self, question: str, max_iterations: int = 3, verbose: bool = False
+        self,
+        question: str,
+        max_iterations: int = 3,
+        confidence_threshold: float = 0.8,
+        max_concurrency: int = 1,
+        verbose: bool = False,
     ):
-        """Run multi-agent research on a question."""
+        """Run research via the pydantic-graph pipeline (default)."""
         async with HaikuRAG(db_path=self.db_path) as client:
             try:
-                # Create orchestrator with default config or fallback to QA
-                orchestrator = ResearchOrchestrator()
-
                 if verbose:
-                    self.console.print(
-                        f"[bold cyan]Starting research with {orchestrator.provider}:{orchestrator.model}[/bold cyan]"
-                    )
+                    self.console.print("[bold cyan]Starting research[/bold cyan]")
                     self.console.print(f"[bold blue]Question:[/bold blue] {question}")
                     self.console.print()
 
-                # Conduct research
-                report = await orchestrator.conduct_research(
+                graph = build_research_graph()
+                state = ResearchState(
                     question=question,
-                    client=client,
+                    context=ResearchContext(original_question=question),
                     max_iterations=max_iterations,
-                    verbose=verbose,
-                    console=self.console if verbose else None,
+                    confidence_threshold=confidence_threshold,
+                    max_concurrency=max_concurrency,
+                )
+                deps = ResearchDeps(
+                    client=client, console=self.console if verbose else None
                 )
 
+                start = PlanNode(
+                    provider=Config.RESEARCH_PROVIDER or Config.QA_PROVIDER,
+                    model=Config.RESEARCH_MODEL or Config.QA_MODEL,
+                )
+                # Prefer graph.run; fall back to iter if unavailable
+                report = None
+                try:
+                    result = await graph.run(start, state=state, deps=deps)
+                    report = result.output
+                except Exception:
+                    from pydantic_graph import End
+
+                    async with graph.iter(start, state=state, deps=deps) as run:
+                        node = run.next_node
+                        while not isinstance(node, End):
+                            node = await run.next(node)
+                        if run.result:
+                            report = run.result.output
+                if report is None:
+                    raise RuntimeError("Graph did not produce a report")
+
                 # Display the report
                 self.console.print("[bold green]Research Report[/bold green]")
                 self.console.rule()
@@ -115,6 +145,12 @@ class HaikuRAGApp:
                 self.console.print(report.executive_summary)
                 self.console.print()
 
+                # Confidence (from last evaluation)
+                if state.last_eval:
+                    conf = state.last_eval.confidence_score  # type: ignore[attr-defined]
+                    self.console.print(f"[bold cyan]Confidence:[/bold cyan] {conf:.1%}")
+                    self.console.print()
+
                 # Main Findings
                 if report.main_findings:
                     self.console.print("[bold cyan]Main Findings:[/bold cyan]")

haiku/rag/cli.py

@@ -13,10 +13,10 @@ from haiku.rag.logging import configure_cli_logging
 from haiku.rag.migration import migrate_sqlite_to_lancedb
 from haiku.rag.utils import is_up_to_date
 
-logfire.configure(send_to_logfire="if-token-present")
-logfire.instrument_pydantic_ai()
-
-if not Config.ENV == "development":
+if Config.ENV == "development":
+    logfire.configure(send_to_logfire="if-token-present")
+    logfire.instrument_pydantic_ai()
+else:
     warnings.filterwarnings("ignore")
 
 cli = typer.Typer(
@@ -250,6 +250,16 @@ def research(
         "-n",
         help="Maximum search/analyze iterations",
     ),
+    confidence_threshold: float = typer.Option(
+        0.8,
+        "--confidence-threshold",
+        help="Minimum confidence (0-1) to stop",
+    ),
+    max_concurrency: int = typer.Option(
+        1,
+        "--max-concurrency",
+        help="Max concurrent searches per iteration (planned)",
+    ),
     db: Path = typer.Option(
         Config.DEFAULT_DATA_DIR / "haiku.rag.lancedb",
         "--db",
@@ -266,6 +276,8 @@ def research(
         app.research(
             question=question,
             max_iterations=max_iterations,
+            confidence_threshold=confidence_threshold,
+            max_concurrency=max_concurrency,
             verbose=verbose,
         )
     )

haiku/rag/client.py

@@ -388,7 +388,7 @@ class HaikuRAG:
                 all_chunks = adjacent_chunks + [chunk]
 
                 # Get the range of orders for this expanded chunk
-                orders = [c.metadata.get("order", 0) for c in all_chunks]
+                orders = [c.order for c in all_chunks]
                 min_order = min(orders)
                 max_order = max(orders)
 
@@ -398,9 +398,7 @@ class HaikuRAG:
                         "score": score,
                         "min_order": min_order,
                         "max_order": max_order,
-                        "all_chunks": sorted(
-                            all_chunks, key=lambda c: c.metadata.get("order", 0)
-                        ),
+                        "all_chunks": sorted(all_chunks, key=lambda c: c.order),
                     }
                 )
 
@@ -459,7 +457,7 @@ class HaikuRAG:
                 # Merge all_chunks and deduplicate by order
                 all_chunks_dict = {}
                 for chunk in current["all_chunks"] + range_info["all_chunks"]:
-                    order = chunk.metadata.get("order", 0)
+                    order = chunk.order
                     all_chunks_dict[order] = chunk
                 current["all_chunks"] = [
                     all_chunks_dict[order] for order in sorted(all_chunks_dict.keys())

haiku/rag/reranking/mxbai.py

@@ -1,4 +1,4 @@
-from mxbai_rerank import MxbaiRerankV2
+from mxbai_rerank import MxbaiRerankV2  # pyright: ignore[reportMissingImports]
 
 from haiku.rag.config import Config
 from haiku.rag.reranking.base import RerankerBase

haiku/rag/research/__init__.py

@@ -1,37 +1,20 @@
-"""Multi-agent research workflow for advanced RAG queries."""
-
-from haiku.rag.research.base import (
-    BaseResearchAgent,
-    ResearchOutput,
-    SearchAnswer,
-    SearchResult,
-)
 from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
-from haiku.rag.research.evaluation_agent import (
-    AnalysisEvaluationAgent,
-    EvaluationResult,
+from haiku.rag.research.graph import (
+    PlanNode,
+    ResearchDeps,
+    ResearchState,
+    build_research_graph,
 )
-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
+from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
 
 __all__ = [
-    # Base classes
-    "BaseResearchAgent",
     "ResearchDependencies",
     "ResearchContext",
-    "SearchResult",
-    "ResearchOutput",
-    # Specialized agents
     "SearchAnswer",
-    "SearchSpecialistAgent",
-    "PresearchSurveyAgent",
-    "AnalysisEvaluationAgent",
     "EvaluationResult",
-    "SynthesisAgent",
     "ResearchReport",
-    # Orchestrator
-    "ResearchOrchestrator",
-    "ResearchPlan",
+    "ResearchDeps",
+    "ResearchState",
+    "PlanNode",
+    "build_research_graph",
 ]

haiku/rag/research/common.py

@@ -0,0 +1,53 @@
+from typing import Any
+
+from pydantic_ai import format_as_xml
+from pydantic_ai.models.openai import OpenAIChatModel
+from pydantic_ai.providers.ollama import OllamaProvider
+from pydantic_ai.providers.openai import OpenAIProvider
+
+from haiku.rag.config import Config
+from haiku.rag.research.dependencies import ResearchContext
+
+
+def get_model(provider: str, model: str) -> Any:
+    if provider == "ollama":
+        return OpenAIChatModel(
+            model_name=model,
+            provider=OllamaProvider(base_url=f"{Config.OLLAMA_BASE_URL}/v1"),
+        )
+    elif provider == "vllm":
+        return OpenAIChatModel(
+            model_name=model,
+            provider=OpenAIProvider(
+                base_url=f"{Config.VLLM_RESEARCH_BASE_URL or Config.VLLM_QA_BASE_URL}/v1",
+                api_key="none",
+            ),
+        )
+    else:
+        return f"{provider}:{model}"
+
+
+def log(console, msg: str) -> None:
+    if console:
+        console.print(msg)
+
+
+def format_context_for_prompt(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,  # pyright: ignore[reportAttributeAccessIssue]
+            }
+            for qa in context.qa_responses
+        ],
+        "insights": context.insights,
+        "gaps": context.gaps,
+    }
+    return format_as_xml(context_data, root_tag="research_context")

haiku/rag/research/dependencies.py

@@ -1,7 +1,8 @@
 from pydantic import BaseModel, Field
+from rich.console import Console
 
 from haiku.rag.client import HaikuRAG
-from haiku.rag.research.base import SearchAnswer
+from haiku.rag.research.models import SearchAnswer
 
 
 class ResearchContext(BaseModel):
@@ -11,7 +12,7 @@ class ResearchContext(BaseModel):
     sub_questions: list[str] = Field(
         default_factory=list, description="Decomposed sub-questions"
     )
-    qa_responses: list["SearchAnswer"] = Field(
+    qa_responses: list[SearchAnswer] = Field(
         default_factory=list, description="Structured QA pairs used during research"
     )
     insights: list[str] = Field(
@@ -21,7 +22,7 @@ class ResearchContext(BaseModel):
         default_factory=list, description="Identified information gaps"
     )
 
-    def add_qa_response(self, qa: "SearchAnswer") -> None:
+    def add_qa_response(self, qa: SearchAnswer) -> None:
         """Add a structured QA response (minimal context already included)."""
         self.qa_responses.append(qa)
 
@@ -43,3 +44,4 @@ class ResearchDependencies(BaseModel):
 
     client: HaikuRAG = Field(description="RAG client for document operations")
     context: ResearchContext = Field(description="Shared research context")
+    console: Console | None = None

haiku/rag/research/graph.py

@@ -0,0 +1,29 @@
+from pydantic_graph import Graph
+
+from haiku.rag.research.models import ResearchReport
+from haiku.rag.research.nodes.evaluate import EvaluateNode
+from haiku.rag.research.nodes.plan import PlanNode
+from haiku.rag.research.nodes.search import SearchDispatchNode
+from haiku.rag.research.nodes.synthesize import SynthesizeNode
+from haiku.rag.research.state import ResearchDeps, ResearchState
+
+__all__ = [
+    "PlanNode",
+    "SearchDispatchNode",
+    "EvaluateNode",
+    "SynthesizeNode",
+    "ResearchState",
+    "ResearchDeps",
+    "build_research_graph",
+]
+
+
+def build_research_graph() -> Graph[ResearchState, ResearchDeps, ResearchReport]:
+    return Graph(
+        nodes=[
+            PlanNode,
+            SearchDispatchNode,
+            EvaluateNode,
+            SynthesizeNode,
+        ]
+    )

haiku/rag/research/models.py

@@ -0,0 +1,70 @@
+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 URIs corresponding to the snippets actually used in the"
+            " answer (one URI 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

@@ -0,0 +1,80 @@
+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

@@ -0,0 +1,63 @@
+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

@@ -0,0 +1,91 @@
+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_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)