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

haiku/rag/qa/prompts.py

@@ -1,60 +1,38 @@
-QA_SYSTEM_PROMPT = """
-You are a knowledgeable assistant that helps users find information from a document knowledge base.
-
-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
-"""
-
-QA_SYSTEM_PROMPT_WITH_CITATIONS = """
-You are a knowledgeable assistant that helps users find information from a document knowledge base.
-
-IMPORTANT: You MUST use the search_documents tool for every question. Do not answer any question without first searching the knowledge base.
-
-Your process:
-1. IMMEDIATELY call the search_documents tool with relevant keywords from the user's question
-2. Review the search results and their relevance scores
-3. If you need additional context, perform follow-up searches with different keywords
-4. Provide a short and to the point comprehensive answer based only on the retrieved documents
-5. Always include citations for the sources used in your answer
+QA_SYSTEM_PROMPT = """You are a knowledgeable assistant that answers questions using a document knowledge base.
+
+Process:
+1. Call search_documents with relevant keywords from the question
+2. Review the results and their relevance scores
+3. If needed, perform follow-up searches with different keywords (max 3 total)
+4. Provide a concise answer based strictly on the retrieved content
+
+The search tool returns results like:
+[chunk_abc123] (score: 0.85)
+Source: "Document Title" > Section > Subsection
+Type: paragraph
+Content:
+The actual text content here...
+
+[chunk_def456] (score: 0.72)
+Source: "Another Document"
+Type: table
+Content:
+| Column 1 | Column 2 |
+...
+
+Each result includes:
+- chunk_id in brackets and relevance score
+- Source: document title and section hierarchy (when available)
+- Type: content type like paragraph, table, code, list_item (when available)
+- Content: the actual text
+
+In your response, include the chunk IDs you used in cited_chunks.
 
 Guidelines:
