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

haiku/rag/app.py

@@ -194,10 +194,44 @@ class HaikuRAGApp:
             for chunk, score in results:
                 self._rich_print_search_result(chunk, score)
 
-    async def ask(self, question: str, cite: bool = False):
+    async def ask(
+        self,
+        question: str,
+        cite: bool = False,
+        deep: bool = False,
+        verbose: bool = False,
+    ):
         async with HaikuRAG(db_path=self.db_path) as self.client:
             try:
-                answer = await self.client.ask(question, cite=cite)
+                if deep:
+                    from rich.console import Console
+
+                    from haiku.rag.qa.deep.dependencies import DeepQAContext
+                    from haiku.rag.qa.deep.graph import build_deep_qa_graph
+                    from haiku.rag.qa.deep.nodes import DeepQAPlanNode
+                    from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
+
+                    graph = build_deep_qa_graph()
+                    context = DeepQAContext(
+                        original_question=question, use_citations=cite
+                    )
+                    state = DeepQAState(context=context)
+                    deps = DeepQADeps(
+                        client=self.client, console=Console() if verbose else None
+                    )
+
+                    start_node = DeepQAPlanNode(
+                        provider=Config.QA_PROVIDER,
+                        model=Config.QA_MODEL,
+                    )
+
+                    result = await graph.run(
+                        start_node=start_node, state=state, deps=deps
+                    )
+                    answer = result.output.answer
+                else:
+                    answer = await self.client.ask(question, cite=cite)
+
                 self.console.print(f"[bold blue]Question:[/bold blue] {question}")
                 self.console.print()
                 self.console.print("[bold green]Answer:[/bold green]")

haiku/rag/cli.py

@@ -299,11 +299,21 @@ def ask(
         "--cite",
         help="Include citations in the response",
     ),
+    deep: bool = typer.Option(
+        False,
+        "--deep",
+        help="Use deep multi-agent QA for complex questions",
+    ),
+    verbose: bool = typer.Option(
+        False,
+        "--verbose",
+        help="Show verbose progress output (only with --deep)",
+    ),
 ):
     from haiku.rag.app import HaikuRAGApp
 
     app = HaikuRAGApp(db_path=db)
-    asyncio.run(app.ask(question=question, cite=cite))
+    asyncio.run(app.ask(question=question, cite=cite, deep=deep, verbose=verbose))
 
 
 @cli.command("research", help="Run multi-agent research and output a concise report")

haiku/rag/client.py

@@ -46,6 +46,9 @@ class HaikuRAG:
 
     async def __aexit__(self, exc_type, exc_val, exc_tb):  # noqa: ARG002
         """Async context manager exit."""
+        # Wait for any pending vacuum to complete before closing
+        async with self.store._vacuum_lock:
+            pass
         self.close()
         return False
 
@@ -522,19 +525,22 @@ class HaikuRAG:
         merged.append(current)
         return merged
 
-    async def ask(self, question: str, cite: bool = False) -> str:
+    async def ask(
+        self, question: str, cite: bool = False, system_prompt: str | None = None
+    ) -> str:
         """Ask a question using the configured QA agent.
 
         Args:
             question: The question to ask.
             cite: Whether to include citations in the response.
+            system_prompt: Optional custom system prompt for the QA agent.
 
         Returns:
             The generated answer as a string.
         """
         from haiku.rag.qa import get_qa_agent
 
-        qa_agent = get_qa_agent(self, use_citations=cite)
+        qa_agent = get_qa_agent(self, use_citations=cite, system_prompt=system_prompt)
         return await qa_agent.answer(question)
 
     async def rebuild_database(self) -> AsyncGenerator[str, None]:
@@ -617,13 +623,13 @@ class HaikuRAG:
 
         # Final maintenance: centralized vacuum to curb disk usage
         try:
-            self.store.vacuum()
+            await self.store.vacuum()
         except Exception:
             pass
 
     async def vacuum(self) -> None:
         """Optimize and clean up old versions across all tables."""
-        self.store.vacuum()
+        await self.store.vacuum()
 
     def close(self):
         """Close the underlying store connection."""

haiku/rag/config.py

@@ -57,6 +57,11 @@ class AppConfig(BaseModel):
     # and error out when the database does not already exist.
     DISABLE_DB_AUTOCREATE: bool = False
 
+    # Vacuum retention threshold in seconds. Only versions older than this
+    # threshold will be removed during vacuum operations. Default is 60 seconds
+    # to allow concurrent connections to safely use recent versions.
+    VACUUM_RETENTION_SECONDS: int = 60
+
     @field_validator("MONITOR_DIRECTORIES", mode="before")
     @classmethod
     def parse_monitor_directories(cls, v):

