From 1fb01a9ff809802869437b39dd2a55775d9672dc Mon Sep 17 00:00:00 2001 From: Rajat Ahuja Date: Thu, 4 Dec 2025 17:45:39 -0500 Subject: [PATCH] fix: use Async clients --- src/dependencies.py | 18 +++++++ src/vector_store/__init__.py | 59 +++++++++++++++------ src/vector_store/lancedb.py | 93 ++++++++++++++++++++------------- src/vector_store/turbopuffer.py | 24 ++++----- 4 files changed, 131 insertions(+), 63 deletions(-) diff --git a/src/dependencies.py b/src/dependencies.py index 77699150..c83a6685 100644 --- a/src/dependencies.py +++ b/src/dependencies.py @@ -1,5 +1,6 @@ import uuid from contextlib import asynccontextmanager +from typing import TYPE_CHECKING from fastapi import Depends from sqlalchemy import text @@ -8,6 +9,9 @@ from sqlalchemy.ext.asyncio import AsyncSession from src.config import settings from src.db import SessionLocal, request_context +if TYPE_CHECKING: + from src.vector_store import VectorStore + async def get_db(): """FastAPI Dependency Generator for Database""" @@ -61,3 +65,17 @@ async def tracked_db(operation_name: str | None = None): db: AsyncSession = Depends(get_db) + + +def get_vector_store_dep() -> "VectorStore": + """FastAPI dependency for vector store. + + This is a thin wrapper around get_vector_store() to allow for + proper dependency injection in FastAPI routes. + """ + from src.vector_store import get_vector_store + + return get_vector_store() + + +vector_store: "VectorStore" = Depends(get_vector_store_dep) diff --git a/src/vector_store/__init__.py b/src/vector_store/__init__.py index 0b3e0cf5..3f4cf2fb 100644 --- a/src/vector_store/__init__.py +++ b/src/vector_store/__init__.py @@ -161,9 +161,38 @@ class VectorStore(ABC): _vector_store_instance: VectorStore | None = None +def _create_vector_store() -> VectorStore: + """ + Create a new vector store instance based on configuration. + + Returns: + The vector store instance based on configuration. + + Raises: + ValueError: If the configured vector store type is invalid. + """ + store_type = settings.VECTOR_STORE.TYPE + + if store_type == "turbopuffer": + from src.vector_store.turbopuffer import TurbopufferVectorStore + + return TurbopufferVectorStore() + elif store_type == "lancedb": + from src.vector_store.lancedb import LanceDBVectorStore + + return LanceDBVectorStore() + else: + raise ValueError(f"Unknown vector store type: {store_type}") + + def get_vector_store() -> VectorStore: """ - Get the configured vector store instance (singleton). + FastAPI dependency that provides the configured vector store instance (singleton). + + This function is designed to be used as a FastAPI dependency: + vector_store: VectorStore = Depends(get_vector_store) + + It can also be called directly for non-request contexts (e.g., background tasks). Returns: The vector store instance based on configuration. @@ -173,28 +202,26 @@ def get_vector_store() -> VectorStore: """ global _vector_store_instance - if _vector_store_instance is not None: - return _vector_store_instance - - store_type = settings.VECTOR_STORE.TYPE - - if store_type == "turbopuffer": - from src.vector_store.turbopuffer import TurbopufferVectorStore - - _vector_store_instance = TurbopufferVectorStore() - elif store_type == "lancedb": - from src.vector_store.lancedb import LanceDBVectorStore - - _vector_store_instance = LanceDBVectorStore() - else: - raise ValueError(f"Unknown vector store type: {store_type}") + if _vector_store_instance is None: + _vector_store_instance = _create_vector_store() return _vector_store_instance +def reset_vector_store() -> None: + """ + Reset the vector store singleton instance. + + This is primarily useful for testing to ensure a fresh instance is created. + """ + global _vector_store_instance + _vector_store_instance = None + + __all__ = [ "VectorStore", "VectorRecord", "QueryResult", "get_vector_store", + "reset_vector_store", ] diff --git a/src/vector_store/lancedb.py b/src/vector_store/lancedb.py index c145ae5d..95ea5691 100644 --- a/src/vector_store/lancedb.py +++ b/src/vector_store/lancedb.py @@ -6,10 +6,11 @@ for use in self-hosted deployments of Honcho. """ import logging -from typing import Any +from typing import Any, cast import lancedb import pyarrow as pa +from lancedb import AsyncConnection, AsyncTable from src.config import settings @@ -27,27 +28,37 @@ class LanceDBVectorStore(VectorStore): """ LanceDB implementation of the VectorStore interface. - Uses LanceDB's embedded mode for local vector storage. + Uses LanceDB's async embedded mode for local vector storage. Each namespace corresponds to a LanceDB table. """ - _db: lancedb.DBConnection + _db: AsyncConnection | None = None + _db_path: str def __init__(self): """Initialize the LanceDB vector store.""" super().__init__() - self._db = lancedb.connect(settings.VECTOR_STORE.LANCEDB_PATH) + self._db_path = settings.VECTOR_STORE.LANCEDB_PATH + self._db = None - def _get_table(self, namespace: str) -> lancedb.table.Table | None: + async def _get_db(self) -> AsyncConnection: + """Get or create the async database connection.""" + if self._db is None: + self._db = await lancedb.connect_async(self._db_path) + return self._db + + async def _get_table(self, namespace: str) -> AsyncTable | None: """Get a table if it exists, otherwise return None.""" - if namespace in self._db.table_names(): - return self._db.open_table(namespace) + db = await self._get_db() + table_names = await db.table_names() + if namespace in table_names: + return await db.open_table(namespace) return None - def _get_or_create_table( + async def _get_or_create_table( self, namespace: str, sample_data: list[dict[str, Any]] | None = None - ) -> lancedb.table.Table: - """_get_or_create_table + ) -> AsyncTable: + """ Get existing table or create if not exists. Args: @@ -55,14 +66,16 @@ class LanceDBVectorStore(VectorStore): sample_data: Optional sample data to infer schema from Returns: - LanceDB table + LanceDB async table """ - if namespace in self._db.table_names(): - return self._db.open_table(namespace) + db = await self._get_db() + table_names = await db.table_names() + if namespace in table_names: + return await db.open_table(namespace) # Create table with sample data if provided if sample_data: - return self._db.create_table(namespace, data=sample_data) + return await db.create_table(namespace, data=sample_data) # Create empty table with base schema schema = pa.schema( # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] @@ -71,7 +84,7 @@ class LanceDBVectorStore(VectorStore): pa.field("vector", pa.list_(pa.float32(), VECTOR_DIMENSION)), # pyright: ignore[reportUnknownMemberType] ] ) - return self._db.create_table(namespace, schema=schema) # pyright: ignore[reportUnknownArgumentType] + return await db.create_table(namespace, schema=schema) # pyright: ignore[reportUnknownArgumentType] def _row_to_dict(self, vector: VectorRecord) -> dict[str, Any]: """Convert a VectorRecord to a dict for LanceDB.""" @@ -98,10 +111,15 @@ class LanceDBVectorStore(VectorStore): """ try: row = self._row_to_dict(vector) - table = self._get_or_create_table(namespace, sample_data=[row]) + table = await self._get_or_create_table(namespace, sample_data=[row]) # Use merge_insert for upsert behavior - table.merge_insert("id").when_matched_update_all().execute([row]) + await ( + table.merge_insert("id") + .when_matched_update_all() + .when_not_matched_insert_all() + .execute([row]) + ) logger.debug(f"Upserted vector {vector.id} to namespace {namespace}") except Exception: @@ -127,12 +145,15 @@ class LanceDBVectorStore(VectorStore): try: rows = [self._row_to_dict(v) for v in vectors] - table = self._get_or_create_table(namespace, sample_data=rows) + table = await self._get_or_create_table(namespace, sample_data=rows) # Use merge_insert for upsert behavior - table.merge_insert( - "id" - ).when_matched_update_all().when_not_matched_insert_all().execute(rows) + await ( + 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: @@ -163,44 +184,44 @@ class LanceDBVectorStore(VectorStore): Returns: List of QueryResult objects, ordered by similarity (most similar first) """ - table = self._get_table(namespace) + table = await self._get_table(namespace) if table is None: logger.debug(f"Table {namespace} does not exist, returning empty results") return [] try: - # Build query (LanceDB types are incomplete, so type checker reports false positive) - query = table.search(embedding).distance_type("cosine").limit(top_k) # pyright: ignore[reportAttributeAccessIssue, reportUnknownVariableType, reportUnknownMemberType] + # Build query + query = table.vector_search(embedding).distance_type("cosine").limit(top_k) # Apply filters if provided if filters: where_clause = self._build_where_clause(filters) if where_clause: - query = query.where(where_clause) # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] + query = query.where(where_clause) # Execute query - results = query.to_list() # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType] + # LanceDB async API returns list of dicts with incomplete type annotations + results = cast(list[dict[str, Any]], await query.to_list()) # Convert to QueryResult objects query_results: list[QueryResult] = [] - for row in results: # pyright: ignore[reportUnknownVariableType] - dist = float(row.get("_distance", 0.0)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] + for row in results: + dist = float(row.get("_distance", 0.0)) # Filter by max_distance if specified if max_distance is not None and dist > max_distance: continue # Extract metadata (everything except id, vector, _distance) - # Type annotations for dict comprehension to satisfy type checker metadata: dict[str, Any] = { k: v - for k, v in row.items() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] + for k, v in row.items() if k not in ("id", "vector", "_distance") } query_results.append( QueryResult( - id=str(row["id"]), # pyright: ignore[reportUnknownArgumentType] + id=str(row["id"]), score=dist, metadata=metadata, ) @@ -255,7 +276,7 @@ class LanceDBVectorStore(VectorStore): if not ids: return - table = self._get_table(namespace) + table = await self._get_table(namespace) if table is None: logger.debug(f"Table {namespace} does not exist, nothing to delete") return @@ -264,7 +285,7 @@ class LanceDBVectorStore(VectorStore): # Build IN clause with properly escaped IDs escaped_ids = [f"'{id.replace(chr(39), chr(39) + chr(39))}'" for id in ids] in_clause = ", ".join(escaped_ids) - table.delete(f"id IN ({in_clause})") + await table.delete(f"id IN ({in_clause})") logger.debug(f"Deleted {len(ids)} vectors from namespace {namespace}") except Exception: logger.exception( @@ -280,8 +301,10 @@ class LanceDBVectorStore(VectorStore): namespace: The namespace (table) to delete """ try: - if namespace in self._db.table_names(): - self._db.drop_table(namespace) + db = await self._get_db() + table_names = await db.table_names() + if namespace in table_names: + await db.drop_table(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 58fd4f07..d6676612 100644 --- a/src/vector_store/turbopuffer.py +++ b/src/vector_store/turbopuffer.py @@ -9,8 +9,8 @@ import logging from collections.abc import Sequence from typing import Any, Literal -from turbopuffer import NotFoundError, Turbopuffer -from turbopuffer.lib.namespace import Namespace +from turbopuffer import AsyncTurbopuffer, NotFoundError +from turbopuffer.lib.namespace import AsyncNamespace from turbopuffer.types import Filter from src.config import settings @@ -30,13 +30,13 @@ class TurbopufferVectorStore(VectorStore): """ Turbopuffer implementation of the VectorStore interface. - Uses Turbopuffer's Python SDK for vector operations. + Uses Turbopuffer's async 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 """ - tpuf: Turbopuffer + tpuf: AsyncTurbopuffer def __init__(self): """ @@ -51,12 +51,12 @@ class TurbopufferVectorStore(VectorStore): "VECTOR_STORE_TURBOPUFFER_API_KEY must be set for Turbopuffer vector store" ) - # Initialize the Turbopuffer client + # Initialize the async Turbopuffer client # Region can be configured via VECTOR_STORE_TURBOPUFFER_REGION or TURBOPUFFER_REGION env var region = settings.VECTOR_STORE.TURBOPUFFER_REGION or "gcp-us-east4" - self.tpuf = Turbopuffer(api_key=api_key, region=region) + self.tpuf = AsyncTurbopuffer(api_key=api_key, region=region) - def _get_namespace(self, namespace: str) -> Namespace: + def _get_namespace(self, namespace: str) -> AsyncNamespace: """Get a Turbopuffer namespace object.""" return self.tpuf.namespace(namespace) @@ -85,7 +85,7 @@ class TurbopufferVectorStore(VectorStore): **attributes, } - ns.write( + await ns.write( upsert_rows=[row], distance_metric=DISTANCE_METRIC, ) @@ -122,7 +122,7 @@ class TurbopufferVectorStore(VectorStore): ] try: - ns.write( + await ns.write( upsert_rows=rows, distance_metric=DISTANCE_METRIC, ) @@ -178,7 +178,7 @@ class TurbopufferVectorStore(VectorStore): if filter_condition is not None: query_kwargs["filters"] = filter_condition - response = ns.query(**query_kwargs) + response = await ns.query(**query_kwargs) query_results: list[QueryResult] = [] for row in response.rows or []: @@ -269,7 +269,7 @@ class TurbopufferVectorStore(VectorStore): ns = self._get_namespace(namespace) try: - ns.write(deletes=ids) + await ns.write(deletes=ids) except NotFoundError: # Namespace doesn't exist - nothing to delete logger.debug(f"Namespace {namespace} does not exist, nothing to delete") @@ -289,7 +289,7 @@ class TurbopufferVectorStore(VectorStore): ns = self._get_namespace(namespace) try: - ns.delete_all() + await ns.delete_all() logger.debug(f"Deleted all vectors from namespace {namespace}") except NotFoundError: # Namespace doesn't exist - nothing to delete