-- Base your answers strictly on the provided document content
-- 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.
-- ALWAYS include citations at the end of your response using the format below
-
-Citation Format:
-After your answer, include a "Citations:" section that lists:
-- The document title (if available) or URI from each search result used
-- A brief excerpt (first 50-100 characters) of the content that supported your answer
-- Format: "Citations:\n- [document title or URI]: [content_excerpt]..."
-
-Example response format:
-[Your answer here]
-
-Citations:
-- /path/to/document1.pdf: "This document explains that AFMAN stands for Air Force Manual..."
-- /path/to/document2.pdf: "The manual provides guidance on military procedures and..."
-
-Be concise, and always maintain accuracy over completeness. Prefer short, direct answers that are well-supported by the documents.
-/no_think
+- Base answers strictly on retrieved content - do not use external knowledge
+- Use the Source and Type metadata to understand context
+- If multiple results are relevant, synthesize them coherently
+- If information is insufficient, say: "I cannot find enough information in the knowledge base to answer this question."
+- Be concise and direct - avoid elaboration unless asked
+- Higher scores indicate more relevant results
 """

haiku/rag/reranking/__init__.py

@@ -24,7 +24,7 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
 
     reranker: RerankerBase | None = None
 
-    if config.reranking.provider == "mxbai":
+    if config.reranking.model and config.reranking.model.provider == "mxbai":
         try:
             from haiku.rag.reranking.mxbai import MxBAIReranker
 
@@ -33,7 +33,7 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
         except ImportError:
             reranker = None
 
-    elif config.reranking.provider == "cohere":
+    elif config.reranking.model and config.reranking.model.provider == "cohere":
         try:
             from haiku.rag.reranking.cohere import CohereReranker
 
@@ -41,20 +41,23 @@ def get_reranker(config: AppConfig = Config) -> RerankerBase | None:
         except ImportError:
             reranker = None
 
-    elif config.reranking.provider == "vllm":
+    elif config.reranking.model and config.reranking.model.provider == "vllm":
         try:
             from haiku.rag.reranking.vllm import VLLMReranker
 
-            reranker = VLLMReranker(config.reranking.model)
+            base_url = config.reranking.model.base_url
+            if not base_url:
+                raise ValueError("vLLM reranker requires base_url in reranking.model")
+            reranker = VLLMReranker(config.reranking.model.name, base_url)
         except ImportError:
             reranker = None
 
-    elif config.reranking.provider == "zeroentropy":
+    elif config.reranking.model and config.reranking.model.provider == "zeroentropy":
         try:
             from haiku.rag.reranking.zeroentropy import ZeroEntropyReranker
 
             # Use configured model or default to zerank-1
-            model = config.reranking.model or "zerank-1"
+            model = config.reranking.model.name or "zerank-1"
             reranker = ZeroEntropyReranker(model)
         except ImportError:
             reranker = None

haiku/rag/reranking/base.py

@@ -3,7 +3,7 @@ from haiku.rag.store.models.chunk import Chunk
 
 
 class RerankerBase:
-    _model: str = Config.reranking.model
+    _model: str | None = Config.reranking.model.name if Config.reranking.model else None
 
     async def rerank(
         self, query: str, chunks: list[Chunk], top_n: int = 10

haiku/rag/reranking/cohere.py

@@ -9,10 +9,10 @@ except ImportError as e:
     ) from e
 
 
-class CohereReranker(RerankerBase):
+class CohereReranker(RerankerBase):  # pragma: no cover
     def __init__(self):
         # Cohere SDK reads CO_API_KEY from environment by default
-        self._client = cohere.ClientV2()
+        self._client = cohere.AsyncClientV2()
 
     async def rerank(
         self, query: str, chunks: list[Chunk], top_n: int = 10
@@ -22,8 +22,9 @@ class CohereReranker(RerankerBase):
 
         documents = [chunk.content for chunk in chunks]
 
-        response = self._client.rerank(
-            model=self._model, query=query, documents=documents, top_n=top_n
+        model_name = self._model or "rerank-v3.5"
+        response = await self._client.rerank(
+            model=model_name, query=query, documents=documents, top_n=top_n
         )
 
         reranked_chunks = []

haiku/rag/reranking/mxbai.py

@@ -7,9 +7,12 @@ from haiku.rag.store.models.chunk import Chunk
 
 class MxBAIReranker(RerankerBase):
     def __init__(self):
-        self._client = MxbaiRerankV2(
-            Config.reranking.model, disable_transformers_warnings=True
+        model_name = (
+            Config.reranking.model.name
+            if Config.reranking.model
+            else "mixedbread-ai/mxbai-rerank-base-v2"
         )
+        self._client = MxbaiRerankV2(model_name, disable_transformers_warnings=True)
 
     async def rerank(
         self, query: str, chunks: list[Chunk], top_n: int = 10

haiku/rag/reranking/vllm.py

@@ -1,14 +1,13 @@
 import httpx
 
-from haiku.rag.config import Config
 from haiku.rag.reranking.base import RerankerBase
 from haiku.rag.store.models.chunk import Chunk
 
 
-class VLLMReranker(RerankerBase):
-    def __init__(self, model: str):
+class VLLMReranker(RerankerBase):  # pragma: no cover
+    def __init__(self, model: str, base_url: str):
         self._model = model
-        self._base_url = Config.providers.vllm.rerank_base_url
+        self._base_url = base_url
 
     async def rerank(
         self, query: str, chunks: list[Chunk], top_n: int = 10

haiku/rag/reranking/zeroentropy.py

@@ -1,10 +1,10 @@
-from zeroentropy import ZeroEntropy
+from zeroentropy import AsyncZeroEntropy
 
 from haiku.rag.reranking.base import RerankerBase
 from haiku.rag.store.models.chunk import Chunk
 
 
-class ZeroEntropyReranker(RerankerBase):
+class ZeroEntropyReranker(RerankerBase):  # pragma: no cover
     """Zero Entropy reranker implementation using the zerank-1 model."""
 
     def __init__(self, model: str = "zerank-1"):
@@ -15,7 +15,7 @@ class ZeroEntropyReranker(RerankerBase):
         """
         self._model = model
         # Zero Entropy SDK reads ZEROENTROPY_API_KEY from environment by default
