haiku.rag 0.11.2__tar.gz → 0.11.4__tar.gz

{haiku_rag-0.11.2 → haiku_rag-0.11.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.11.2
+Version: 0.11.4
 Summary: Agentic Retrieval Augmented Generation (RAG) with LanceDB
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT
@@ -78,6 +78,12 @@ haiku-rag ask "Who is the author of haiku.rag?"
 # Ask questions with citations
 haiku-rag ask "Who is the author of haiku.rag?" --cite
 
+# Deep QA (multi-agent question decomposition)
+haiku-rag ask "Who is the author of haiku.rag?" --deep --cite
+
+# Deep QA with verbose output
+haiku-rag ask "Who is the author of haiku.rag?" --deep --verbose
+
 # Multi‑agent research (iterative plan/search/evaluate)
 haiku-rag research \
   "What are the main drivers and trends of global temperature anomalies since 1990?" \

{haiku_rag-0.11.2 → haiku_rag-0.11.4}/README.md

@@ -40,6 +40,12 @@ haiku-rag ask "Who is the author of haiku.rag?"
 # Ask questions with citations
 haiku-rag ask "Who is the author of haiku.rag?" --cite
 
+# Deep QA (multi-agent question decomposition)
+haiku-rag ask "Who is the author of haiku.rag?" --deep --cite
+
+# Deep QA with verbose output
+haiku-rag ask "Who is the author of haiku.rag?" --deep --verbose
+
 # Multi‑agent research (iterative plan/search/evaluate)
 haiku-rag research \
   "What are the main drivers and trends of global temperature anomalies since 1990?" \

{haiku_rag-0.11.2 → haiku_rag-0.11.4}/pyproject.toml

@@ -2,7 +2,7 @@
 
 name = "haiku.rag"
 description = "Agentic Retrieval Augmented Generation (RAG) with LanceDB"
-version = "0.11.2"
+version = "0.11.4"
 authors = [{ name = "Yiorgis Gozadinos", email = "ggozadinos@gmail.com" }]
 license = { text = "MIT" }
 readme = { file = "README.md", content-type = "text/markdown" }

{haiku_rag-0.11.2 → haiku_rag-0.11.4}/src/evaluations/benchmark.py

@@ -15,6 +15,7 @@ from rich.progress import Progress
 from evaluations.config import DatasetSpec, RetrievalSample
 from evaluations.datasets import DATASETS
 from evaluations.llm_judge import ANSWER_EQUIVALENCE_RUBRIC
+from evaluations.prompts import WIX_SUPPORT_PROMPT
 from haiku.rag import logging  # noqa: F401
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import Config
@@ -61,7 +62,6 @@ async def populate_db(spec: DatasetSpec) -> None:
                     metadata=payload.metadata,
                 )
                 progress.advance(task)
-            rag.store.vacuum()
 
 
 def _is_relevant_match(retrieved_uri: str | None, sample: RetrievalSample) -> bool:
@@ -80,6 +80,11 @@ async def run_retrieval_benchmark(spec: DatasetSpec) -> dict[str, float] | None:
         3: 0.0,
         5: 0.0,
     }
+    success_totals = {
+        1: 0.0,
+        3: 0.0,
+        5: 0.0,
+    }
     total_queries = 0
 
     with Progress() as progress:
@@ -109,15 +114,16 @@ async def run_retrieval_benchmark(spec: DatasetSpec) -> dict[str, float] | None:
                     if retrieved_doc and retrieved_doc.uri:
                         retrieved_uris.append(retrieved_doc.uri)
 
-                # Compute per-query recall@K by counting how many relevant
-                # documents are retrieved within the first K results and
-                # averaging these fractions across all queries.
+                # Compute metrics for each cutoff
                 for cutoff in (1, 3, 5):
                     top_k = set(retrieved_uris[:cutoff])
                     relevant = set(sample.expected_uris)
                     if relevant:
                         matched = len(top_k & relevant)
+                        # Recall: fraction of relevant docs retrieved
                         recall_totals[cutoff] += matched / len(relevant)
+                        # Success: binary - did we get at least one relevant doc?
+                        success_totals[cutoff] += 1.0 if matched > 0 else 0.0
 
                 progress.advance(task)
 
