diff --git a/src/haiku/rag/store/engine.py b/src/haiku/rag/store/engine.py index 4d736834..517b139f 100644 --- a/src/haiku/rag/store/engine.py +++ b/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: "") diff --git a/src/haiku/rag/store/models/document.py b/src/haiku/rag/store/models/document.py index 39a46271..e23d5d3c 100644 --- a/src/haiku/rag/store/models/document.py +++ b/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) diff --git a/src/haiku/rag/store/repositories/document.py b/src/haiku/rag/store/repositories/document.py index bf8160e8..e0689ad7 100644 --- a/src/haiku/rag/store/repositories/document.py +++ b/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, }, diff --git a/src/haiku/rag/store/upgrades/__init__.py b/src/haiku/rag/store/upgrades/__init__.py index 5888cdba..cf9b8e19 100644 --- a/src/haiku/rag/store/upgrades/__init__.py +++ b/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) diff --git a/src/haiku/rag/store/upgrades/v0_10_2.py b/src/haiku/rag/store/upgrades/v0_10_2.py new file mode 100644 index 00000000..bb9b5113 --- /dev/null +++ b/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", +)