haiku.rag 0.9.2__py3-none-any.whl → 0.10.0__py3-none-any.whl

haiku/rag/store/upgrades/__init__.py

@@ -1 +1,60 @@
-upgrades = []
+import logging
+from collections.abc import Callable
+from dataclasses import dataclass
+
+from packaging.version import Version, parse
+
+from haiku.rag.store.engine import Store
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class Upgrade:
+    """Represents a database upgrade step."""
+
+    version: str
+    apply: Callable[[Store], None]
+    description: str = ""
+
+
+# Registry of upgrade steps (ordered by version)
+upgrades: list[Upgrade] = []
+
+
+def run_pending_upgrades(store: Store, from_version: str, to_version: str) -> None:
+    """Run upgrades where from_version < step.version <= to_version."""
+    v_from: Version = parse(from_version)
+    v_to: Version = parse(to_version)
+
+    # Ensure that tests/development run available code upgrades even if the
+    # installed package version hasn't been bumped to include them yet.
+    if upgrades:
+        highest_step_version: Version = max(parse(u.version) for u in upgrades)
+        if highest_step_version > v_to:
+            v_to = highest_step_version
+
+    # Determine applicable steps
+    sorted_steps = sorted(upgrades, key=lambda u: parse(u.version))
+    applicable = [s for s in sorted_steps if v_from < parse(s.version) <= v_to]
+    if applicable:
+        logger.info("%d upgrade step(s) pending", len(applicable))
+
+    # Apply in ascending order
+    for idx, step in enumerate(applicable, start=1):
+        logger.info(
+            "Applying upgrade %s: %s (%d/%d)",
+            step.version,
+            step.description or "",
+            idx,
+            len(applicable),
+        )
+        step.apply(store)
+        logger.info("Completed upgrade %s", step.version)
+
+
+from .v0_9_3 import upgrade_fts_phrase as upgrade_0_9_3_fts  # noqa: E402
+from .v0_9_3 import upgrade_order as upgrade_0_9_3_order  # noqa: E402
+
+upgrades.append(upgrade_0_9_3_order)
+upgrades.append(upgrade_0_9_3_fts)

haiku/rag/store/upgrades/v0_9_3.py

@@ -0,0 +1,112 @@
+import json
+
+from lancedb.pydantic import LanceModel, Vector
+from pydantic import Field
+
+from haiku.rag.store.engine import Store
+from haiku.rag.store.upgrades import Upgrade
+
+
+def _infer_vector_dim(store: Store) -> int:
+    """Infer vector dimension from existing data; fallback to embedder config."""
+    try:
+        arrow = store.chunks_table.search().limit(1).to_arrow()
+        rows = arrow.to_pylist()
+        if rows:
+            vec = rows[0].get("vector")
+            if isinstance(vec, list) and vec:
+                return len(vec)
+    except Exception:
+        pass
+    # Fallback to configured embedder vector dim
+    return getattr(store.embedder, "_vector_dim", 1024)
+
+
+def _apply_chunk_order(store: Store) -> None:
+    """Add integer 'order' column to chunks and backfill from metadata."""
+
+    vector_dim = _infer_vector_dim(store)
+
+    class ChunkRecordV2(LanceModel):
+        id: str
+        document_id: str
+        content: str
+        metadata: str = Field(default="{}")
+        order: int = Field(default=0)
+        vector: Vector(vector_dim) = Field(  # type: ignore
+            default_factory=lambda: [0.0] * vector_dim
+        )
+
+    # Read existing chunks
+    try:
+        chunks_arrow = store.chunks_table.search().to_arrow()
+        rows = chunks_arrow.to_pylist()
+    except Exception:
+        rows = []
+
+    new_chunk_records: list[ChunkRecordV2] = []
+    for row in rows:
+        md_raw = row.get("metadata") or "{}"
+        try:
+            md = json.loads(md_raw) if isinstance(md_raw, str) else md_raw
+        except Exception:
+            md = {}
+        # Extract and normalize order
+        order_val = 0
+        try:
+            if isinstance(md, dict) and "order" in md:
+                order_val = int(md["order"])  # type: ignore[arg-type]
+        except Exception:
+            order_val = 0
+
+        if isinstance(md, dict) and "order" in md:
+            md = {k: v for k, v in md.items() if k != "order"}
+
+        vec = row.get("vector") or [0.0] * vector_dim
+
+        new_chunk_records.append(
+            ChunkRecordV2(
+                id=row.get("id"),
+                document_id=row.get("document_id"),
+                content=row.get("content", ""),
+                metadata=json.dumps(md),
+                order=order_val,
+                vector=vec,
+            )
+        )
+
+    # Recreate chunks table with new schema
+    try:
+        store.db.drop_table("chunks")
+    except Exception:
+        pass
+
+    store.chunks_table = store.db.create_table("chunks", schema=ChunkRecordV2)
+    store.chunks_table.create_fts_index("content", replace=True)
+
+    if new_chunk_records:
+        store.chunks_table.add(new_chunk_records)
+
+
+upgrade_order = Upgrade(
+    version="0.9.3",
+    apply=_apply_chunk_order,
+    description="Add 'order' column to chunks and backfill from metadata",
+)
+
+
+def _apply_fts_phrase_support(store: Store) -> None:
+    """Recreate FTS index with phrase query support and no stop-word removal."""
+    try:
+        store.chunks_table.create_fts_index(
+            "content", replace=True, with_position=True, remove_stop_words=False
+        )
+    except Exception:
+        pass
+
+
+upgrade_fts_phrase = Upgrade(
+    version="0.9.3",
+    apply=_apply_fts_phrase_support,
+    description="Enable FTS phrase queries (with positions) and keep stop-words",
+)

