haiku.rag 0.11.2__py3-none-any.whl → 0.11.4__py3-none-any.whl

haiku/rag/qa/deep/nodes.py

@@ -0,0 +1,303 @@
+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, End, GraphRunContext
+
+from haiku.rag.graph.common import get_model, log
+from haiku.rag.graph.models import ResearchPlan, SearchAnswer
+from haiku.rag.graph.prompts import PLAN_PROMPT, SEARCH_AGENT_PROMPT
+from haiku.rag.qa.deep.dependencies import DeepQADependencies
+from haiku.rag.qa.deep.models import DeepQAAnswer, DeepQAEvaluation
+from haiku.rag.qa.deep.prompts import (
+    DECISION_PROMPT,
+    SYNTHESIS_PROMPT,
+    SYNTHESIS_PROMPT_WITH_CITATIONS,
+)
+from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
+
+
+@dataclass
+class DeepQAPlanNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
+    provider: str
+    model: str
+
+    async def run(
+        self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
+    ) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]:
+        state = ctx.state
+        deps = ctx.deps
+
+        log(deps, state, "\n[bold cyan]📋 Planning approach...[/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=DeepQADependencies,
+        )
+
+        @plan_agent.tool
+        async def gather_context(
+            ctx2: RunContext[DeepQADependencies], 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 approach for answering the main question.\n\n"
+            f"Main question: {state.context.original_question}"
+        )
+
+        agent_deps = DeepQADependencies(
+            client=deps.client,
+            context=state.context,
+            console=deps.console,
+        )
+        plan_result = await plan_agent.run(prompt, deps=agent_deps)
+        state.context.sub_questions = list(plan_result.output.sub_questions)[
+            : state.max_sub_questions
+        ]
+
+        log(deps, state, "\n[bold green]✅ Plan Created:[/bold green]")
+        log(
+            deps,
+            state,
+            f"   [bold]Main Question:[/bold] {state.context.original_question}",
+        )
+        log(deps, state, "   [bold]Sub-questions:[/bold]")
+        for i, sq in enumerate(state.context.sub_questions, 1):
+            log(deps, state, f"      {i}. {sq}")
+
+        return DeepQASearchDispatchNode(self.provider, self.model)
+
+
+@dataclass
+class DeepQASearchDispatchNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
+    provider: str
+    model: str
+
+    async def run(
+        self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
+    ) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]:
+        state = ctx.state
+        deps = ctx.deps
+
+        if not state.context.sub_questions:
+            return DeepQADecisionNode(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.context.sub_questions and len(batch) < take:
+            batch.append(state.context.sub_questions.pop(0))
+
+        async def answer_one(sub_q: str) -> SearchAnswer | None:
+            log(
+                deps,
+                state,
+                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=DeepQADependencies,
+            )
+
+            @agent.tool
+            async def search_and_answer(
+                ctx2: RunContext[DeepQADependencies], 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 = DeepQADependencies(
+                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, state, 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)
+            preview = ans.answer[:150] + ("…" if len(ans.answer) > 150 else "")
+            log(deps, state, f"   [green]✓[/green] {preview}")
+
+        return DeepQASearchDispatchNode(self.provider, self.model)
+
+
+@dataclass
+class DeepQADecisionNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
+    provider: str
+    model: str
+
+    async def run(
+        self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
+    ) -> BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]:
+        state = ctx.state
+        deps = ctx.deps
+
+        log(
+            deps,
+            state,
+            "\n[bold cyan]📊 Evaluating information sufficiency...[/bold cyan]",
+        )
+
+        agent = Agent(
+            model=get_model(self.provider, self.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,
+            console=deps.console,
+        )
+        result = await agent.run(prompt, deps=agent_deps)
+        evaluation = result.output
+
+        state.iterations += 1
+
+        log(deps, state, f"   [bold]Assessment:[/bold] {evaluation.reasoning}")
+        status = "[green]Yes[/green]" if evaluation.is_sufficient else "[red]No[/red]"
+        log(deps, state, f"   Sufficient: {status}")
+
+        # Add new questions if not sufficient
+        for new_q in evaluation.new_questions:
+            if new_q not in state.context.sub_questions:
+                state.context.sub_questions.append(new_q)
+
+        if evaluation.new_questions:
+            log(deps, state, "   [cyan]New questions:[/cyan]")
+            for question in evaluation.new_questions:
+                log(deps, state, f"   • {question}")
+
+        # Decide next step
+        if evaluation.is_sufficient or state.iterations >= state.max_iterations:
+            if state.iterations >= state.max_iterations:
+                log(
+                    deps,
+                    state,
+                    f"\n[bold yellow]⚠️  Reached max iterations ({state.max_iterations})[/bold yellow]",
+                )
+            log(deps, state, "\n[bold green]✅ Moving to synthesis.[/bold green]")
+            return DeepQASynthesizeNode(self.provider, self.model)
+
+        log(
+            deps,
+            state,
+            f"\n[bold cyan]🔄 Starting iteration {state.iterations + 1}...[/bold cyan]",
+        )
+        return DeepQASearchDispatchNode(self.provider, self.model)
+
+
+@dataclass
+class DeepQASynthesizeNode(BaseNode[DeepQAState, DeepQADeps, DeepQAAnswer]):
+    provider: str
+    model: str
+
+    async def run(
+        self, ctx: GraphRunContext[DeepQAState, DeepQADeps]
+    ) -> End[DeepQAAnswer]:
+        state = ctx.state
+        deps = ctx.deps
+
+        log(
+            deps,
+            state,
+            "\n[bold cyan]📝 Synthesizing final answer...[/bold cyan]",
+        )
+
+        prompt_template = (
+            SYNTHESIS_PROMPT_WITH_CITATIONS
+            if state.context.use_citations
+            else SYNTHESIS_PROMPT
+        )
+
+        agent = Agent(
+            model=get_model(self.provider, self.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,
+            console=deps.console,
+        )
+        result = await agent.run(prompt, deps=agent_deps)
+
+        log(deps, state, "[bold green]✅ Answer complete![/bold green]")
+        return End(result.output)

haiku/rag/qa/deep/prompts.py

@@ -0,0 +1,57 @@
+SYNTHESIS_PROMPT = """You are an expert at synthesizing information into clear, concise answers.
+
+Task:
+- Combine the gathered information from sub-questions into a single comprehensive answer
+- Answer the original question directly and completely
+- Base your answer strictly on the provided evidence
+- Be clear, accurate, and well-structured
+
+Output format:
+- answer: The complete answer to the original question (2-4 paragraphs)
+- sources: List of document titles/URIs used (extract from the sub-answers)
+
+Guidelines:
+- Start directly with the answer - no preamble like "Based on the research..."
+- Use a clear, professional tone
+- Organize information logically
+- If evidence is incomplete, state limitations clearly
+- Do not include any claims not supported by the gathered information"""
+
+SYNTHESIS_PROMPT_WITH_CITATIONS = """You are an expert at synthesizing information into clear, concise answers with proper citations.
+
+Task:
+- Combine the gathered information from sub-questions into a single comprehensive answer
+- Answer the original question directly and completely
+- Base your answer strictly on the provided evidence
+- Include inline citations using [Source Title] format
+
+Output format:
+- answer: The complete answer with inline citations (2-4 paragraphs)
+- sources: List of document titles/URIs used (extract from the sub-answers)
+
+Guidelines:
+- Start directly with the answer - no preamble like "Based on the research..."
+- Add citations after each claim: [Source Title]
+- Use a clear, professional tone
+- Organize information logically
+- If evidence is incomplete, state limitations clearly
+- Do not include any claims not supported by the gathered information"""
+
+DECISION_PROMPT = """You are an expert at evaluating whether gathered information is sufficient to answer a question.
+
+Task:
+- Review the original question and all gathered sub-question answers
+- Determine if we have enough information to provide a comprehensive answer
+- If insufficient, suggest specific new sub-questions to fill the gaps
+
+Output format:
+- is_sufficient: Boolean indicating if we can answer the question comprehensively
+- reasoning: Clear explanation of your assessment
+- new_questions: List of specific follow-up questions needed (empty if sufficient)
+
+Guidelines:
+- Be strict but reasonable in your assessment
+- Focus on whether core aspects of the question are addressed
+- New questions should be specific and distinct from what's been asked
+- Limit new questions to 2-3 maximum
+- Consider whether additional searches would meaningfully improve the answer"""

haiku/rag/qa/deep/state.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass
+
+from rich.console import Console
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.qa.deep.dependencies import DeepQAContext
+
+
+@dataclass
+class DeepQADeps:
+    client: HaikuRAG
+    console: Console | None = None
+
+    def emit_log(self, message: str, state: "DeepQAState | None" = None) -> None:
+        if self.console:
+            self.console.print(message)
+
+
+@dataclass
+class DeepQAState:
+    context: DeepQAContext
+    max_sub_questions: int = 3
+    max_iterations: int = 2
+    max_concurrency: int = 3
+    iterations: int = 0

haiku/rag/research/__init__.py

@@ -1,28 +1,3 @@
+from haiku.rag.graph.models import SearchAnswer
 from haiku.rag.research.dependencies import ResearchContext, ResearchDependencies
-from haiku.rag.research.graph import (
-    PlanNode,
-    ResearchDeps,
-    ResearchState,
-    build_research_graph,
-)
-from haiku.rag.research.models import EvaluationResult, ResearchReport, SearchAnswer
-from haiku.rag.research.stream import (
-    ResearchStateSnapshot,
-    ResearchStreamEvent,
-    stream_research_graph,
-)
-
-__all__ = [
-    "ResearchDependencies",
-    "ResearchContext",
-    "SearchAnswer",
-    "EvaluationResult",
-    "ResearchReport",
-    "ResearchDeps",
-    "ResearchState",
-    "PlanNode",
-    "build_research_graph",
-    "stream_research_graph",
-    "ResearchStreamEvent",
-    "ResearchStateSnapshot",
-]
+from haiku.rag.research.models import EvaluationResult, ResearchReport

haiku/rag/research/common.py

@@ -1,39 +1,8 @@
-from typing import TYPE_CHECKING, 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
 from haiku.rag.research.models import InsightAnalysis
 
-if TYPE_CHECKING:  # pragma: no cover
-    from haiku.rag.research.state import ResearchDeps, ResearchState
-
-
-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(deps: "ResearchDeps", state: "ResearchState", msg: str) -> None:
-    deps.emit_log(msg, state)
-
 
 def format_context_for_prompt(context: ResearchContext) -> str:
     """Format the research context as XML for inclusion in prompts."""

haiku/rag/research/dependencies.py

@@ -4,11 +4,11 @@ from pydantic import BaseModel, Field
 from rich.console import Console
 
 from haiku.rag.client import HaikuRAG
+from haiku.rag.graph.models import SearchAnswer
 from haiku.rag.research.models import (
     GapRecord,
     InsightAnalysis,
     InsightRecord,
-    SearchAnswer,
 )
 from haiku.rag.research.stream import ResearchStream
 

haiku/rag/research/graph.py

@@ -1,23 +1,12 @@
 from pydantic_graph import Graph
 
+from haiku.rag.graph.nodes.analysis import AnalyzeInsightsNode, DecisionNode
+from haiku.rag.graph.nodes.plan import PlanNode
+from haiku.rag.graph.nodes.search import SearchDispatchNode
+from haiku.rag.graph.nodes.synthesize import SynthesizeNode
 from haiku.rag.research.models import ResearchReport
-from haiku.rag.research.nodes.analysis import AnalyzeInsightsNode, DecisionNode
-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",
-    "AnalyzeInsightsNode",
-    "DecisionNode",
-    "SynthesizeNode",
-    "ResearchState",
-    "ResearchDeps",
-    "build_research_graph",
-]
-
 
 def build_research_graph() -> Graph[ResearchState, ResearchDeps, ResearchReport]:
     return Graph(

haiku/rag/research/models.py

@@ -131,31 +131,6 @@ class InsightAnalysis(BaseModel):
     )
 
 
-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."""
 

haiku/rag/research/prompts.py

@@ -1,49 +1,3 @@
-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."""
-
 INSIGHT_AGENT_PROMPT = """You are the insight aggregation specialist for the
 research workflow.
 

haiku/rag/store/engine.py

@@ -1,3 +1,4 @@
+import asyncio
 import json
 import logging
 from datetime import timedelta
@@ -51,6 +52,7 @@ class Store:
     def __init__(self, db_path: Path, skip_validation: bool = False):
         self.db_path: Path = db_path
         self.embedder = get_embedder()
+        self._vacuum_lock = asyncio.Lock()
 
         # Create the ChunkRecord model with the correct vector dimension
         self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
@@ -78,14 +80,40 @@ class Store:
         if not skip_validation:
             self._validate_configuration()
 
-    def vacuum(self) -> None:
-        """Optimize and clean up old versions across all tables to reduce disk usage."""
+    async def vacuum(self, retention_seconds: int | None = None) -> None:
+        """Optimize and clean up old versions across all tables to reduce disk usage.
+
+        Args:
+            retention_seconds: Retention threshold in seconds. Only versions older
+                              than this will be removed. If None, uses Config.VACUUM_RETENTION_SECONDS.
+
+        Note:
+            If vacuum is already running, this method returns immediately without blocking.
+            Use asyncio.create_task(store.vacuum()) for non-blocking background execution.
+        """
         if self._has_cloud_config() and str(Config.LANCEDB_URI).startswith("db://"):
             return
 
-        # Perform maintenance per table using optimize() with cleanup_older_than 0
-        for table in [self.documents_table, self.chunks_table, self.settings_table]:
-            table.optimize(cleanup_older_than=timedelta(0))
+        # Skip if already running (non-blocking)
+        if self._vacuum_lock.locked():
+            return
+
+        async with self._vacuum_lock:
+            try:
+                # Evaluate config at runtime to allow dynamic changes
+                if retention_seconds is None:
+                    retention_seconds = Config.VACUUM_RETENTION_SECONDS
+                # Perform maintenance per table using optimize() with configurable retention
+                retention = timedelta(seconds=retention_seconds)
+                for table in [
+                    self.documents_table,
+                    self.chunks_table,
+                    self.settings_table,
+                ]:
+                    table.optimize(cleanup_older_than=retention)
+            except (RuntimeError, OSError) as e:
+                # Handle resource errors gracefully
+                logger.debug(f"Vacuum skipped due to resource constraints: {e}")
 
     def _connect_to_lancedb(self, db_path: Path):
         """Establish connection to LanceDB (local, cloud, or object storage)."""