haiku.rag-slim 0.16.0__py3-none-any.whl

haiku/rag/graph/common/nodes.py

@@ -0,0 +1,265 @@
+"""Common node implementations for graph workflows."""
+
+import asyncio
+from collections.abc import Awaitable, Callable
+from typing import Any, Protocol
+
+from pydantic_ai import Agent, RunContext
+from pydantic_ai.format_prompt import format_as_xml
+from pydantic_ai.output import ToolOutput
+from pydantic_graph.beta import StepContext
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.graph.agui.emitter import AGUIEmitter
+from haiku.rag.graph.common import get_model
+from haiku.rag.graph.common.models import ResearchPlan, SearchAnswer
+from haiku.rag.graph.common.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
+
+
+class GraphContext(Protocol):
+    """Protocol for graph context objects."""
+
+    original_question: str
+    sub_questions: list[str]
+
+    def add_qa_response(self, qa: SearchAnswer) -> None:
+        """Add a QA response to context."""
+        ...
+
+
+class GraphState(Protocol):
+    """Protocol for graph state objects."""
+
+    context: GraphContext
+    max_concurrency: int
+
+
+class GraphDeps(Protocol):
+    """Protocol for graph dependencies."""
+
+    client: HaikuRAG
+    agui_emitter: AGUIEmitter[Any, Any] | None
+    semaphore: asyncio.Semaphore | None
+
+
+class GraphAgentDeps(Protocol):
+    """Protocol for agent dependencies."""
+
+    client: HaikuRAG
+    context: GraphContext
+
+
+def create_plan_node[AgentDepsT: GraphAgentDeps](
+    provider: str,
+    model: str,
+    deps_type: type[AgentDepsT],
+    activity_message: str = "Creating plan",
+    output_retries: int | None = None,
+) -> Callable[[StepContext[Any, Any, None]], Awaitable[None]]:
+    """Create a plan node for any graph.
+
+    Args:
+        provider: Model provider (e.g., 'openai', 'anthropic')
+        model: Model name
+        deps_type: Type of dependencies for the agent (e.g., ResearchDependencies, DeepQADependencies)
+        activity_message: Message to show during planning activity
+        output_retries: Number of output retries for the agent (optional)
+
+    Returns:
+        Async function that can be used as a graph step
+    """
+
+    async def plan(ctx: StepContext[Any, Any, None], /) -> None:
+        state: GraphState = ctx.state  # type: ignore[assignment]
+        deps: GraphDeps = ctx.deps  # type: ignore[assignment]
+
+        if deps.agui_emitter:
+            deps.agui_emitter.start_step("plan")
+            deps.agui_emitter.update_activity("planning", activity_message)
+
+        try:
+            # Build agent configuration
+            agent_config = {
+                "model": get_model(provider, model),
+                "output_type": ResearchPlan,
+                "instructions": (
+                    PLAN_PROMPT
+                    + "\n\nUse the gather_context tool once on the main question before planning."
+                ),
+                "retries": 3,
+                "deps_type": deps_type,
+            }
+            if output_retries is not None:
+                agent_config["output_retries"] = output_retries
+
+            plan_agent = Agent(**agent_config)
+
+            @plan_agent.tool
+            async def gather_context(
+                ctx2: RunContext[AgentDepsT], 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)
+
+            # Tool is registered via decorator above
+            _ = gather_context
+
+            prompt = (
+                "Plan a focused approach for the main question.\n\n"
+                f"Main question: {state.context.original_question}"
+            )
+
+            # Create agent dependencies
+            agent_deps = deps_type(client=deps.client, context=state.context)  # type: ignore[call-arg]
+            plan_result = await plan_agent.run(prompt, deps=agent_deps)
+            state.context.sub_questions = list(plan_result.output.sub_questions)
+
+            # State now contains the plan - emit state update and narrate
+            if deps.agui_emitter:
+                deps.agui_emitter.update_state(state)
+                count = len(state.context.sub_questions)
+                deps.agui_emitter.update_activity(
+                    "planning", f"Created plan with {count} sub-questions"
+                )
+        finally:
+            if deps.agui_emitter:
+                deps.agui_emitter.finish_step()
+
+    return plan
+
+
+def create_search_node[AgentDepsT: GraphAgentDeps](
+    provider: str,
+    model: str,
+    deps_type: type[AgentDepsT],
+    with_step_wrapper: bool = True,
+    success_message_format: str = "Answered: {sub_q}",
+    handle_exceptions: bool = False,
+) -> Callable[[StepContext[Any, Any, str]], Awaitable[SearchAnswer]]:
+    """Create a search_one node for any graph.
+
+    Args:
+        provider: Model provider
+        model: Model name
+        deps_type: Type of dependencies for the agent
+        with_step_wrapper: Whether to wrap with agui_emitter start/finish step
+        success_message_format: Format string for success activity message
+        handle_exceptions: Whether to handle exceptions with fallback answer
+
+    Returns:
+        Async function that can be used as a graph step
+    """
+
+    async def search_one(ctx: StepContext[Any, Any, str], /) -> SearchAnswer:
+        state: GraphState = ctx.state  # type: ignore[assignment]
+        deps: GraphDeps = ctx.deps  # type: ignore[assignment]
+        sub_q = ctx.inputs
+
+        # Create unique step name from question text
+        step_name = f"search: {sub_q}"
+
+        if deps.agui_emitter and with_step_wrapper:
+            deps.agui_emitter.start_step(step_name)
+
+        try:
+            # Create semaphore if not already provided
+            if deps.semaphore is None:
+                deps.semaphore = asyncio.Semaphore(state.max_concurrency)
+
+            # Use semaphore to control concurrency
+            async with deps.semaphore:
+                return await _do_search(
+                    state,
+                    deps,
+                    sub_q,
+                    provider,
+                    model,
+                    deps_type,
+                    success_message_format,
+                    handle_exceptions,
+                )
+        finally:
+            if deps.agui_emitter and with_step_wrapper:
+                deps.agui_emitter.finish_step()
+
+    return search_one
+
+
+async def _do_search[AgentDepsT: GraphAgentDeps](
+    state: GraphState,
+    deps: GraphDeps,
+    sub_q: str,
+    provider: str,
+    model: str,
+    deps_type: type[AgentDepsT],
+    success_message_format: str,
+    handle_exceptions: bool,
+) -> SearchAnswer:
+    """Internal search implementation."""
+    if deps.agui_emitter:
+        deps.agui_emitter.update_activity("searching", f"Searching: {sub_q}")
+
+    agent = Agent(
+        model=get_model(provider, model),
+        output_type=ToolOutput(SearchAnswer, max_retries=3),
+        instructions=SEARCH_AGENT_PROMPT,
+        retries=3,
+        deps_type=deps_type,
+    )
+
+    @agent.tool
+    async def search_and_answer(
+        ctx2: RunContext[AgentDepsT], 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")
+
+    # Tool is registered via decorator above
+    _ = search_and_answer
+
+    agent_deps = deps_type(client=deps.client, context=state.context)  # type: ignore[call-arg]
+
+    try:
+        result = await agent.run(sub_q, deps=agent_deps)
+        answer = result.output
+        if answer:
+            state.context.add_qa_response(answer)
+            # State updated with new answer - emit state update and narrate
+            if deps.agui_emitter:
+                deps.agui_emitter.update_state(state)
+                # Format the success message
+                if "{confidence" in success_message_format:
+                    message = success_message_format.format(
+                        sub_q=sub_q, confidence=answer.confidence
+                    )
+                else:
+                    message = success_message_format.format(sub_q=sub_q)
+                deps.agui_emitter.update_activity("searching", message)
+        return answer
+    except Exception as e:
+        if handle_exceptions:
+            # Narrate the error
+            if deps.agui_emitter:
+                deps.agui_emitter.update_activity("searching", f"Search failed: {e}")
+            failure_answer = SearchAnswer(
+                query=sub_q,
+                answer=f"Search failed after retries: {str(e)}",
+                confidence=0.0,
+            )
+            return failure_answer
+        else:
+            raise

haiku/rag/graph/common/prompts.py

@@ -0,0 +1,46 @@
+"""Common prompts used across different graph implementations."""
+
+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."""

haiku/rag/graph/common/utils.py

@@ -0,0 +1,44 @@
+"""Common utilities for all graph implementations."""
+
+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
+
+
+def get_model(provider: str, model: str) -> OpenAIChatModel | str:
+    """
+    Get a model instance for the specified provider and model name.
+
+    Args:
+        provider: The model provider ("ollama", "vllm", or other)
+        model: The model name
+
+    Returns:
+        A configured model instance
+
+    Raises:
+        ValueError: If the provider is unknown
+    """
+    if provider == "ollama":
+        return OpenAIChatModel(
+            model_name=model,
+            provider=OllamaProvider(base_url=f"{Config.providers.ollama.base_url}/v1"),
+        )
+    elif provider == "vllm":
+        return OpenAIChatModel(
+            model_name=model,
+            provider=OpenAIProvider(
+                base_url=f"{Config.providers.vllm.research_base_url or Config.providers.vllm.qa_base_url}/v1",
+                api_key="none",
+            ),
+        )
+    elif provider in ("openai", "anthropic", "gemini", "groq", "bedrock"):
+        # These providers use string format
+        return f"{provider}:{model}"
+    else:
+        raise ValueError(
+            f"Unknown model provider: {provider}. "
+            f"Supported providers: ollama, vllm, openai, anthropic, gemini, groq, bedrock"
+        )

haiku/rag/graph/deep_qa/__init__.py

@@ -0,0 +1 @@
+from haiku.rag.graph.deep_qa.models import DeepQAAnswer

haiku/rag/graph/deep_qa/dependencies.py

@@ -0,0 +1,27 @@
+from pydantic import BaseModel, Field
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.graph.common.models import SearchAnswer
+
+
+class DeepQAContext(BaseModel):
+    original_question: str = Field(description="The original question")
+    sub_questions: list[str] = Field(
+        default_factory=list, description="Decomposed sub-questions"
+    )
+    qa_responses: list[SearchAnswer] = Field(
+        default_factory=list, description="QA pairs collected during answering"
+    )
+    use_citations: bool = Field(
+        default=False, description="Whether to include citations in the answer"
+    )
+
+    def add_qa_response(self, qa: SearchAnswer) -> None:
+        self.qa_responses.append(qa)
+
+
+class DeepQADependencies(BaseModel):
+    model_config = {"arbitrary_types_allowed": True}
+
+    client: HaikuRAG = Field(description="RAG client for document operations")
+    context: DeepQAContext = Field(description="Shared QA context")

haiku/rag/graph/deep_qa/graph.py

@@ -0,0 +1,243 @@
+from pydantic_ai import Agent
+from pydantic_ai.format_prompt import format_as_xml
+from pydantic_graph.beta import Graph, GraphBuilder, StepContext
+from pydantic_graph.beta.join import reduce_list_append
+
+from haiku.rag.config import Config
+from haiku.rag.config.models import AppConfig
+from haiku.rag.graph.common import get_model
+from haiku.rag.graph.common.models import SearchAnswer
+from haiku.rag.graph.common.nodes import create_plan_node, create_search_node
+from haiku.rag.graph.deep_qa.dependencies import DeepQADependencies
+from haiku.rag.graph.deep_qa.models import DeepQAAnswer, DeepQAEvaluation
+from haiku.rag.graph.deep_qa.prompts import (
+    DECISION_PROMPT,
+    SYNTHESIS_PROMPT,
+    SYNTHESIS_PROMPT_WITH_CITATIONS,
+)
+from haiku.rag.graph.deep_qa.state import DeepQADeps, DeepQAState
+
+
+def build_deep_qa_graph(
+    config: AppConfig = Config,
+) -> Graph[DeepQAState, DeepQADeps, None, DeepQAAnswer]:
+    """Build the Deep QA graph.
+
+    Args:
+        config: AppConfig object (uses config.qa for provider, model, and graph parameters)
+
+    Returns:
+        Configured Deep QA graph
+    """
+    provider = config.qa.provider
+    model = config.qa.model
+    g = GraphBuilder(
+        state_type=DeepQAState,
+        deps_type=DeepQADeps,
+        output_type=DeepQAAnswer,
+    )
+
+    # Create and register the plan node using the factory
+    plan = g.step(
+        create_plan_node(
+            provider=provider,
+            model=model,
+            deps_type=DeepQADependencies,  # type: ignore[arg-type]
+            activity_message="Planning approach",
+            output_retries=None,  # Deep QA doesn't use output_retries
+        )
+    )  # type: ignore[arg-type]
+
+    # Create and register the search_one node using the factory
+    search_one = g.step(
+        create_search_node(
+            provider=provider,
+            model=model,
+            deps_type=DeepQADependencies,  # type: ignore[arg-type]
+            with_step_wrapper=False,  # Deep QA doesn't wrap with agui_emitter step
+            success_message_format="Answered: {sub_q}",
+            handle_exceptions=True,
+        )
+    )  # type: ignore[arg-type]
+
+    @g.step
+    async def get_batch(
+        ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
+    ) -> list[str] | None:
+        """Get all remaining questions for this iteration."""
+        state = ctx.state
+
+        if not state.context.sub_questions:
+            return None
+
+        # Take ALL remaining questions - max_concurrency controls parallel execution within .map()
+        batch = list(state.context.sub_questions)
+        state.context.sub_questions.clear()
+        return batch
+
+    @g.step
+    async def decide(
+        ctx: StepContext[DeepQAState, DeepQADeps, list[SearchAnswer]],
+    ) -> bool:
+        state = ctx.state
+        deps = ctx.deps
+
+        if deps.agui_emitter:
+            deps.agui_emitter.start_step("decide")
+            deps.agui_emitter.update_activity(
+                "evaluating", "Evaluating information sufficiency"
+            )
+
+        try:
+            agent = Agent(
+                model=get_model(provider, model),
+                output_type=DeepQAEvaluation,
+                instructions=DECISION_PROMPT,
+                retries=3,
+                deps_type=DeepQADependencies,
+            )
+
+            context_data = {
+                "original_question": state.context.original_question,
+                "gathered_answers": [
+                    {
+                        "question": qa.query,
+                        "answer": qa.answer,
+                        "sources": qa.sources,
+                    }
+                    for qa in state.context.qa_responses
+                ],
+            }
+            context_xml = format_as_xml(context_data, root_tag="gathered_information")
+
+            prompt = (
+                "Evaluate whether we have sufficient information to answer the question.\n\n"
+                f"{context_xml}"
+            )
+
+            agent_deps = DeepQADependencies(
+                client=deps.client,
+                context=state.context,
+            )
+            result = await agent.run(prompt, deps=agent_deps)
+            evaluation = result.output
+
+            state.iterations += 1
+
+            for new_q in evaluation.new_questions:
+                if new_q not in state.context.sub_questions:
+                    state.context.sub_questions.append(new_q)
+
+            if deps.agui_emitter:
+                deps.agui_emitter.update_state(state)
+                status = "sufficient" if evaluation.is_sufficient else "insufficient"
+                deps.agui_emitter.update_activity(
+                    "evaluating",
+                    f"Information {status} after {state.iterations} iteration(s)",
+                )
+
+            should_continue = (
+                not evaluation.is_sufficient and state.iterations < state.max_iterations
+            )
+
+            return should_continue
+        finally:
+            if deps.agui_emitter:
+                deps.agui_emitter.finish_step()
+
+    @g.step
+    async def synthesize(
+        ctx: StepContext[DeepQAState, DeepQADeps, None | bool],
+    ) -> DeepQAAnswer:
+        state = ctx.state
+        deps = ctx.deps
+
+        if deps.agui_emitter:
+            deps.agui_emitter.start_step("synthesize")
+            deps.agui_emitter.update_activity(
+                "synthesizing", "Synthesizing final answer"
+            )
+
+        try:
+            prompt_template = (
+                SYNTHESIS_PROMPT_WITH_CITATIONS
+                if state.context.use_citations
+                else SYNTHESIS_PROMPT
+            )
+
+            agent = Agent(
+                model=get_model(provider, model),
+                output_type=DeepQAAnswer,
+                instructions=prompt_template,
+                retries=3,
+                deps_type=DeepQADependencies,
+            )
+
+            context_data = {
+                "original_question": state.context.original_question,
+                "sub_answers": [
+                    {
+                        "question": qa.query,
+                        "answer": qa.answer,
+                        "sources": qa.sources,
+                    }
+                    for qa in state.context.qa_responses
+                ],
+            }
+            context_xml = format_as_xml(context_data, root_tag="gathered_information")
+
+            prompt = f"Synthesize a comprehensive answer to the original question.\n\n{context_xml}"
+
+            agent_deps = DeepQADependencies(
+                client=deps.client,
+                context=state.context,
+            )
+            result = await agent.run(prompt, deps=agent_deps)
+
+            if deps.agui_emitter:
+                deps.agui_emitter.update_activity("synthesizing", "Answer complete")
+
+            return result.output
+        finally:
+            if deps.agui_emitter:
+                deps.agui_emitter.finish_step()
+
+    # Build the graph structure
+    collect_answers = g.join(
+        reduce_list_append,
+        initial_factory=list[SearchAnswer],
+    )
+
+    g.add(
+        g.edge_from(g.start_node).to(plan),
+        g.edge_from(plan).to(get_batch),
+    )
+
+    # Branch based on whether we have questions
+    g.add(
+        g.edge_from(get_batch).to(
+            g.decision()
+            .branch(g.match(list).label("Has questions").map().to(search_one))
+            .branch(g.match(type(None)).label("No questions").to(synthesize))
+        ),
+        g.edge_from(search_one).to(collect_answers),
+        g.edge_from(collect_answers).to(decide),
+    )
+
+    # Branch based on decision
+    g.add(
+        g.edge_from(decide).to(
+            g.decision()
+            .branch(
+                g.match(bool, matches=lambda x: x).label("Continue QA").to(get_batch)
+            )
+            .branch(
+                g.match(bool, matches=lambda x: not x)
+                .label("Done with QA")
+                .to(synthesize)
+            )
+        ),
+        g.edge_from(synthesize).to(g.end_node),
+    )
+
+    return g.build()

haiku/rag/graph/deep_qa/models.py

@@ -0,0 +1,20 @@
+from pydantic import BaseModel, Field
+
+
+class DeepQAEvaluation(BaseModel):
+    is_sufficient: bool = Field(
+        description="Whether we have sufficient information to answer the question"
+    )
+    reasoning: str = Field(description="Explanation of the sufficiency assessment")
+    new_questions: list[str] = Field(
+        description="Additional sub-questions needed if insufficient",
+        default_factory=list,
+    )
+
+
+class DeepQAAnswer(BaseModel):
+    answer: str = Field(description="The comprehensive answer to the question")
+    sources: list[str] = Field(
+        description="Document titles or URIs used to generate the answer",
+        default_factory=list,
+    )