-        self._client = ZeroEntropy()
+        self._client = AsyncZeroEntropy()
 
     async def rerank(
         self, query: str, chunks: list[Chunk], top_n: int = 10
@@ -37,8 +37,9 @@ class ZeroEntropyReranker(RerankerBase):
         documents = [chunk.content for chunk in chunks]
 
         # Call Zero Entropy reranking API
-        response = self._client.models.rerank(
-            model=self._model,
+        model_name = self._model or "zerank-1"
+        response = await self._client.models.rerank(
+            model=model_name,
             query=query,
             documents=documents,
         )

haiku/rag/store/__init__.py

@@ -1,4 +1,5 @@
 from .engine import Store
+from .exceptions import ReadOnlyError
 from .models import Chunk, Document
 
-__all__ = ["Store", "Chunk", "Document"]
+__all__ = ["Store", "Chunk", "Document", "ReadOnlyError"]

haiku/rag/store/engine.py

@@ -1,9 +1,10 @@
 import asyncio
 import json
 import logging
-from datetime import timedelta
+from datetime import datetime, timedelta
 from importlib import metadata
 from pathlib import Path
+from typing import Any
 from uuid import uuid4
 
 import lancedb
@@ -12,6 +13,7 @@ from pydantic import Field
 
 from haiku.rag.config import AppConfig, Config
 from haiku.rag.embeddings import get_embedder
+from haiku.rag.store.exceptions import ReadOnlyError
 
 logger = logging.getLogger(__name__)
 
@@ -22,6 +24,8 @@ class DocumentRecord(LanceModel):
     uri: str | None = None
     title: str | None = None
     metadata: str = Field(default="{}")
+    docling_document_json: str | None = None
+    docling_version: str | None = None
     created_at: str = Field(default_factory=lambda: "")
     updated_at: str = Field(default_factory=lambda: "")
 
@@ -36,6 +40,7 @@ def create_chunk_model(vector_dim: int):
         id: str = Field(default_factory=lambda: str(uuid4()))
         document_id: str
         content: str
+        content_fts: str = Field(default="")
         metadata: str = Field(default="{}")
         order: int = Field(default=0)
         vector: Vector(vector_dim) = Field(default_factory=lambda: [0.0] * vector_dim)  # type: ignore
@@ -54,39 +59,67 @@ class Store:
         db_path: Path,
         config: AppConfig = Config,
         skip_validation: bool = False,
-        allow_create: bool = True,
+        create: bool = False,
+        read_only: bool = False,
+        before: datetime | None = None,
     ):
         self.db_path: Path = db_path
         self._config = config
+        self._before = before
+        # Time-travel mode is always read-only
+        self._read_only = read_only or (before is not None)
         self.embedder = get_embedder(config=self._config)
         self._vacuum_lock = asyncio.Lock()
 
         # Create the ChunkRecord model with the correct vector dimension
         self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
 
-        # Local filesystem handling for DB directory
+        # Check if database exists (for local filesystem only)
+        is_new_db = False
         if not self._has_cloud_config():
-            if not allow_create:
-                # Read operations should not create the database
-                if not db_path.exists():
+            if not db_path.exists():
+                if not create:
                     raise FileNotFoundError(
-                        f"Database does not exist: {db_path}. Use a write operation (add, add-src) to create it."
+                        f"Database does not exist at {self.db_path.absolute()}. "
+                        "Use 'haiku-rag init' to create a new database."
                     )
-            else:
-                # Write operations - ensure parent directories exist
+                is_new_db = True
+                # Ensure parent directories exist for new databases
                 if not db_path.parent.exists():
                     Path.mkdir(db_path.parent, parents=True)
 
         # Connect to LanceDB
         self.db = self._connect_to_lancedb(db_path)
 
-        # Initialize tables
-        self.create_or_update_db()
+        # Initialize tables (creates them if they don't exist)
+        self._init_tables()
+
+        # Checkout tables to historical state if before is specified
+        if before is not None:
+            self._checkout_tables_before(before)
+
+        # Run upgrades only on existing databases, set version for new ones
+        # Skip upgrades in read-only mode (they would fail anyway)
+        if not self._read_only:
+            if is_new_db:
+                self._set_initial_version()
+            else:
+                self._run_upgrades()
 
         # Validate config compatibility after connection is established
         if not skip_validation:
             self._validate_configuration()
 
+    @property
+    def is_read_only(self) -> bool:
+        """Whether the store is in read-only mode."""
+        return self._read_only
+
+    def _assert_writable(self) -> None:
+        """Raise ReadOnlyError if the store is in read-only mode."""
+        if self._read_only:
+            raise ReadOnlyError("Cannot modify database in read-only mode")
+
     async def vacuum(self, retention_seconds: int | None = None) -> None:
         """Optimize and clean up old versions across all tables to reduce disk usage.
 
@@ -97,7 +130,12 @@ class Store:
         Note:
             If vacuum is already running, this method returns immediately without blocking.
             Use asyncio.create_task(store.vacuum()) for non-blocking background execution.
+
+        Raises:
+            ReadOnlyError: If the store is in read-only mode.
         """
+        self._assert_writable()
+
         if self._has_cloud_config() and str(self._config.lancedb.uri).startswith(
             "db://"
         ):
@@ -145,6 +183,87 @@ class Store:
             and self._config.lancedb.region
         )
 
