From b207403eb49b0cca66974d9baf73146ab1986d18 Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Thu, 4 Dec 2025 16:42:47 -0500 Subject: [PATCH] fix: turbopuffer --- .env.template | 9 +++++++ config.toml.example | 7 ++++++ src/config.py | 2 -- src/crud/document.py | 9 +++---- src/dreamer/dreamer.py | 6 +++++ src/vector_store/__init__.py | 17 +++++++------ src/vector_store/lancedb.py | 26 -------------------- src/vector_store/turbopuffer.py | 20 ++++++++++++--- tests/crud/test_workspace.py | 2 -- tests/integration/test_message_embeddings.py | 18 +++++--------- tests/integration/test_representation.py | 6 +---- 11 files changed, 59 insertions(+), 63 deletions(-) diff --git a/.env.template b/.env.template index d1ad3da8..222b61d2 100644 --- a/.env.template +++ b/.env.template @@ -180,3 +180,12 @@ LLM_ANTHROPIC_API_KEY=your-anthropic-api-key-here # CACHE_NAMESPACE="honcho" # CACHE_DEFAULT_TTL_SECONDS=300 # CACHE_DEFAULT_LOCK_TTL_SECONDS=5 + +# ============================================================================= +# Vector Store Settings +# ============================================================================= +# VECTOR_STORE_TYPE="lancedb" +# VECTOR_STORE_NAMESPACE="honcho" +# VECTOR_STORE_TURBOPUFFER_API_KEY= +# VECTOR_STORE_TURBOPUFFER_REGION= +# VECTOR_STORE_LANCEDB_PATH="./lancedb_data" diff --git a/config.toml.example b/config.toml.example index 8ef1b5e1..1d5b8d8d 100644 --- a/config.toml.example +++ b/config.toml.example @@ -141,3 +141,10 @@ URL = "redis://localhost:6379/0" NAMESPACE="honcho" DEFAULT_TTL_SECONDS = 300 DEFAULT_LOCK_TTL_SECONDS = 5 + +# Vector store settings +# VECTOR_STORE_TYPE = "lancedb" +# VECTOR_STORE_NAMESPACE = "honcho" +# VECTOR_STORE_TURBOPUFFER_API_KEY = +# VECTOR_STORE_TURBOPUFFER_REGION = +# VECTOR_STORE_LANCEDB_PATH = "./lancedb_data" diff --git a/src/config.py b/src/config.py index 8df92c25..6e2ca48d 100644 --- a/src/config.py +++ b/src/config.py @@ -373,8 +373,6 @@ class VectorStoreSettings(HonchoSettings): # Turbopuffer-specific settings TURBOPUFFER_API_KEY: str | None = None - # Turbopuffer region (e.g., "gcp-us-east4", "aws-us-east-1") - # Can also be set via TURBOPUFFER_REGION environment variable TURBOPUFFER_REGION: str | None = None # LanceDB-specific settings (local embedded mode) diff --git a/src/crud/document.py b/src/crud/document.py index be04b438..4334f854 100644 --- a/src/crud/document.py +++ b/src/crud/document.py @@ -283,17 +283,16 @@ async def is_rejected_duplicate( logger.warning( f"[DUPLICATE DETECTION] Deleting existing in favor of new. new='{doc.content}', existing='{existing_doc.content}'." ) - # Delete from database - await db.delete(existing_doc) - await db.flush() # Flush to make deletion visible in this transaction - - # Delete from vector store vector_store = get_vector_store() namespace = vector_store.get_document_namespace( workspace_name, observer, observed ) await vector_store.delete_many(namespace, [existing_doc.id]) + # Delete from database after vector store succeeds + await db.delete(existing_doc) + await db.flush() # Flush to make deletion visible in this transaction + return False # Don't reject the new document # Existing document has more information, reject the new one diff --git a/src/dreamer/dreamer.py b/src/dreamer/dreamer.py index 5da9d1b6..65d7964e 100644 --- a/src/dreamer/dreamer.py +++ b/src/dreamer/dreamer.py @@ -17,6 +17,7 @@ from src.utils.representation import ( ExplicitObservation, Representation, ) +from src.vector_store import get_vector_store logger = logging.getLogger(__name__) @@ -168,6 +169,11 @@ async def _consolidate_cluster( # delete old documents for doc in cluster: await db.delete(doc) + doc_ids = [doc.id for doc in cluster] + + vector_store = get_vector_store() + namespace = vector_store.get_document_namespace(workspace_name, observer, observed) + await vector_store.delete_many(namespace, doc_ids) await db.commit() diff --git a/src/vector_store/__init__.py b/src/vector_store/__init__.py index 5437537a..0b3e0cf5 100644 --- a/src/vector_store/__init__.py +++ b/src/vector_store/__init__.py @@ -32,11 +32,12 @@ class VectorStore(ABC): Abstract base class for vector store implementations. All vector operations are namespace-scoped. Namespaces map to: - - Document embeddings: {prefix}:{workspace}:{observer}:{observed} (per collection) - - Message embeddings: {prefix}:{workspace}:messages (per workspace) + - Document embeddings: {prefix}.{workspace}.{observer}.{observed} (per collection) + - Message embeddings: {prefix}.{workspace}.messages (per workspace) - Note: Colon (:) is used as the delimiter to avoid collisions since it's not - allowed in workspace/peer IDs (which only allow [A-Za-z0-9_-]). + Note: Period (.) is used as the delimiter since vector stores (Turbopuffer, LanceDB) + only allow [A-Za-z0-9-_.] in namespace names, and period is not allowed in + workspace/peer IDs (which only allow [A-Za-z0-9_-]). """ namespace_prefix: str @@ -60,9 +61,9 @@ class VectorStore(ABC): observed: Name of the observed peer Returns: - Namespace string in format: {prefix}:{workspace}:{observer}:{observed} + Namespace string in format: {prefix}.{workspace}.{observer}.{observed} """ - return f"{self.namespace_prefix}:{workspace_name}:{observer}:{observed}" + return f"{self.namespace_prefix}.{workspace_name}.{observer}.{observed}" def get_message_namespace(self, workspace_name: str) -> str: """ @@ -72,9 +73,9 @@ class VectorStore(ABC): workspace_name: Name of the workspace Returns: - Namespace string in format: {prefix}:{workspace}:messages + Namespace string in format: {prefix}.{workspace}.messages """ - return f"{self.namespace_prefix}:{workspace_name}:messages" + return f"{self.namespace_prefix}.{workspace_name}.messages" # === Core operations === @abstractmethod diff --git a/src/vector_store/lancedb.py b/src/vector_store/lancedb.py index 547c421b..c145ae5d 100644 --- a/src/vector_store/lancedb.py +++ b/src/vector_store/lancedb.py @@ -29,11 +29,6 @@ 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 @@ -43,25 +38,6 @@ 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(): @@ -151,7 +127,6 @@ 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 @@ -307,7 +282,6 @@ class LanceDBVectorStore(VectorStore): try: if namespace in self._db.table_names(): self._db.drop_table(namespace) - logger.debug(f"Deleted namespace {namespace}") else: logger.debug(f"Namespace {namespace} does not exist, nothing to delete") except Exception: diff --git a/src/vector_store/turbopuffer.py b/src/vector_store/turbopuffer.py index 30a05884..58fd4f07 100644 --- a/src/vector_store/turbopuffer.py +++ b/src/vector_store/turbopuffer.py @@ -9,7 +9,7 @@ import logging from collections.abc import Sequence from typing import Any, Literal -from turbopuffer import Turbopuffer +from turbopuffer import NotFoundError, Turbopuffer from turbopuffer.lib.namespace import Namespace from turbopuffer.types import Filter @@ -32,8 +32,8 @@ class TurbopufferVectorStore(VectorStore): Uses Turbopuffer's Python SDK for vector operations. Each namespace corresponds to either: - - A document collection: {prefix}-{workspace}-{observer}-{observed} - - A workspace's message embeddings: {prefix}-{workspace}-messages + - A document collection: {prefix}.{workspace}.{observer}.{observed} + - A workspace's message embeddings: {prefix}.{workspace}.messages """ tpuf: Turbopuffer @@ -211,6 +211,14 @@ class TurbopufferVectorStore(VectorStore): ) return query_results + except NotFoundError: + # Namespace doesn't exist yet - no vectors have been written + # Return empty results (same behavior as LanceDB for missing tables) + logger.debug( + f"Namespace {namespace} does not exist, returning empty results" + ) + return [] + except Exception: logger.exception(f"Failed to query namespace {namespace}") raise @@ -262,6 +270,9 @@ class TurbopufferVectorStore(VectorStore): try: ns.write(deletes=ids) + except NotFoundError: + # Namespace doesn't exist - nothing to delete + logger.debug(f"Namespace {namespace} does not exist, nothing to delete") except Exception: logger.exception( f"Failed to delete {len(ids)} vectors from namespace {namespace}" @@ -280,6 +291,9 @@ class TurbopufferVectorStore(VectorStore): try: ns.delete_all() logger.debug(f"Deleted all vectors from namespace {namespace}") + except NotFoundError: + # Namespace doesn't exist - nothing to delete + logger.debug(f"Namespace {namespace} does not exist, nothing to delete") except Exception: logger.exception(f"Failed to delete namespace {namespace}") raise diff --git a/tests/crud/test_workspace.py b/tests/crud/test_workspace.py index d8ea00a5..f93545b5 100644 --- a/tests/crud/test_workspace.py +++ b/tests/crud/test_workspace.py @@ -199,7 +199,6 @@ class TestWorkspaceCRUD: observed=test_peer.name, session_name=session.name, content="Test document content", - embedding=[0.1] * 1536, # Mock embedding vector ) db_session.add(document) await db_session.flush() @@ -451,7 +450,6 @@ class TestWorkspaceCRUD: observed=peer2.name, session_name=session1.name, content="Test document", - embedding=[0.1] * 1536, # Mock embedding vector ) db_session.add(document) diff --git a/tests/integration/test_message_embeddings.py b/tests/integration/test_message_embeddings.py index 35ef7980..dd51aeb9 100644 --- a/tests/integration/test_message_embeddings.py +++ b/tests/integration/test_message_embeddings.py @@ -64,16 +64,13 @@ async def test_message_embedding_created_when_setting_enabled( result = await db_session.execute(stmt) embedding_record = result.scalar_one_or_none() - # Verify the embedding was created + # Verify the embedding record was created (embedding vectors are now stored externally) assert embedding_record is not None assert embedding_record.message_id == created_message.public_id assert embedding_record.content == test_message_content assert embedding_record.workspace_name == test_workspace.name assert embedding_record.session_name == test_session.name assert embedding_record.peer_name == test_peer.name - # Verify embedding vector exists and is not empty - assert embedding_record.embedding is not None - assert len(embedding_record.embedding) > 0 @pytest.mark.asyncio @@ -176,15 +173,13 @@ async def test_multiple_message_embeddings_created_when_setting_enabled( result = await db_session.execute(stmt) embedding_record = result.scalar_one_or_none() - # Verify the embedding was created + # Verify the embedding record was created (embedding vectors are now stored externally) assert embedding_record is not None assert embedding_record.message_id == created_message.public_id assert embedding_record.content == messages[i].content assert embedding_record.workspace_name == test_workspace.name assert embedding_record.session_name == test_session.name assert embedding_record.peer_name == test_peer.name - assert embedding_record.embedding is not None - assert len(embedding_record.embedding) > 0 @pytest.mark.asyncio @@ -322,7 +317,8 @@ async def test_message_chunking_creates_multiple_embeddings( result = await db_session.execute(stmt) embedding_records = list(result.scalars().all()) - # Verify multiple embeddings were created (one per chunk) + # Verify multiple embedding records were created (one per chunk) + # Embedding vectors are now stored externally in the vector store assert len(embedding_records) == 3 # Should have 3 embeddings for 3 chunks for _, embedding_record in enumerate(embedding_records): @@ -333,7 +329,5 @@ async def test_message_chunking_creates_multiple_embeddings( assert embedding_record.workspace_name == test_workspace.name assert embedding_record.session_name == test_session.name assert embedding_record.peer_name == test_peer.name - assert embedding_record.embedding is not None - assert len(embedding_record.embedding) == 1536 - # Each chunk should have a different embedding vector (0.1, 0.2, 0.3) - assert embedding_record.embedding[0] in [0.1, 0.2, 0.3] + # chunk_index should be set for each chunk + assert embedding_record.chunk_index is not None diff --git a/tests/integration/test_representation.py b/tests/integration/test_representation.py index f3e112b5..5d62b748 100644 --- a/tests/integration/test_representation.py +++ b/tests/integration/test_representation.py @@ -277,7 +277,6 @@ class TestDocumentCreationWorkflow: level="explicit", times_derived=5, internal_metadata={}, - embedding=[0.1] * 1536, ) db_session.add(highly_derived_doc) @@ -291,7 +290,6 @@ class TestDocumentCreationWorkflow: level="explicit", times_derived=2, internal_metadata={}, - embedding=[0.2] * 1536, ) db_session.add(less_derived_doc) @@ -313,7 +311,7 @@ class TestDocumentCreationWorkflow: async def test_representation_from_documents(self): """Test converting documents to representation""" - # Create test documents + # Create test documents (embedding vectors are now stored externally) explicit_doc = models.Document( workspace_name="test_workspace", observer="test_peer", @@ -324,7 +322,6 @@ class TestDocumentCreationWorkflow: "message_ids": [(1, 1)], }, session_name="test_session", - embedding=[0.1] * 1536, created_at=datetime(2025, 1, 1, 10, 0, 0, tzinfo=timezone.utc), ) @@ -339,7 +336,6 @@ class TestDocumentCreationWorkflow: "premises": ["User said they like programming"], }, session_name="test_session", - embedding=[0.2] * 1536, created_at=datetime(2025, 1, 1, 10, 1, 0, tzinfo=timezone.utc), )