DB migration and schema updates to introduce title in documents

src/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: "")

src/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)

src/haiku/rag/store/repositories/document.py

@@ -34,6 +34,7 @@ class DocumentRepository:
             id=record.id,
             content=record.content,
             uri=record.uri,
+            title=getattr(record, "title", None),
             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,
             },

src/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_2 import upgrade_add_title as upgrade_0_10_2_add_title  # noqa: E402
 
 upgrades.append(upgrade_0_9_3_order)
 upgrades.append(upgrade_0_9_3_fts)
+upgrades.append(upgrade_0_10_2_add_title)

src/haiku/rag/store/upgrades/v0_10_2.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.2",
+    apply=_apply_add_document_title,
+    description="Add nullable 'title' column to documents table",
+)