+    def get_stats(self) -> dict:
+        """Get comprehensive table statistics.
+
+        Returns:
+            Dictionary with statistics for documents and chunks tables including:
+            - Row counts
+            - Storage sizes
+            - Vector index status and statistics
+        """
+        stats_dict: dict = {
+            "documents": {"exists": False},
+            "chunks": {"exists": False},
+        }
+
+        # Documents table stats
+        doc_stats: dict = self.documents_table.stats()  # type: ignore[assignment]
+        stats_dict["documents"] = {
+            "exists": True,
+            "num_rows": doc_stats.get("num_rows", 0),
+            "total_bytes": doc_stats.get("total_bytes", 0),
+        }
+
+        # Chunks table stats
+        chunk_stats: dict = self.chunks_table.stats()  # type: ignore[assignment]
+        stats_dict["chunks"] = {
+            "exists": True,
+            "num_rows": chunk_stats.get("num_rows", 0),
+            "total_bytes": chunk_stats.get("total_bytes", 0),
+        }
+
+        # Vector index stats
+        indices = self.chunks_table.list_indices()
+        has_vector_index = any("vector" in str(idx).lower() for idx in indices)
+        stats_dict["chunks"]["has_vector_index"] = has_vector_index
+
+        if has_vector_index:
+            index_stats = self.chunks_table.index_stats("vector_idx")
+            if index_stats is not None:
+                stats_dict["chunks"]["num_indexed_rows"] = index_stats.num_indexed_rows
+                stats_dict["chunks"]["num_unindexed_rows"] = (
+                    index_stats.num_unindexed_rows
+                )
+
+        return stats_dict
+
+    def _ensure_vector_index(self) -> None:
+        """Create or rebuild vector index on chunks table.
+
+        Cloud deployments auto-create indexes, so we skip for those.
+        For self-hosted, creates an IVF_PQ index. If an index exists,
+        it will be replaced (using replace=True parameter).
+        Note: Index creation requires sufficient training data.
+        """
+        if self._has_cloud_config():
+            return
+
+        try:
+            # Check if table has enough data (indexes require training data)
+            row_count = self.chunks_table.count_rows()
+            if row_count < 256:
+                logger.debug(
+                    f"Skipping vector index creation: need at least 256 rows, have {row_count}"
+                )
+                return
+
+            # Create or replace index (replace=True is the default)
+            logger.info("Creating vector index on chunks table...")
+            self.chunks_table.create_index(
+                metric=self._config.search.vector_index_metric,
+                index_type="IVF_PQ",
+                replace=True,  # Explicit: replace existing index
+            )
+
+            # Wait for index creation to complete
+            # Index name is column_name + "_idx"
+            self.chunks_table.wait_for_index(["vector_idx"], timeout=timedelta(hours=1))
+
+            logger.info("Vector index created successfully")
+        except Exception as e:
+            logger.warning(f"Could not create vector index: {e}")
+
     def _validate_configuration(self) -> None:
         """Validate that the configuration is compatible with the database."""
         from haiku.rag.store.repositories.settings import SettingsRepository