haiku/rag/graph/__init__.py

@@ -0,0 +1 @@
+from haiku.rag.graph.models import ResearchPlan, SearchAnswer

haiku/rag/graph/base.py

@@ -0,0 +1,31 @@
+from typing import Protocol, runtime_checkable
+
+from pydantic import BaseModel, Field
+from rich.console import Console
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.graph.models import SearchAnswer
+
+
+@runtime_checkable
+class GraphContext(Protocol):
+    """Protocol for graph context objects."""
+
+    original_question: str
+    sub_questions: list[str]
+    qa_responses: list[SearchAnswer]
+
+    def add_qa_response(self, qa: SearchAnswer) -> None: ...
+
+
+class BaseGraphDeps(BaseModel):
+    """Base dependencies for graph nodes."""
+
+    model_config = {"arbitrary_types_allowed": True}
+
+    client: HaikuRAG = Field(description="RAG client for document operations")
+    console: Console | None = None
+
+    def emit_log(self, message: str) -> None:
+        if self.console:
+            self.console.print(message)

haiku/rag/graph/common.py

@@ -0,0 +1,33 @@
+from typing import Any, Protocol
+
+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
+
+
+class HasEmitLog(Protocol):
+    def emit_log(self, message: str, state: Any = None) -> None: ...
+
+
+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: HasEmitLog, state: Any, message: str) -> None:
+    deps.emit_log(message, state)

haiku/rag/graph/models.py

@@ -0,0 +1,24 @@
+from pydantic import BaseModel, Field
+
+
+class ResearchPlan(BaseModel):
+    main_question: str
+    sub_questions: list[str]
+
+
+class SearchAnswer(BaseModel):
+    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,
+    )

haiku/rag/{research → graph}/nodes/analysis.py

@@ -3,15 +3,13 @@ from dataclasses import dataclass
 from pydantic_ai import Agent
 from pydantic_graph import BaseNode, GraphRunContext
 
+from haiku.rag.graph.common import get_model, log
 from haiku.rag.research.common import (
     format_analysis_for_prompt,
     format_context_for_prompt,
-    get_model,
-    log,
 )
 from haiku.rag.research.dependencies import ResearchDependencies
 from haiku.rag.research.models import EvaluationResult, InsightAnalysis, ResearchReport
-from haiku.rag.research.nodes.synthesize import SynthesizeNode
 from haiku.rag.research.prompts import DECISION_AGENT_PROMPT, INSIGHT_AGENT_PROMPT
 from haiku.rag.research.state import ResearchDeps, ResearchState
 
@@ -89,6 +87,8 @@ class AnalyzeInsightsNode(BaseNode[ResearchState, ResearchDeps, ResearchReport])
             for question in analysis.new_questions:
                 log(deps, state, f"   • {question}")
 
+        from haiku.rag.graph.nodes.analysis import DecisionNode
+
         return DecisionNode(self.provider, self.model)
 
 
@@ -169,7 +169,8 @@ class DecisionNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         status = "[green]Yes[/green]" if output.is_sufficient else "[red]No[/red]"
         log(deps, state, f"   Sufficient: {status}")
 