{haiku_rag-0.9.2.dist-info → haiku_rag-0.10.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haiku.rag
-Version: 0.9.2
+Version: 0.10.0
 Summary: Agentic Retrieval Augmented Generation (RAG) with LanceDB
 Author-email: Yiorgis Gozadinos <ggozadinos@gmail.com>
 License: MIT
@@ -23,6 +23,7 @@ Requires-Dist: fastmcp>=2.12.3
 Requires-Dist: httpx>=0.28.1
 Requires-Dist: lancedb>=0.25.0
 Requires-Dist: pydantic-ai>=1.0.8
+Requires-Dist: pydantic-graph>=1.0.8
 Requires-Dist: pydantic>=2.11.9
 Requires-Dist: python-dotenv>=1.1.1
 Requires-Dist: rich>=14.1.0
@@ -48,6 +49,7 @@ Retrieval-Augmented Generation (RAG) library built on LanceDB.
 - **Local LanceDB**: No external servers required, supports also LanceDB cloud storage, S3, Google Cloud & Azure
 - **Multiple embedding providers**: Ollama, VoyageAI, OpenAI, vLLM
 - **Multiple QA providers**: Any provider/model supported by Pydantic AI
+- **Research graph (multi‑agent)**: Plan → Search → Evaluate → Synthesize with agentic AI
 - **Native hybrid search**: Vector + full-text search with native LanceDB RRF reranking
 - **Reranking**: Default search result reranking with MixedBread AI, Cohere, or vLLM
 - **Question answering**: Built-in QA agents on your documents
@@ -75,6 +77,14 @@ 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
 
+# Multi‑agent research (iterative plan/search/evaluate)
+haiku-rag research \
+  "What are the main drivers and trends of global temperature anomalies since 1990?" \
+  --max-iterations 2 \
+  --confidence-threshold 0.8 \
+  --max-concurrency 3 \
+  --verbose
+
 # Rebuild database (re-chunk and re-embed all documents)
 haiku-rag rebuild
 
@@ -90,6 +100,13 @@ haiku-rag serve
 
 ```python
 from haiku.rag.client import HaikuRAG
+from haiku.rag.research import (
+    ResearchContext,
+    ResearchDeps,
+    ResearchState,
+    build_research_graph,
+    PlanNode,
+)
 
 async with HaikuRAG("database.lancedb") as client:
     # Add document
@@ -107,6 +124,25 @@ async with HaikuRAG("database.lancedb") as client:
     # Ask questions with citations
     answer = await client.ask("Who is the author of haiku.rag?", cite=True)
     print(answer)
+
+    # Multi‑agent research pipeline (Plan → Search → Evaluate → Synthesize)
+    graph = build_research_graph()
+    state = ResearchState(
+        question=(
+            "What are the main drivers and trends of global temperature "
+            "anomalies since 1990?"
+        ),
+        context=ResearchContext(original_question="…"),
+        max_iterations=2,
+        confidence_threshold=0.8,
+        max_concurrency=3,
+    )
+    deps = ResearchDeps(client=client)
+    start = PlanNode(provider=None, model=None)
+    result = await graph.run(start, state=state, deps=deps)
+    report = result.output
+    print(report.title)
+    print(report.executive_summary)
 ```
 
 ## MCP Server

haiku_rag-0.10.0.dist-info/RECORD

@@ -0,0 +1,53 @@
+haiku/rag/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+haiku/rag/app.py,sha256=m5agkPrJhbzEbdC01CU_GR2Gj4voFuAGmxR7DS2K9is,12934
+haiku/rag/chunker.py,sha256=PVe6ysv8UlacUd4Zb3_8RFWIaWDXnzBAy2VDJ4TaUsE,1555
+haiku/rag/cli.py,sha256=oXEQoRTlzrrJ9hC27_Dht9ElBb9q_wTEESnXdNy3eW8,10257
+haiku/rag/client.py,sha256=QgJQu7g7JjAzWN6R10NeDqpFf89Dml_LiWce4QRHLHc,21177
+haiku/rag/config.py,sha256=SPEIv2IElZmZh4Wsp8gk7ViRW5ZzD-UGmIqRAXscDdI,2134
+haiku/rag/logging.py,sha256=dm65AwADpcQsH5OAPtRA-4hsw0w5DK-sGOvzYkj6jzw,1720
+haiku/rag/mcp.py,sha256=bR9Y-Nz-hvjiql20Y0KE0hwNGwyjmPGX8K9d-qmXptY,4683
+haiku/rag/migration.py,sha256=M--KnSF3lxgKjxmokb4vuzGH-pV8eg0C_8e7jvPqW8Y,11058
+haiku/rag/monitor.py,sha256=r386nkhdlsU8UECwIuVwnrSlgMk3vNIuUZGNIzkZuec,2770
+haiku/rag/reader.py,sha256=qkPTMJuQ_o4sK-8zpDl9WFYe_MJ7aL_gUw6rczIpW-g,3274
+haiku/rag/utils.py,sha256=aiuPu_rrfpyIvJJq0o5boUIIvCdNzdpKwAIPYYn3iG8,4965
+haiku/rag/embeddings/__init__.py,sha256=44IfDITGIFTflGT6UEmiYOwpWFVbYv5smLY59D0YeCs,1419
+haiku/rag/embeddings/base.py,sha256=BnSviKrlzjv3L0sZJs_T-pxfawd-bcTak-rsX-D2f3A,497
+haiku/rag/embeddings/ollama.py,sha256=LuLlHH6RGoO9_gFCIlbmesuXOj017gTw6z-p8Ez0CfE,595
+haiku/rag/embeddings/openai.py,sha256=fIFCk-jpUtaW0xsnrQnJ824O0UCjaGG2sgvBzREhilc,503
+haiku/rag/embeddings/vllm.py,sha256=vhaUnCn6VMkfSluLhWKtSV-sekFaPsp4pKo2N7-SBCY,626
+haiku/rag/embeddings/voyageai.py,sha256=UW-MW4tJKnPB6Fs2P7A3yt-ZeRm46H9npckchSriPX8,661
+haiku/rag/qa/__init__.py,sha256=Sl7Kzrg9CuBOcMF01wc1NtQhUNWjJI0MhIHfCWrb8V4,434
+haiku/rag/qa/agent.py,sha256=f4Keh-ESgctNbTg96QL95HYjINVLOcxa8t8crx92MMk,3081
+haiku/rag/qa/prompts.py,sha256=LhRfDtO8Pb06lpr4PpwEaKUYItZ5OiIkeqcCogcssHY,3347
+haiku/rag/reranking/__init__.py,sha256=IRXHs4qPu6VbGJQpzSwhgtVWWumURH_vEoVFE-extlo,894
+haiku/rag/reranking/base.py,sha256=LM9yUSSJ414UgBZhFTgxGprlRqzfTe4I1vgjricz2JY,405
+haiku/rag/reranking/cohere.py,sha256=1iTdiaa8vvb6oHVB2qpWzUOVkyfUcimVSZp6Qr4aq4c,1049
+haiku/rag/reranking/mxbai.py,sha256=uveGFIdmNmepd2EQsvYr64wv0ra2_wB845hdSZXy5Cw,908
+haiku/rag/reranking/vllm.py,sha256=xVGH9ss-ISWdJ5SKUUHUbTqBo7PIEmA_SQv0ScdJ6XA,1479
+haiku/rag/research/__init__.py,sha256=t4JAmIXcKaWqvpFGX5yaehsNrfblskEMn-4mDmdKn9c,502
+haiku/rag/research/common.py,sha256=EUnsA6VZ3-WMweXESuUYezH1ALit8N38064bsZFqtBE,1688
+haiku/rag/research/dependencies.py,sha256=ZiSQdV6jHti4DuUp4WCaJL73TqYDr5vC8ppB34M2cNg,1639
+haiku/rag/research/graph.py,sha256=m3vDP1nPXWzfS7VeTQzmTOk-lFpoaTvKHvRIF2mbxvs,798
+haiku/rag/research/models.py,sha256=klE2qGF5fom5gJRQzQUbnoGYaXusNKeJ9veeXoYDD5Q,2308
+haiku/rag/research/prompts.py,sha256=v_DZNaKk88CDEF8qt9c-puO6QF-NyBQKnl_mO1pMauY,5013
+haiku/rag/research/state.py,sha256=vFwO8c2JmwwfkELE5Mwjt9Oat-bHn5tayf31MIG2SRs,623
+haiku/rag/research/nodes/evaluate.py,sha256=Cp2J-jXYZothiQV3zRZFaCsBLaUU0Tm_-ri-hlgQQII,2897
+haiku/rag/research/nodes/plan.py,sha256=9AkTls01Q3zTLKGgIgSCX9X4VYC8IWjEWii8A_f77YQ,2439
+haiku/rag/research/nodes/search.py,sha256=lHgDCCL7hQdpQeMK-HVzsF_hH_pIv44xxSIiv1JuvYo,3513
+haiku/rag/research/nodes/synthesize.py,sha256=4acKduqWnE11ML7elUksKLozxzWJTkBLSJ2li_YMxgY,1736
+haiku/rag/store/__init__.py,sha256=hq0W0DAC7ysqhWSP2M2uHX8cbG6kbr-sWHxhq6qQcY0,103
+haiku/rag/store/engine.py,sha256=-3MZJYft2XTWaLuyKha8DKhWQeU5E5CBeskXXF5fXso,9555
+haiku/rag/store/models/__init__.py,sha256=s0E72zneGlowvZrFWaNxHYjOAUjgWdLxzdYsnvNRVlY,88
+haiku/rag/store/models/chunk.py,sha256=Ww_hj3DMwJLNM33l1GvIP84yzDFc6cxfiWcotUfWSYg,383
+haiku/rag/store/models/document.py,sha256=zSSpt6pyrMJAIXGQvIcqojcqUzwZnhp3WxVokaWxNRc,396
+haiku/rag/store/repositories/__init__.py,sha256=Olv5dLfBQINRV3HrsfUpjzkZ7Qm7goEYyMNykgo_DaY,291
+haiku/rag/store/repositories/chunk.py,sha256=O2SEhQy3ZptWjwwpxS-L8KNq2tEqEBqheHfLw-M_FqA,15012
+haiku/rag/store/repositories/document.py,sha256=m11SamQoGYs5ODfmarJGU1yIcqtgmnba-5bGOPQuYrI,7773
+haiku/rag/store/repositories/settings.py,sha256=7XMBMavU8zRgdBoQzQg0Obfa7UKjuVnBugidTC6sEW0,5548
+haiku/rag/store/upgrades/__init__.py,sha256=gDOxiq3wdZPr3JoenjNYxx0cpgZJhbaFKNX2fzXRq1Q,1852
+haiku/rag/store/upgrades/v0_9_3.py,sha256=NrjNilQSgDtFWRbL3ZUtzQzJ8tf9u0dDRJtnDFwwbdw,3322
+haiku_rag-0.10.0.dist-info/METADATA,sha256=QLc8BBJ4WCNEvseyYpWNfkuUfmdxGywD6Jtn0OTsrc0,5879
+haiku_rag-0.10.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+haiku_rag-0.10.0.dist-info/entry_points.txt,sha256=G1U3nAkNd5YDYd4v0tuYFbriz0i-JheCsFuT9kIoGCI,48
+haiku_rag-0.10.0.dist-info/licenses/LICENSE,sha256=eXZrWjSk9PwYFNK9yUczl3oPl95Z4V9UXH7bPN46iPo,1065
+haiku_rag-0.10.0.dist-info/RECORD,,

haiku/rag/research/base.py

@@ -1,130 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any
-
-from pydantic import BaseModel, Field
-from pydantic_ai import Agent
-from pydantic_ai.models.openai import OpenAIChatModel
-from pydantic_ai.output import ToolOutput
-from pydantic_ai.providers.ollama import OllamaProvider
-from pydantic_ai.providers.openai import OpenAIProvider
-from pydantic_ai.run import AgentRunResult
-
-from haiku.rag.config import Config
-
-if TYPE_CHECKING:
-    from haiku.rag.research.dependencies import ResearchDependencies
-
-
-class BaseResearchAgent[T](ABC):
-    """Base class for all research agents."""
-
-    def __init__(
-        self,
-        provider: str,
-        model: str,
-        output_type: type[T],
-    ):
-        self.provider = provider
-        self.model = model
-        self.output_type = output_type
-
-        model_obj = self._get_model(provider, model)
-
-        # Import deps type lazily to avoid circular import during module load
-        from haiku.rag.research.dependencies import ResearchDependencies
-
-        # If the agent is expected to return plain text, pass `str` directly.
-        # Otherwise, wrap the model with ToolOutput for robust tool-handling retries.
-        agent_output_type: Any
-        if self.output_type is str:  # plain text output
-            agent_output_type = str
-        else:
-            agent_output_type = ToolOutput(self.output_type, max_retries=3)
-
-        self._agent = Agent(
-            model=model_obj,
-            deps_type=ResearchDependencies,
-            output_type=agent_output_type,
-            system_prompt=self.get_system_prompt(),
-        )
-
-        # Register tools
-        self.register_tools()
-
-    def _get_model(self, provider: str, model: str):
-        """Get the appropriate model object for the provider."""
-        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:
-            # For all other providers, use the provider:model format
-            return f"{provider}:{model}"
-
-    @abstractmethod
-    def get_system_prompt(self) -> str:
-        """Return the system prompt for this agent."""
-        pass
-
-    @abstractmethod
-    def register_tools(self) -> None:
-        """Register agent-specific tools."""
-        pass
-
-    async def run(
-        self, prompt: str, deps: "ResearchDependencies", **kwargs
-    ) -> AgentRunResult[T]:
-        """Execute the agent."""
-        return await self._agent.run(prompt, deps=deps, **kwargs)
-
-    @property
-    def agent(self) -> Agent[Any, T]:
-        """Access the underlying Pydantic AI agent."""
-        return self._agent
-
-
-class SearchResult(BaseModel):
-    """Standard search result format."""
-
-    content: str
-    score: float
-    document_uri: str
-    metadata: dict[str, Any] = Field(default_factory=dict)
-
-
-class ResearchOutput(BaseModel):
-    """Standard research output format."""
-
-    summary: str
-    detailed_findings: list[str]
-    sources: list[str]
-    confidence: float
-
-
-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 URIs corresponding to the snippets actually used in the"
-            " answer (one URI per snippet; omit if none)"
-        ),
-        default_factory=list,
-    )

