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

haiku/rag/qa/agent.py

@@ -12,7 +12,9 @@ from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIO
 class SearchResult(BaseModel):
     content: str = Field(description="The document text content")
     score: float = Field(description="Relevance score (higher is more relevant)")
-    document_uri: str = Field(description="Source URI/path of the document")
+    document_uri: str = Field(
+        description="Source title (if available) or URI/path of the document"
+    )
 
 
 class Dependencies(BaseModel):
@@ -59,7 +61,7 @@ class QuestionAnswerAgent:
                 SearchResult(
                     content=chunk.content,
                     score=score,
-                    document_uri=chunk.document_uri or "",
+                    document_uri=(chunk.document_title or chunk.document_uri or ""),
                 )
                 for chunk, score in expanded_results
             ]

haiku/rag/qa/prompts.py

@@ -44,9 +44,9 @@ Guidelines:
 
 Citation Format:
 After your answer, include a "Citations:" section that lists:
-- The document URI from each search result used
+- The document title (if available) or URI from each search result used
 - A brief excerpt (first 50-100 characters) of the content that supported your answer
-- Format: "Citations:\n- [document_uri]: [content_excerpt]..."
+- Format: "Citations:\n- [document title or URI]: [content_excerpt]..."
 
 Example response format:
 [Your answer here]

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,9 +1,8 @@
 from pydantic import BaseModel, Field
-from pydantic_ai import format_as_xml
 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):
@@ -13,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(
@@ -23,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)
 
@@ -46,24 +45,3 @@ class ResearchDependencies(BaseModel):
     client: HaikuRAG = Field(description="RAG client for document operations")
     context: ResearchContext = Field(description="Shared research context")
     console: Console | None = None
-
-
-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,
-            }
-            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/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 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

@@ -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,93 @@
+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

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