@@ -152,9 +271,8 @@ class Store:
         settings_repo = SettingsRepository(self)
         settings_repo.validate_config_compatibility()
 
-    def create_or_update_db(self):
-        """Create the database tables."""
-
+    def _init_tables(self):
+        """Initialize database tables (create if they don't exist)."""
         # Get list of existing tables
         existing_tables = self.db.table_names()
 
@@ -171,9 +289,9 @@ class Store:
             self.chunks_table = self.db.open_table("chunks")
         else:
             self.chunks_table = self.db.create_table("chunks", schema=self.ChunkRecord)
-            # Create FTS index on the new table with phrase query support
+            # Create FTS index on content_fts (contextualized content) for better search
             self.chunks_table.create_fts_index(
-                "content", replace=True, with_position=True, remove_stop_words=False
+                "content_fts", replace=True, with_position=True, remove_stop_words=False
             )
 
         # Create or get settings table
@@ -189,34 +307,21 @@ class Store:
                 [SettingsRecord(id="settings", settings=json.dumps(settings_data))]
             )
 
-        # Run pending upgrades based on stored version and package version
+    def _set_initial_version(self):
+        """Set the initial version for a new database."""
+        self.set_haiku_version(metadata.version("haiku.rag-slim"))
+
+    def _run_upgrades(self):
+        """Run pending database upgrades."""
         try:
             from haiku.rag.store.upgrades import run_pending_upgrades
 
             current_version = metadata.version("haiku.rag-slim")
             db_version = self.get_haiku_version()
 
-            if db_version != "0.0.0":
-                run_pending_upgrades(self, db_version, current_version)
-
-            # After upgrades complete (or if none), set stored version
-            # to the greater of the installed package version and the
-            # highest available upgrade step version in code.
-            try:
-                from packaging.version import parse as _v
-
-                from haiku.rag.store.upgrades import upgrades as _steps
-
-                highest_step = max((_v(u.version) for u in _steps), default=None)
-                effective_version = (
-                    str(max(_v(current_version), highest_step))
-                    if highest_step is not None
-                    else current_version
-                )
-            except Exception:
-                effective_version = current_version
+            run_pending_upgrades(self, db_version, current_version)
 
-            self.set_haiku_version(effective_version)
+            self.set_haiku_version(current_version)
         except Exception as e:
             # Avoid hard failure on initial connection; log and continue so CLI remains usable.
             logger.warning(
@@ -241,7 +346,12 @@ class Store:
         return "0.0.0"
 
     def set_haiku_version(self, version: str) -> None:
-        """Updates the user version in settings."""
+        """Updates the user version in settings.
+
+        Raises:
+            ReadOnlyError: If the store is in read-only mode.
+        """
+        self._assert_writable()
         settings_records = list(
             self.settings_table.search().limit(1).to_pydantic(SettingsRecord)
         )
@@ -267,7 +377,12 @@ class Store:
             )
 
     def recreate_embeddings_table(self) -> None:
-        """Recreate the chunks table with current vector dimensions."""
+        """Recreate the chunks table with current vector dimensions.
+
+        Raises:
+            ReadOnlyError: If the store is in read-only mode.
+        """
+        self._assert_writable()
         # Drop and recreate chunks table
         try:
             self.db.drop_table("chunks")
@@ -278,9 +393,9 @@ class Store:
         self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
         self.chunks_table = self.db.create_table("chunks", schema=self.ChunkRecord)
 
-        # Create FTS index on the new table with phrase query support
+        # Create FTS index on content_fts (contextualized content) for better search
         self.chunks_table.create_fts_index(
-            "content", replace=True, with_position=True, remove_stop_words=False
+            "content_fts", replace=True, with_position=True, remove_stop_words=False
         )
 
     def close(self):
@@ -297,7 +412,12 @@ class Store:
         }
 
     def restore_table_versions(self, versions: dict[str, int]) -> bool:
-        """Restore tables to the provided versions using LanceDB's API."""
+        """Restore tables to the provided versions using LanceDB's API.
+
+        Raises:
+            ReadOnlyError: If the store is in read-only mode.
+        """
+        self._assert_writable()
         self.documents_table.restore(int(versions["documents"]))
         self.chunks_table.restore(int(versions["chunks"]))
         self.settings_table.restore(int(versions["settings"]))
@@ -307,3 +427,83 @@ class Store:
     def _connection(self):
         """Compatibility property for repositories expecting _connection."""
         return self
+
+    def _checkout_tables_before(self, before: datetime) -> None:
+        """Checkout all tables to their state at or before the given datetime.
+
+        Args:
+            before: The datetime to checkout to
+
+        Raises:
+            ValueError: If no version exists before the given datetime
+        """
+        # LanceDB stores timestamps as naive datetimes in local time.
+        # Convert 'before' to naive local time for comparison.
+        if before.tzinfo is not None:
+            # Convert to local time and make naive
+            before_local = before.astimezone().replace(tzinfo=None)
+        else:
+            # Already naive, assume local time
+            before_local = before
+
+        tables = [
+            ("documents", self.documents_table),
+            ("chunks", self.chunks_table),
+            ("settings", self.settings_table),
+        ]
+
+        for table_name, table in tables:
+            versions = table.list_versions()
+            # Find the latest version at or before the target datetime
+            # Versions are sorted by version number, not timestamp, so we need to check all
+            best_version = None
+            best_timestamp = None
+
+            for v in versions:
+                # LanceDB version timestamps are naive datetime objects in local time
+                v_timestamp = v["timestamp"]
+                # Make sure it's naive for comparison
+                if v_timestamp.tzinfo is not None:
+                    v_timestamp = v_timestamp.replace(tzinfo=None)
+
+                if v_timestamp <= before_local:
+                    if best_timestamp is None or v_timestamp > best_timestamp:
+                        best_version = v["version"]
+                        best_timestamp = v_timestamp
+
+            if best_version is None:
+                # Find the earliest version to report in error message
+                if versions:
+                    earliest = min(versions, key=lambda v: v["timestamp"])
+                    earliest_ts = earliest["timestamp"]
+                    raise ValueError(
+                        f"No data exists before {before}. "
+                        f"Database was created on {earliest_ts}"
+                    )
+                else:
+                    raise ValueError(
+                        f"No data exists before {before}. Table has no versions."
+                    )
+
+            # Checkout to the found version
+            table.checkout(best_version)
+
+    def list_table_versions(self, table_name: str) -> list[dict[str, Any]]:
+        """List version history for a table.
+
+        Args:
+            table_name: Name of the table ("documents", "chunks", or "settings")
+
+        Returns:
+            List of version info dicts with "version" and "timestamp" keys
+        """
+        table_map = {
+            "documents": self.documents_table,
+            "chunks": self.chunks_table,
+            "settings": self.settings_table,
+        }
+        table = table_map.get(table_name)
+        if table is None:
+            raise ValueError(f"Unknown table: {table_name}")
+
+        return list(table.list_versions())