haiku/rag/research/evaluation_agent.py

@@ -1,42 +0,0 @@
-from pydantic import BaseModel, Field
-
-from haiku.rag.research.base import BaseResearchAgent
-from haiku.rag.research.prompts import EVALUATION_AGENT_PROMPT
-
-
-class EvaluationResult(BaseModel):
-    """Result of analysis and evaluation."""
-
-    key_insights: list[str] = Field(
-        description="Main insights extracted from the research so far"
-    )
-    new_questions: list[str] = Field(
-        description="New sub-questions to add to the research (max 3)",
-        max_length=3,
-        default=[],
-    )
-    confidence_score: float = Field(
-        description="Confidence level in the completeness of research (0-1)",
-        ge=0.0,
-        le=1.0,
-    )
-    is_sufficient: bool = Field(
-        description="Whether the research is sufficient to answer the original question"
-    )
-    reasoning: str = Field(
-        description="Explanation of why the research is or isn't complete"
-    )
-
-
-class AnalysisEvaluationAgent(BaseResearchAgent[EvaluationResult]):
-    """Agent that analyzes findings and evaluates research completeness."""
-
-    def __init__(self, provider: str, model: str) -> None:
-        super().__init__(provider, model, output_type=EvaluationResult)
-
-    def get_system_prompt(self) -> str:
-        return EVALUATION_AGENT_PROMPT
-
-    def register_tools(self) -> None:
-        """No additional tools needed - uses LLM capabilities directly."""
-        pass