-        from haiku.rag.research.nodes.search import SearchDispatchNode
+        from haiku.rag.graph.nodes.search import SearchDispatchNode
+        from haiku.rag.graph.nodes.synthesize import SynthesizeNode
 
         if (
             output.is_sufficient

haiku/rag/{research → graph}/nodes/plan.py

@@ -3,11 +3,11 @@ 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.graph.common import get_model, log
+from haiku.rag.graph.models import ResearchPlan
+from haiku.rag.graph.prompts import PLAN_PROMPT
 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.models import ResearchReport
 from haiku.rag.research.state import ResearchDeps, ResearchState
 
 
@@ -67,4 +67,6 @@ class PlanNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         for i, sq in enumerate(state.context.sub_questions, 1):
             log(deps, state, f"      {i}. {sq}")
 
+        from haiku.rag.graph.nodes.search import SearchDispatchNode
+
         return SearchDispatchNode(self.provider, self.model)

haiku/rag/{research → graph}/nodes/search.py

@@ -7,10 +7,11 @@ 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.graph.common import get_model, log
+from haiku.rag.graph.models import SearchAnswer
+from haiku.rag.graph.prompts import SEARCH_AGENT_PROMPT
 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.models import ResearchReport
 from haiku.rag.research.state import ResearchDeps, ResearchState
 
 
@@ -25,7 +26,7 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
         state = ctx.state
         deps = ctx.deps
         if not state.context.sub_questions:
-            from haiku.rag.research.nodes.analysis import AnalyzeInsightsNode
+            from haiku.rag.graph.nodes.analysis import AnalyzeInsightsNode
 
             return AnalyzeInsightsNode(self.provider, self.model)
 

haiku/rag/{research → graph}/nodes/synthesize.py

@@ -3,10 +3,9 @@ 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.graph.common import get_model, log
+from haiku.rag.research.common import format_context_for_prompt
+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

haiku/rag/graph/prompts.py

@@ -0,0 +1,45 @@
+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/migration.py

@@ -27,7 +27,7 @@ class SQLiteToLanceDBMigrator:
         self.lancedb_path = lancedb_path
         self.console = Console()
 
-    def migrate(self) -> bool:
+    async def migrate(self) -> bool:
         """Perform the migration."""
         try:
             self.console.print(
@@ -94,7 +94,7 @@ class SQLiteToLanceDBMigrator:
             # Optimize and cleanup using centralized vacuum
             self.console.print("[cyan]Optimizing LanceDB...[/cyan]")
             try:
-                lance_store.vacuum()
+                await lance_store.vacuum()
                 self.console.print("[green]✅ Optimization completed[/green]")
             except Exception as e:
                 self.console.print(
@@ -313,4 +313,4 @@ async def migrate_sqlite_to_lancedb(
         lancedb_path = sqlite_path.parent / (sqlite_path.stem + ".lancedb")
 
     migrator = SQLiteToLanceDBMigrator(sqlite_path, lancedb_path)
-    return migrator.migrate()
+    return await migrator.migrate()

haiku/rag/qa/__init__.py

@@ -3,7 +3,11 @@ from haiku.rag.config import Config
 from haiku.rag.qa.agent import QuestionAnswerAgent
 
 
-def get_qa_agent(client: HaikuRAG, use_citations: bool = False) -> QuestionAnswerAgent:
+def get_qa_agent(
+    client: HaikuRAG,
+    use_citations: bool = False,
+    system_prompt: str | None = None,
+) -> QuestionAnswerAgent:
     provider = Config.QA_PROVIDER
     model_name = Config.QA_MODEL
 
@@ -12,4 +16,5 @@ def get_qa_agent(client: HaikuRAG, use_citations: bool = False) -> QuestionAnswe
         provider=provider,
         model=model_name,
         use_citations=use_citations,
+        system_prompt=system_prompt,
     )

haiku/rag/qa/agent.py

@@ -30,18 +30,21 @@ class QuestionAnswerAgent:
         model: str,
         use_citations: bool = False,
         q: float = 0.0,
+        system_prompt: str | None = None,
     ):
         self._client = client
 
-        system_prompt = (
-            QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT
-        )
+        if system_prompt is None:
+            system_prompt = (
+                QA_SYSTEM_PROMPT_WITH_CITATIONS if use_citations else QA_SYSTEM_PROMPT
+            )
         model_obj = self._get_model(provider, model)
 
         self._agent = Agent(
             model=model_obj,
             deps_type=Dependencies,
             system_prompt=system_prompt,
+            retries=3,
         )
 
         @self._agent.tool

haiku/rag/qa/deep/__init__.py

@@ -0,0 +1 @@
+from haiku.rag.qa.deep.models import DeepQAAnswer

haiku/rag/qa/deep/dependencies.py

@@ -0,0 +1,29 @@
+from pydantic import BaseModel, Field
+from rich.console import Console
+
+from haiku.rag.client import HaikuRAG
+from haiku.rag.graph.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")
+    console: Console | None = None

haiku/rag/qa/deep/graph.py

@@ -0,0 +1,21 @@
+from pydantic_graph import Graph
+
+from haiku.rag.qa.deep.models import DeepQAAnswer
+from haiku.rag.qa.deep.nodes import (
+    DeepQADecisionNode,
+    DeepQAPlanNode,
+    DeepQASearchDispatchNode,
+    DeepQASynthesizeNode,
+)
+from haiku.rag.qa.deep.state import DeepQADeps, DeepQAState
+
+
+def build_deep_qa_graph() -> Graph[DeepQAState, DeepQADeps, DeepQAAnswer]:
+    return Graph(
+        nodes=[
+            DeepQAPlanNode,
+            DeepQASearchDispatchNode,
+            DeepQADecisionNode,
+            DeepQASynthesizeNode,
+        ]
+    )

haiku/rag/qa/deep/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,
+    )