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

haiku/rag/client.py

@@ -33,8 +33,6 @@ class HaikuRAG:
             db_path: Path to the database file.
             skip_validation: Whether to skip configuration validation on database load.
         """
-        if not db_path.parent.exists():
-            Path.mkdir(db_path.parent, parents=True)
         self.store = Store(db_path, skip_validation=skip_validation)
         self.document_repository = DocumentRepository(self.store)
         self.chunk_repository = ChunkRepository(self.store)
@@ -52,6 +50,7 @@ class HaikuRAG:
         self,
         docling_document,
         uri: str | None = None,
+        title: str | None = None,
         metadata: dict | None = None,
         chunks: list[Chunk] | None = None,
     ) -> Document:
@@ -60,6 +59,7 @@ class HaikuRAG:
         document = Document(
             content=content,
             uri=uri,
+            title=title,
             metadata=metadata or {},
         )
         return await self.document_repository._create_with_docling(
@@ -70,6 +70,7 @@ class HaikuRAG:
         self,
         content: str,
         uri: str | None = None,
+        title: str | None = None,
         metadata: dict | None = None,
         chunks: list[Chunk] | None = None,
     ) -> Document:
@@ -90,6 +91,7 @@ class HaikuRAG:
         document = Document(
             content=content,
             uri=uri,
+            title=title,
             metadata=metadata or {},
         )
         return await self.document_repository._create_with_docling(
@@ -97,7 +99,7 @@ class HaikuRAG:
         )
 
     async def create_document_from_source(
-        self, source: str | Path, metadata: dict = {}
+        self, source: str | Path, title: str | None = None, metadata: dict | None = None
     ) -> Document:
         """Create or update a document from a file path or URL.
 
@@ -118,11 +120,16 @@ class HaikuRAG:
             httpx.RequestError: If URL request fails
         """
 
+        # Normalize metadata
+        metadata = metadata or {}
+
         # Check if it's a URL
         source_str = str(source)
         parsed_url = urlparse(source_str)
         if parsed_url.scheme in ("http", "https"):
-            return await self._create_or_update_document_from_url(source_str, metadata)
+            return await self._create_or_update_document_from_url(
+                source_str, title=title, metadata=metadata
+            )
         elif parsed_url.scheme == "file":
             # Handle file:// URI by converting to path
             source_path = Path(parsed_url.path)
@@ -138,37 +145,51 @@ class HaikuRAG:
         uri = source_path.absolute().as_uri()
         md5_hash = hashlib.md5(source_path.read_bytes()).hexdigest()
 
+        # Get content type from file extension (do before early return)
+        content_type, _ = mimetypes.guess_type(str(source_path))
+        if not content_type:
+            content_type = "application/octet-stream"
+        # Merge metadata with contentType and md5
+        metadata.update({"contentType": content_type, "md5": md5_hash})
+
         # Check if document already exists
         existing_doc = await self.get_document_by_uri(uri)
         if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
-            # MD5 unchanged, return existing document
+            # MD5 unchanged; update title/metadata if provided
+            updated = False
+            if title is not None and title != existing_doc.title:
+                existing_doc.title = title
+                updated = True
+            if metadata:
+                existing_doc.metadata = {**(existing_doc.metadata or {}), **metadata}
+                updated = True
+            if updated:
+                return await self.document_repository.update(existing_doc)
             return existing_doc
 
+        # Parse file only when content changed or new document
         docling_document = FileReader.parse_file(source_path)
 
-        # Get content type from file extension
-        content_type, _ = mimetypes.guess_type(str(source_path))
-        if not content_type:
-            content_type = "application/octet-stream"
-
-        # Merge metadata with contentType and md5
-        metadata.update({"contentType": content_type, "md5": md5_hash})
-
         if existing_doc:
             # Update existing document
             existing_doc.content = docling_document.export_to_markdown()
             existing_doc.metadata = metadata
+            if title is not None:
+                existing_doc.title = title
             return await self.document_repository._update_with_docling(
                 existing_doc, docling_document
             )
         else:
             # Create new document using DoclingDocument
             return await self._create_document_with_docling(
-                docling_document=docling_document, uri=uri, metadata=metadata
+                docling_document=docling_document,
+                uri=uri,
+                title=title,
+                metadata=metadata,
             )
 
     async def _create_or_update_document_from_url(
-        self, url: str, metadata: dict = {}
+        self, url: str, title: str | None = None, metadata: dict | None = None
     ) -> Document:
         """Create or update a document from a URL by downloading and parsing the content.
 