@@ -129,16 +135,28 @@ async def run_retrieval_benchmark(spec: DatasetSpec) -> dict[str, float] | None:
     recall_at_3 = recall_totals[3] / total_queries
     recall_at_5 = recall_totals[5] / total_queries
 
+    success_at_1 = success_totals[1] / total_queries
+    success_at_3 = success_totals[3] / total_queries
+    success_at_5 = success_totals[5] / total_queries
+
     console.print("\n=== Retrieval Benchmark Results ===", style="bold cyan")
     console.print(f"Total queries: {total_queries}")
-    console.print(f"Recall@1: {recall_at_1:.4f}")
-    console.print(f"Recall@3: {recall_at_3:.4f}")
-    console.print(f"Recall@5: {recall_at_5:.4f}")
+    console.print("\nRecall@K (fraction of relevant docs retrieved):")
+    console.print(f"  Recall@1: {recall_at_1:.4f}")
+    console.print(f"  Recall@3: {recall_at_3:.4f}")
+    console.print(f"  Recall@5: {recall_at_5:.4f}")
+    console.print("\nSuccess@K (queries with at least one relevant doc):")
+    console.print(f"  Success@1: {success_at_1:.4f} ({success_at_1 * 100:.1f}%)")
+    console.print(f"  Success@3: {success_at_3:.4f} ({success_at_3 * 100:.1f}%)")
+    console.print(f"  Success@5: {success_at_5:.4f} ({success_at_5 * 100:.1f}%)")
 
     return {
         "recall@1": recall_at_1,
         "recall@3": recall_at_3,
         "recall@5": recall_at_5,
+        "success@1": success_at_1,
+        "success@3": success_at_3,
+        "success@5": success_at_5,
     }
 
 
@@ -187,7 +205,8 @@ async def run_qa_benchmark(
         )
 
         async with HaikuRAG(spec.db_path) as rag:
-            qa = get_qa_agent(rag)
+            system_prompt = WIX_SUPPORT_PROMPT if spec.key == "wix" else None
+            qa = get_qa_agent(rag, system_prompt=system_prompt)
 
             async def answer_question(question: str) -> str:
                 return await qa.answer(question)

haiku_rag-0.11.4/src/evaluations/prompts.py

@@ -0,0 +1,22 @@
+WIX_SUPPORT_PROMPT = """
+You are a WIX technical support expert helping users with questions about the WIX platform.
+
+Your process:
+1. When a user asks a question, use the search_documents tool to find relevant information
+2. Search with specific keywords and phrases from the user's question
+3. Review the search results and their relevance scores
+4. If you need additional context, perform follow-up searches with different keywords
+5. Provide a short and to the point comprehensive answer based only on the retrieved documents
+
+Guidelines:
+- Base your answers strictly on the provided document content
+- Quote or reference specific information when possible
+- If multiple documents contain relevant information, synthesize them coherently
+- Indicate when information is incomplete or when you need to search for additional context
+- If the retrieved documents don't contain sufficient information, clearly state: "I cannot find enough information in the knowledge base to answer this question."
+- For complex questions, consider breaking them down and performing multiple searches
+- Stick to the answer, do not ellaborate or provide context unless explicitly asked for it.
+
+Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
+/no_think
+"""

{haiku_rag-0.11.2 → haiku_rag-0.11.4}/src/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-0.11.2 → haiku_rag-0.11.4}/src/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-0.11.2 → haiku_rag-0.11.4}/src/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-0.11.2 → haiku_rag-0.11.4}/src/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-0.11.4/src/haiku/rag/graph/__init__.py

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

haiku_rag-0.11.4/src/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-0.11.4/src/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-0.11.4/src/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-0.11.2/src/haiku/rag/research → haiku_rag-0.11.4/src/haiku/rag/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-0.11.2/src/haiku/rag/research → haiku_rag-0.11.4/src/haiku/rag/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-0.11.2/src/haiku/rag/research → haiku_rag-0.11.4/src/haiku/rag/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-0.11.2/src/haiku/rag/research → haiku_rag-0.11.4/src/haiku/rag/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-0.11.4/src/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-0.11.2 → haiku_rag-0.11.4}/src/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-0.11.2 → haiku_rag-0.11.4}/src/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-0.11.2 → haiku_rag-0.11.4}/src/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-0.11.4/src/haiku/rag/qa/deep/__init__.py

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

haiku_rag-0.11.4/src/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-0.11.4/src/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,
+        ]
+    )