RomanNum3ral
/
honcho
mirror of https://github.com/plastic-labs/honcho.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

fix: use Async clients

This commit is contained in:
Rajat Ahuja 2025-12-04 17:45:39 -05:00
parent 2178fc91ce
commit 1fb01a9ff8
4 changed files with 131 additions and 63 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

18
src/dependencies.py
View File

@ -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)

59
src/vector_store/__init__.py
View File

@ -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",
]

93
src/vector_store/lancedb.py
View File

@ -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:

24
src/vector_store/turbopuffer.py
View File

@ -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