@@ -188,20 +209,35 @@ class HaikuRAG:
             ValueError: If the content cannot be parsed
             httpx.RequestError: If URL request fails
         """
+        metadata = metadata or {}
+
         async with httpx.AsyncClient() as client:
             response = await client.get(url)
             response.raise_for_status()
 
             md5_hash = hashlib.md5(response.content).hexdigest()
 
+            # Get content type early (used for potential no-op update)
+            content_type = response.headers.get("content-type", "").lower()
+
             # Check if document already exists
             existing_doc = await self.get_document_by_uri(url)
             if existing_doc and existing_doc.metadata.get("md5") == md5_hash:
-                # MD5 unchanged, return existing document
+                # MD5 unchanged; update title/metadata if provided
+                updated = False
+                if title is not None and title != existing_doc.title:
+                    existing_doc.title = title
+                    updated = True
+                metadata.update({"contentType": content_type, "md5": md5_hash})
+                if metadata:
+                    existing_doc.metadata = {
+                        **(existing_doc.metadata or {}),
+                        **metadata,
+                    }
+                    updated = True
+                if updated:
+                    return await self.document_repository.update(existing_doc)
                 return existing_doc
-
-            # Get content type to determine file extension
-            content_type = response.headers.get("content-type", "").lower()
             file_extension = self._get_extension_from_content_type_or_url(
                 url, content_type
             )
@@ -228,12 +264,17 @@ class HaikuRAG:
             if existing_doc:
                 existing_doc.content = docling_document.export_to_markdown()
                 existing_doc.metadata = metadata
+                if title is not None:
+                    existing_doc.title = title
                 return await self.document_repository._update_with_docling(
                     existing_doc, docling_document
                 )
             else:
                 return await self._create_document_with_docling(
-                    docling_document=docling_document, uri=url, metadata=metadata
+                    docling_document=docling_document,
+                    uri=url,
+                    title=title,
+                    metadata=metadata,
                 )
 
     def _get_extension_from_content_type_or_url(
@@ -418,6 +459,7 @@ class HaikuRAG:
                     content="".join(combined_content_parts),
                     metadata=original_chunk.metadata,
                     document_uri=original_chunk.document_uri,
+                    document_title=original_chunk.document_title,
                     document_meta=original_chunk.document_meta,
                 )
 
@@ -524,7 +566,7 @@ class HaikuRAG:
 
                     # Try to re-create from source (this creates the document with chunks)
                     new_doc = await self.create_document_from_source(
-                        doc.uri, doc.metadata or {}
+                        source=doc.uri, metadata=doc.metadata or {}
                     )
 
                     assert new_doc.id is not None, "New document ID should not be None"

haiku/rag/config.py

@@ -53,6 +53,10 @@ class AppConfig(BaseModel):
     ANTHROPIC_API_KEY: str = ""
     COHERE_API_KEY: str = ""
 
+    # If true, refuse to auto-create a new LanceDB database or tables
+    # and error out when the database does not already exist.
+    DISABLE_DB_AUTOCREATE: bool = False
+
     @field_validator("MONITOR_DIRECTORIES", mode="before")
     @classmethod
     def parse_monitor_directories(cls, v):

haiku/rag/mcp.py

@@ -17,6 +17,7 @@ class DocumentResult(BaseModel):
     id: str | None
     content: str
     uri: str | None = None
+    title: str | None = None
     metadata: dict[str, Any] = {}
     created_at: str
     updated_at: str
@@ -28,13 +29,15 @@ def create_mcp_server(db_path: Path) -> FastMCP:
 
     @mcp.tool()
     async def add_document_from_file(
-        file_path: str, metadata: dict[str, Any] | None = None
+        file_path: str,
+        metadata: dict[str, Any] | None = None,
+        title: str | None = None,
     ) -> str | None:
         """Add a document to the RAG system from a file path."""
         try:
             async with HaikuRAG(db_path) as rag:
                 document = await rag.create_document_from_source(
-                    Path(file_path), metadata or {}
+                    Path(file_path), title=title, metadata=metadata or {}
                 )
                 return document.id
         except Exception:
@@ -42,24 +45,31 @@ def create_mcp_server(db_path: Path) -> FastMCP:
 
     @mcp.tool()
     async def add_document_from_url(
-        url: str, metadata: dict[str, Any] | None = None
+        url: str, metadata: dict[str, Any] | None = None, title: str | None = None
     ) -> str | None:
         """Add a document to the RAG system from a URL."""
         try:
             async with HaikuRAG(db_path) as rag:
-                document = await rag.create_document_from_source(url, metadata or {})
+                document = await rag.create_document_from_source(
+                    url, title=title, metadata=metadata or {}
+                )
                 return document.id
         except Exception:
             return None
 
     @mcp.tool()
     async def add_document_from_text(
-        content: str, uri: str | None = None, metadata: dict[str, Any] | None = None
+        content: str,
+        uri: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        title: str | None = None,
     ) -> str | None:
         """Add a document to the RAG system from text content."""
         try:
             async with HaikuRAG(db_path) as rag:
-                document = await rag.create_document(content, uri, metadata or {})
+                document = await rag.create_document(
+                    content, uri, title=title, metadata=metadata or {}
+                )
                 return document.id
         except Exception:
             return None
@@ -102,6 +112,7 @@ def create_mcp_server(db_path: Path) -> FastMCP:
                     id=document.id,
                     content=document.content,
                     uri=document.uri,
+                    title=document.title,
                     metadata=document.metadata,
                     created_at=str(document.created_at),
                     updated_at=str(document.updated_at),
@@ -123,6 +134,7 @@ def create_mcp_server(db_path: Path) -> FastMCP:
                         id=doc.id,
                         content=doc.content,
                         uri=doc.uri,
+                        title=doc.title,
                         metadata=doc.metadata,
                         created_at=str(doc.created_at),
                         updated_at=str(doc.updated_at),

haiku/rag/migration.py

@@ -51,7 +51,7 @@ class SQLiteToLanceDBMigrator:
 
                 sqlite_conn.enable_load_extension(True)
                 sqlite_vec.load(sqlite_conn)
-                self.console.print("[blue]Loaded sqlite-vec extension[/blue]")
+                self.console.print("[cyan]Loaded sqlite-vec extension[/cyan]")
             except Exception as e:
                 self.console.print(
                     f"[yellow]Warning: Could not load sqlite-vec extension: {e}[/yellow]"
@@ -92,7 +92,7 @@ class SQLiteToLanceDBMigrator:
             sqlite_conn.close()
 
             # Optimize and cleanup using centralized vacuum
-            self.console.print("[blue]Optimizing LanceDB...[/blue]")
+            self.console.print("[cyan]Optimizing LanceDB...[/cyan]")
             try:
                 lance_store.vacuum()
                 self.console.print("[green]✅ Optimization completed[/green]")

haiku/rag/qa/agent.py

@@ -12,7 +12,9 @@ from haiku.rag.qa.prompts import QA_SYSTEM_PROMPT, QA_SYSTEM_PROMPT_WITH_CITATIO
 class SearchResult(BaseModel):
     content: str = Field(description="The document text content")
     score: float = Field(description="Relevance score (higher is more relevant)")
-    document_uri: str = Field(description="Source URI/path of the document")
+    document_uri: str = Field(
+        description="Source title (if available) or URI/path of the document"
+    )
 
 
 class Dependencies(BaseModel):
@@ -59,7 +61,7 @@ class QuestionAnswerAgent:
                 SearchResult(
                     content=chunk.content,
                     score=score,
-                    document_uri=chunk.document_uri or "",
+                    document_uri=(chunk.document_title or chunk.document_uri or ""),
                 )
                 for chunk, score in expanded_results
             ]

haiku/rag/qa/prompts.py

@@ -44,9 +44,9 @@ Guidelines:
 
 Citation Format:
 After your answer, include a "Citations:" section that lists:
-- The document URI from each search result used
+- 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_uri]: [content_excerpt]..."
+- Format: "Citations:\n- [document title or URI]: [content_excerpt]..."
 
 Example response format:
 [Your answer here]

haiku/rag/research/models.py

@@ -19,8 +19,8 @@ class SearchAnswer(BaseModel):
     )
     sources: list[str] = Field(
         description=(
-            "Document URIs corresponding to the snippets actually used in the"
-            " answer (one URI per snippet; omit if none)"
+            "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/nodes/search.py

@@ -59,7 +59,9 @@ class SearchDispatchNode(BaseNode[ResearchState, ResearchDeps, ResearchReport]):
                     {
                         "text": chunk.content,
                         "score": score,
-                        "document_uri": (chunk.document_uri or ""),
+                        "document_uri": (
+                            chunk.document_title or chunk.document_uri or ""
+                        ),
                     }
                     for chunk, score in expanded
                 ]

haiku/rag/research/prompts.py

@@ -27,13 +27,14 @@ Tasks:
 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_uri`.
+  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_uris for the snippets
-  you used (one URI per snippet; same order as context). Context must be text‑only.
+- 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.
 

haiku/rag/store/__init__.py

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

haiku/rag/store/engine.py

@@ -19,6 +19,7 @@ class DocumentRecord(LanceModel):
     id: str = Field(default_factory=lambda: str(uuid4()))
     content: str
     uri: str | None = None
+    title: str | None = None
     metadata: str = Field(default="{}")
     created_at: str = Field(default_factory=lambda: "")
     updated_at: str = Field(default_factory=lambda: "")
@@ -54,6 +55,19 @@ class Store:
         # Create the ChunkRecord model with the correct vector dimension
         self.ChunkRecord = create_chunk_model(self.embedder._vector_dim)
 
+        # Local filesystem handling for DB directory
+        if not self._has_cloud_config():
+            if Config.DISABLE_DB_AUTOCREATE:
+                # LanceDB uses a directory path for local databases; enforce presence
+                if not db_path.exists():
+                    raise FileNotFoundError(
+                        f"LanceDB path does not exist: {db_path}. Auto-creation is disabled."
+                    )
+            else:
+                # Ensure parent directories exist when autocreation allowed
+                if not db_path.parent.exists():
+                    Path.mkdir(db_path.parent, parents=True)
+
         # Connect to LanceDB
         self.db = self._connect_to_lancedb(db_path)
 

haiku/rag/store/models/__init__.py

@@ -1,4 +1,4 @@
 from .chunk import Chunk
 from .document import Document
 
-__all__ = ["Chunk", "Document"]
+__all__ = ["Chunk", "Document"]

haiku/rag/store/models/chunk.py

@@ -12,5 +12,6 @@ class Chunk(BaseModel):
     metadata: dict = {}
     order: int = 0
     document_uri: str | None = None
+    document_title: str | None = None
     document_meta: dict = {}
     embedding: list[float] | None = None

haiku/rag/store/models/document.py

@@ -11,6 +11,7 @@ class Document(BaseModel):
     id: str | None = None
     content: str
     uri: str | None = None
+    title: str | None = None
     metadata: dict = {}
     created_at: datetime = Field(default_factory=datetime.now)
     updated_at: datetime = Field(default_factory=datetime.now)

haiku/rag/store/repositories/chunk.py

@@ -317,6 +317,7 @@ class ChunkRepository:
         )
 
         doc_uri = doc_results[0].uri if doc_results else None
+        doc_title = doc_results[0].title if doc_results else None
         doc_meta = doc_results[0].metadata if doc_results else "{}"
 
         chunks: list[Chunk] = []
@@ -330,6 +331,7 @@ class ChunkRepository:
                     metadata=md,
                     order=rec.order,
                     document_uri=doc_uri,
+                    document_title=doc_title,
                     document_meta=json.loads(doc_meta),
                 )
             )
@@ -398,6 +400,7 @@ class ChunkRepository:
             # Get document info from pre-fetched map
             doc = documents_map.get(chunk_record.document_id)
             doc_uri = doc.uri if doc else None
+            doc_title = doc.title if doc else None
             doc_meta = doc.metadata if doc else "{}"
 
             md = json.loads(chunk_record.metadata)
@@ -409,6 +412,7 @@ class ChunkRepository:
                 metadata=md,
                 order=chunk_record.order,
                 document_uri=doc_uri,
+                document_title=doc_title,
                 document_meta=json.loads(doc_meta),
             )
 

haiku/rag/store/repositories/document.py

@@ -34,6 +34,7 @@ class DocumentRepository:
             id=record.id,
             content=record.content,
             uri=record.uri,
+            title=record.title,
             metadata=json.loads(record.metadata),
             created_at=datetime.fromisoformat(record.created_at)
             if record.created_at
@@ -56,6 +57,7 @@ class DocumentRepository:
             id=doc_id,
             content=entity.content,
             uri=entity.uri,
+            title=entity.title,
             metadata=json.dumps(entity.metadata),
             created_at=now,
             updated_at=now,
@@ -97,6 +99,7 @@ class DocumentRepository:
             values={
                 "content": entity.content,
                 "uri": entity.uri,
+                "title": entity.title,
                 "metadata": json.dumps(entity.metadata),
                 "updated_at": now,
             },

haiku/rag/store/upgrades/__init__.py

@@ -55,6 +55,8 @@ def run_pending_upgrades(store: Store, from_version: str, to_version: str) -> No
 
 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
+from .v0_10_1 import upgrade_add_title as upgrade_0_10_1_add_title  # noqa: E402
 
 upgrades.append(upgrade_0_9_3_order)
 upgrades.append(upgrade_0_9_3_fts)
+upgrades.append(upgrade_0_10_1_add_title)

haiku/rag/store/upgrades/v0_10_1.py

@@ -0,0 +1,64 @@
+import json
+
+from lancedb.pydantic import LanceModel
+from pydantic import Field
+
+from haiku.rag.store.engine import Store
+from haiku.rag.store.upgrades import Upgrade
+
+
+def _apply_add_document_title(store: Store) -> None:
+    """Add a nullable 'title' column to the documents table."""
+
+    # Read existing rows using Arrow for schema-agnostic access
+    try:
+        docs_arrow = store.documents_table.search().to_arrow()
+        rows = docs_arrow.to_pylist()
+    except Exception:
+        rows = []
+
+    class DocumentRecordV2(LanceModel):
+        id: str
+        content: str
+        uri: str | None = None
+        title: str | None = None
+        metadata: str = Field(default="{}")
+        created_at: str = Field(default_factory=lambda: "")
+        updated_at: str = Field(default_factory=lambda: "")
+
+    # Drop and recreate documents table with the new schema
+    try:
+        store.db.drop_table("documents")
+    except Exception:
+        pass
+
+    store.documents_table = store.db.create_table("documents", schema=DocumentRecordV2)
+
+    # Reinsert previous rows with title=None
+    if rows:
+        backfilled = []
+        for row in rows:
+            backfilled.append(
+                DocumentRecordV2(
+                    id=row.get("id"),
+                    content=row.get("content", ""),
+                    uri=row.get("uri"),
+                    title=None,
+                    metadata=(
+                        row.get("metadata")
+                        if isinstance(row.get("metadata"), str)
+                        else json.dumps(row.get("metadata") or {})
+                    ),
+                    created_at=row.get("created_at", ""),
+                    updated_at=row.get("updated_at", ""),
+                )
+            )
+
+        store.documents_table.add(backfilled)
+
+
+upgrade_add_title = Upgrade(
+    version="0.10.1",
+    apply=_apply_add_document_title,
+    description="Add nullable 'title' column to documents table",
+)

haiku/rag/utils.py

@@ -9,10 +9,6 @@ from io import BytesIO
 from pathlib import Path
 from types import ModuleType
 
-import httpx
-from docling.document_converter import DocumentConverter
-from docling_core.types.doc.document import DoclingDocument
-from docling_core.types.io import DocumentStream
 from packaging.version import Version, parse
 
 
@@ -82,6 +78,9 @@ async def is_up_to_date() -> tuple[bool, Version, Version]:
         the running version and the latest version.
     """
 
