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

haiku/rag/graph/common/nodes.py

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

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

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

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

haiku/rag/graph/deep_qa/dependencies.py

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

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

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