From f9a617a911d017c7f6f4f9ab86dc65bb55e44f66 Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Thu, 4 Dec 2025 15:55:22 -0500 Subject: [PATCH] fix: LanceDB --- .gitignore | 1 + ...6_add_chunk_index_to_message_embeddings.py | 54 ----------- .../f1a2b3c4d5e6_make_embeddings_nullable.py | 97 +++++++++++++++++++ src/crud/document.py | 17 ++-- src/vector_store/lancedb.py | 31 +++++- tests/alembic/revisions/__init__.py | 4 +- ..._f1a2b3c4d5e6_make_embeddings_nullable.py} | 0 7 files changed, 137 insertions(+), 67 deletions(-) delete mode 100644 migrations/versions/f1a2b3c4d5e6_add_chunk_index_to_message_embeddings.py create mode 100644 migrations/versions/f1a2b3c4d5e6_make_embeddings_nullable.py rename tests/alembic/revisions/{test_f1a2b3c4d5e6_add_chunk_index_to_message_embeddings.py => test_f1a2b3c4d5e6_make_embeddings_nullable.py} (100%) diff --git a/.gitignore b/.gitignore index 17d9597c..a6d64b56 100644 --- a/.gitignore +++ b/.gitignore @@ -189,3 +189,4 @@ CRUSH.md metrics.jsonl AGENTS.md +lancedb_data/ diff --git a/migrations/versions/f1a2b3c4d5e6_add_chunk_index_to_message_embeddings.py b/migrations/versions/f1a2b3c4d5e6_add_chunk_index_to_message_embeddings.py deleted file mode 100644 index 0275fdab..00000000 --- a/migrations/versions/f1a2b3c4d5e6_add_chunk_index_to_message_embeddings.py +++ /dev/null @@ -1,54 +0,0 @@ -"""add chunk_index to message_embeddings - -This migration adds the chunk_index column to message_embeddings table for tracking -chunked message embeddings in external vector stores (turbopuffer/lancedb). - -Revision ID: f1a2b3c4d5e6 -Revises: baa22cad81e2 -Create Date: 2025-11-24 12:00:00.000000 - -""" - -from collections.abc import Sequence - -import sqlalchemy as sa -from alembic import op - -from migrations.utils import column_exists, get_schema - -# revision identifiers, used by Alembic. -revision: str = "f1a2b3c4d5e6" -down_revision: str | None = "baa22cad81e2" -branch_labels: str | Sequence[str] | None = None -depends_on: str | Sequence[str] | None = None - -schema = get_schema() - - -def upgrade() -> None: - """Add chunk_index column to message_embeddings for tracking chunked message embeddings.""" - inspector = sa.inspect(op.get_bind()) - - # Add chunk_index column to message_embeddings if it doesn't exist - # This is needed to track which chunk of a message this embedding represents - # Vector ID format: {message_public_id}_{chunk_index} - if not column_exists("message_embeddings", "chunk_index", inspector): - op.add_column( - "message_embeddings", - sa.Column( - "chunk_index", - sa.Integer(), - nullable=False, - server_default="0", - ), - schema=schema, - ) - - -def downgrade() -> None: - """Remove chunk_index column from message_embeddings.""" - inspector = sa.inspect(op.get_bind()) - - # Remove chunk_index column if it exists - if column_exists("message_embeddings", "chunk_index", inspector): - op.drop_column("message_embeddings", "chunk_index", schema=schema) diff --git a/migrations/versions/f1a2b3c4d5e6_make_embeddings_nullable.py b/migrations/versions/f1a2b3c4d5e6_make_embeddings_nullable.py new file mode 100644 index 00000000..5bdfeb14 --- /dev/null +++ b/migrations/versions/f1a2b3c4d5e6_make_embeddings_nullable.py @@ -0,0 +1,97 @@ +"""add chunk_index to message_embeddings and make embeddings nullable + +This migration: +1. Adds the chunk_index column to message_embeddings table for tracking + chunked message embeddings in external vector stores (turbopuffer/lancedb). +2. Makes embedding columns nullable in both message_embeddings and documents tables + since embeddings are now stored in external vector stores instead of PostgreSQL. + +Revision ID: f1a2b3c4d5e6 +Revises: baa22cad81e2 +Create Date: 2025-11-24 12:00:00.000000 + +""" + +from collections.abc import Sequence + +import sqlalchemy as sa +from alembic import op +from pgvector.sqlalchemy import Vector + +from migrations.utils import column_exists, get_schema + +# revision identifiers, used by Alembic. +revision: str = "f1a2b3c4d5e6" +down_revision: str | None = "baa22cad81e2" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +schema = get_schema() + + +def upgrade() -> None: + """Add chunk_index column to message_embeddings and make embeddings nullable.""" + inspector = sa.inspect(op.get_bind()) + + # Add chunk_index column to message_embeddings if it doesn't exist + # This is needed to track which chunk of a message this embedding represents + # Vector ID format: {message_public_id}_{chunk_index} + if not column_exists("message_embeddings", "chunk_index", inspector): + op.add_column( + "message_embeddings", + sa.Column( + "chunk_index", + sa.Integer(), + nullable=False, + server_default="0", + ), + schema=schema, + ) + + # Make message_embeddings.embedding nullable since embeddings are now stored + # in external vector stores (turbopuffer/lancedb) instead of PostgreSQL + op.alter_column( + "message_embeddings", + "embedding", + existing_type=Vector(1536), + nullable=True, + schema=schema, + ) + + # Make documents.embedding nullable for the same reason + # (this should already be nullable, but ensure it for consistency) + op.alter_column( + "documents", + "embedding", + existing_type=Vector(1536), + nullable=True, + schema=schema, + ) + + +def downgrade() -> None: + """Remove chunk_index column and revert embedding columns to non-nullable.""" + inspector = sa.inspect(op.get_bind()) + + # Revert documents.embedding back to nullable=True (it was originally nullable=True) + op.alter_column( + "documents", + "embedding", + existing_type=Vector(1536), + nullable=True, # Keep as nullable since it was nullable in the original schema + schema=schema, + ) + + # Revert message_embeddings.embedding back to nullable=False + # Note: This may fail if there are NULL values in the database + op.alter_column( + "message_embeddings", + "embedding", + existing_type=Vector(1536), + nullable=False, + schema=schema, + ) + + # Remove chunk_index column if it exists + if column_exists("message_embeddings", "chunk_index", inspector): + op.drop_column("message_embeddings", "chunk_index", schema=schema) diff --git a/src/crud/document.py b/src/crud/document.py index 7bcb3497..be04b438 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -153,7 +153,8 @@ async def create_documents( Count of new documents """ honcho_documents: list[models.Document] = [] - embeddings_to_store: list[tuple[str, list[float]]] = [] # [(doc_id, embedding)] + # Store (document_model, embedding) pairs - IDs aren't available until after commit + docs_with_embeddings: list[tuple[models.Document, list[float]]] = [] for doc in documents: try: @@ -180,9 +181,9 @@ async def create_documents( ) honcho_documents.append(new_doc) - # Track embedding for vector store (will use document's generated ID) + # Track embedding for vector store (ID will be available after commit) if doc.embedding: - embeddings_to_store.append((new_doc.id, doc.embedding)) + docs_with_embeddings.append((new_doc, doc.embedding)) except Exception as e: logger.error( @@ -194,8 +195,8 @@ async def create_documents( db.add_all(honcho_documents) await db.commit() - # Store embeddings in vector store after documents are committed - if embeddings_to_store: + # Store embeddings in vector store after documents are committed (IDs now available) + if docs_with_embeddings: vector_store = get_vector_store() namespace = vector_store.get_document_namespace( workspace_name, observer, observed @@ -203,12 +204,10 @@ async def create_documents( # Build vector records with metadata for filtering vector_records: list[VectorRecord] = [] - doc_lookup = {doc.id: doc for doc in honcho_documents} - for doc_id, embedding in embeddings_to_store: - doc = doc_lookup[doc_id] + for doc, embedding in docs_with_embeddings: vector_records.append( VectorRecord( - id=doc_id, + id=doc.id, embedding=embedding, metadata={ "workspace_name": workspace_name, diff --git a/src/vector_store/lancedb.py b/src/vector_store/lancedb.py index 49f05c4f..547c421b 100644 --- a/src/vector_store/lancedb.py +++ b/src/vector_store/lancedb.py @@ -29,6 +29,11 @@ class LanceDBVectorStore(VectorStore): Uses LanceDB's embedded mode for local vector storage. Each namespace corresponds to a LanceDB table. + + Note: LanceDB table names can only contain alphanumeric characters, + underscores, hyphens, and periods. We use '.' as the namespace separator + instead of ':' (used by base class) since '.' is not allowed in + workspace/peer IDs. """ _db: lancedb.DBConnection @@ -38,6 +43,25 @@ class LanceDBVectorStore(VectorStore): super().__init__() self._db = lancedb.connect(settings.VECTOR_STORE.LANCEDB_PATH) + # === Namespace helpers (override to use LanceDB-compatible separator) === + def get_document_namespace( + self, workspace_name: str, observer: str, observed: str + ) -> str: + """ + Get the namespace for document embeddings (per collection). + + Uses '.' as separator instead of ':' for LanceDB compatibility. + """ + return f"{self.namespace_prefix}.{workspace_name}.{observer}.{observed}" + + def get_message_namespace(self, workspace_name: str) -> str: + """ + Get the namespace for message embeddings (per workspace). + + Uses '.' as separator instead of ':' for LanceDB compatibility. + """ + return f"{self.namespace_prefix}.{workspace_name}.messages" + def _get_table(self, namespace: str) -> lancedb.table.Table | None: """Get a table if it exists, otherwise return None.""" if namespace in self._db.table_names(): @@ -47,7 +71,7 @@ class LanceDBVectorStore(VectorStore): def _get_or_create_table( self, namespace: str, sample_data: list[dict[str, Any]] | None = None ) -> lancedb.table.Table: - """ + """_get_or_create_table Get existing table or create if not exists. Args: @@ -127,10 +151,13 @@ class LanceDBVectorStore(VectorStore): try: rows = [self._row_to_dict(v) for v in vectors] + print(f"Rows: {rows}") table = self._get_or_create_table(namespace, sample_data=rows) # Use merge_insert for upsert behavior - table.merge_insert("id").when_matched_update_all().execute(rows) + table.merge_insert( + "id" + ).when_matched_update_all().when_not_matched_insert_all().execute(rows) logger.debug(f"Upserted {len(vectors)} vectors to namespace {namespace}") except Exception: diff --git a/tests/alembic/revisions/__init__.py b/tests/alembic/revisions/__init__.py index f9bfc8d1..04b14e98 100644 --- a/tests/alembic/revisions/__init__.py +++ b/tests/alembic/revisions/__init__.py @@ -21,7 +21,7 @@ from . import ( test_d429de0e5338_adopt_peer_paradigm, test_e9b705f9adf9_add_server_defaults_to_timestamp_, test_ec8f94139b02_codify_workspace_name_and_message_id_in_, - test_f1a2b3c4d5e6_add_chunk_index_to_message_embeddings, + test_f1a2b3c4d5e6_make_embeddings_nullable, ) __all__ = [ @@ -45,5 +45,5 @@ __all__ = [ "test_d429de0e5338_adopt_peer_paradigm", "test_e9b705f9adf9_add_server_defaults_to_timestamp_", "test_ec8f94139b02_codify_workspace_name_and_message_id_in_", - "test_f1a2b3c4d5e6_add_chunk_index_to_message_embeddings", + "test_f1a2b3c4d5e6_make_embeddings_nullable", ] diff --git a/tests/alembic/revisions/test_f1a2b3c4d5e6_add_chunk_index_to_message_embeddings.py b/tests/alembic/revisions/test_f1a2b3c4d5e6_make_embeddings_nullable.py similarity index 100% rename from tests/alembic/revisions/test_f1a2b3c4d5e6_add_chunk_index_to_message_embeddings.py rename to tests/alembic/revisions/test_f1a2b3c4d5e6_make_embeddings_nullable.py