+    # Lazy import to avoid pulling httpx (and its deps) on module import
+    import httpx
+
     async with httpx.AsyncClient() as client:
         running_version = parse(metadata.version("haiku.rag"))
         try:
@@ -94,7 +93,7 @@ async def is_up_to_date() -> tuple[bool, Version, Version]:
     return running_version >= pypi_version, running_version, pypi_version
 
 
-def text_to_docling_document(text: str, name: str = "content.md") -> DoclingDocument:
+def text_to_docling_document(text: str, name: str = "content.md"):
     """Convert text content to a DoclingDocument.
 
     Args:
@@ -104,6 +103,10 @@ def text_to_docling_document(text: str, name: str = "content.md") -> DoclingDocu
     Returns:
         A DoclingDocument created from the text content.
     """
+    # Lazy import docling deps to keep import-time light
+    from docling.document_converter import DocumentConverter  # type: ignore
+    from docling_core.types.io import DocumentStream  # type: ignore
+
     bytes_io = BytesIO(text.encode("utf-8"))
     doc_stream = DocumentStream(name=name, stream=bytes_io)
     converter = DocumentConverter()
@@ -160,3 +163,37 @@ def load_callable(path: str):
             f"Attribute '{func_name}' in module '{module_part}' is not callable"
         )
     return func
+
+
+def prefetch_models():
+    """Prefetch runtime models (Docling + Ollama as configured)."""
+    import httpx
+    from docling.utils.model_downloader import download_models
+
+    from haiku.rag.config import Config
+
+    download_models()
+
+    # Collect Ollama models from config
+    required_models: set[str] = set()
+    if Config.EMBEDDINGS_PROVIDER == "ollama":
+        required_models.add(Config.EMBEDDINGS_MODEL)
+    if Config.QA_PROVIDER == "ollama":
+        required_models.add(Config.QA_MODEL)
+    if Config.RESEARCH_PROVIDER == "ollama":
+        required_models.add(Config.RESEARCH_MODEL)
+    if Config.RERANK_PROVIDER == "ollama":
+        required_models.add(Config.RERANK_MODEL)
+
+    if not required_models:
+        return
+
+    base_url = Config.OLLAMA_BASE_URL
+
+    with httpx.Client(timeout=None) as client:
+        for model in sorted(required_models):
+            with client.stream(
+                "POST", f"{base_url}/api/pull", json={"model": model}
+            ) as r:
+                for _ in r.iter_lines():
+                    pass