368 lines
13 KiB
Python
368 lines
13 KiB
Python
"""
|
|
Turbopuffer vector store implementation.
|
|
|
|
This module provides a Turbopuffer-based implementation of the VectorStore interface.
|
|
"""
|
|
|
|
import logging
|
|
import re
|
|
from collections.abc import Sequence
|
|
from typing import Any, Literal, cast
|
|
|
|
from turbopuffer import AsyncTurbopuffer, InternalServerError, NotFoundError
|
|
from turbopuffer.lib.namespace import AsyncNamespace
|
|
from turbopuffer.types import Filter, RowParam
|
|
|
|
from src.config import settings
|
|
from src.exceptions import VectorStoreError
|
|
|
|
from . import VectorQueryResult, VectorRecord, VectorStore
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Type aliases for Turbopuffer's filter formats
|
|
EqFilter = tuple[str, Literal["Eq"], Any]
|
|
InFilter = tuple[str, Literal["In"], Sequence[Any]]
|
|
AndFilter = tuple[Literal["And"], Sequence[Filter]]
|
|
|
|
DISTANCE_METRIC = "cosine_distance"
|
|
|
|
|
|
class TurbopufferVectorStore(VectorStore):
|
|
"""
|
|
Turbopuffer implementation of the VectorStore interface.
|
|
|
|
Namespaces are generated via get_vector_namespace() which hashes
|
|
the variable components to fit Turbopuffer's [A-Za-z0-9-_.]{1,128} limit.
|
|
"""
|
|
|
|
tpuf: AsyncTurbopuffer
|
|
|
|
def __init__(self):
|
|
"""
|
|
Initialize the Turbopuffer vector store.
|
|
"""
|
|
super().__init__()
|
|
|
|
# Configure Turbopuffer client
|
|
api_key = settings.VECTOR_STORE.TURBOPUFFER_API_KEY
|
|
if not api_key:
|
|
raise ValueError(
|
|
"VECTOR_STORE_TURBOPUFFER_API_KEY must be set for Turbopuffer vector store"
|
|
)
|
|
|
|
# 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 = AsyncTurbopuffer(api_key=api_key, region=region)
|
|
|
|
def _get_namespace(self, namespace: str) -> AsyncNamespace:
|
|
"""Get a Turbopuffer namespace object."""
|
|
return self.tpuf.namespace(namespace)
|
|
|
|
async def upsert_many(
|
|
self,
|
|
namespace: str,
|
|
vectors: list[VectorRecord],
|
|
) -> None:
|
|
"""
|
|
Upsert multiple vectors into Turbopuffer.
|
|
|
|
Args:
|
|
namespace: The namespace to store the vectors in
|
|
vectors: List of VectorRecord objects to upsert
|
|
"""
|
|
if not vectors:
|
|
return
|
|
|
|
ns = self._get_namespace(namespace)
|
|
|
|
# The dict literal carries arbitrary metadata fields, which RowParam supports
|
|
# via extra_items=object. basedpyright can't see through the spread, so cast
|
|
# via object per its reportInvalidCast guidance.
|
|
# Spread metadata first so a caller-supplied "id" or "vector" key
|
|
# can never clobber the required upsert fields.
|
|
rows: list[RowParam] = [
|
|
cast(
|
|
RowParam,
|
|
cast(
|
|
object,
|
|
{
|
|
**(v.metadata or {}),
|
|
"id": v.id,
|
|
"vector": v.embedding,
|
|
},
|
|
),
|
|
)
|
|
for v in vectors
|
|
]
|
|
|
|
try:
|
|
await ns.write(
|
|
upsert_rows=cast(Any, rows),
|
|
distance_metric=DISTANCE_METRIC,
|
|
)
|
|
return
|
|
except InternalServerError as exc:
|
|
# Turbopuffer unavailable. SDK implicitly retries 5xx responses,
|
|
# so raise a vector store error and let callers leave writes unsynced.
|
|
logger.warning(
|
|
"Turbopuffer unavailable for upsert to namespace %s (%s after retries)",
|
|
namespace,
|
|
exc.status_code,
|
|
)
|
|
raise VectorStoreError(
|
|
f"Turbopuffer unavailable for upsert to namespace {namespace}"
|
|
) from exc
|
|
except Exception:
|
|
logger.exception(
|
|
f"Failed to upsert {len(vectors)} vectors to namespace {namespace}"
|
|
)
|
|
raise
|
|
|
|
async def query(
|
|
self,
|
|
namespace: str,
|
|
embedding: list[float],
|
|
*,
|
|
top_k: int = 10,
|
|
filters: dict[str, Any] | None = None,
|
|
max_distance: float | None = None,
|
|
include_attributes: bool | list[str] = True,
|
|
) -> list[VectorQueryResult]:
|
|
"""
|
|
Query for similar vectors in Turbopuffer.
|
|
|
|
Args:
|
|
namespace: The namespace to query
|
|
embedding: The query embedding vector
|
|
top_k: Maximum number of results to return
|
|
filters: Optional metadata filters
|
|
max_distance: Optional maximum distance threshold (cosine distance)
|
|
include_attributes: Attributes to include in the response. Passing False
|
|
avoids parsing unused row attributes.
|
|
|
|
Returns:
|
|
List of VectorQueryResult objects, ordered by similarity (most similar first)
|
|
"""
|
|
ns = self._get_namespace(namespace)
|
|
|
|
try:
|
|
# Build filter conditions for Turbopuffer
|
|
filter_condition = self._build_filters(filters) if filters else None
|
|
|
|
# Query using rank_by for vector similarity
|
|
# rank_by must be a tuple: (attribute, "ANN", vector)
|
|
rank_by: tuple[str, Literal["ANN"], Sequence[float]] = (
|
|
"vector",
|
|
"ANN",
|
|
embedding,
|
|
)
|
|
|
|
# Only pass filters if we have them (avoid passing None)
|
|
query_kwargs: dict[str, Any] = {
|
|
"rank_by": rank_by,
|
|
"top_k": top_k,
|
|
"distance_metric": DISTANCE_METRIC,
|
|
"include_attributes": include_attributes,
|
|
}
|
|
if filter_condition is not None:
|
|
query_kwargs["filters"] = filter_condition
|
|
|
|
response = await ns.query(**query_kwargs)
|
|
|
|
query_results: list[VectorQueryResult] = []
|
|
for row in response.rows or []:
|
|
# Distance is accessed via row["$dist"]
|
|
dist: float = float(row["$dist"]) if "$dist" in row else 0.0
|
|
# Filter by max_distance if specified
|
|
if max_distance is not None and dist > max_distance:
|
|
continue
|
|
|
|
# Extract attributes from model_extra (excludes id, vector, $dist)
|
|
row_metadata: dict[str, Any] = {}
|
|
if row.model_extra:
|
|
# Filter out internal fields like $dist
|
|
row_metadata = {
|
|
k: v
|
|
for k, v in row.model_extra.items()
|
|
if not k.startswith("$")
|
|
}
|
|
|
|
query_results.append(
|
|
VectorQueryResult(
|
|
id=str(row.id),
|
|
score=dist,
|
|
metadata=row_metadata,
|
|
)
|
|
)
|
|
|
|
logger.debug(
|
|
f"Query returned {len(query_results)} results from namespace {namespace}"
|
|
)
|
|
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 InternalServerError as exc:
|
|
# Turbopuffer unavailable. SDK implicitly retries 5xx responses,
|
|
# so we should return [].
|
|
logger.warning(
|
|
"Turbopuffer unavailable for query on namespace %s (%s after retries), returning empty results",
|
|
namespace,
|
|
exc.status_code,
|
|
)
|
|
return []
|
|
|
|
except Exception:
|
|
logger.exception(f"Failed to query namespace {namespace}")
|
|
raise
|
|
|
|
def _build_filters(self, filters: dict[str, Any]) -> Filter | None:
|
|
"""
|
|
Convert a filter dict to Turbopuffer filter format.
|
|
|
|
Turbopuffer uses tuples like (attribute, "Eq", value) for filters,
|
|
(attribute, "In", [values]) for membership filters,
|
|
and ("And", [filters]) for combining multiple filters.
|
|
|
|
Supports filter formats:
|
|
- {"key": "value"} -> ("key", "Eq", "value")
|
|
- {"key": {"in": ["a", "b"]}} -> ("key", "In", ["a", "b"])
|
|
|
|
Args:
|
|
filters: Dictionary of attribute -> value filters
|
|
|
|
Returns:
|
|
Turbopuffer Filter or None if no filters
|
|
"""
|
|
if not filters:
|
|
return None
|
|
|
|
filter_list: list[EqFilter | InFilter] = []
|
|
for key, value in filters.items():
|
|
# Check if value is a dict with "in" operator
|
|
if isinstance(value, dict) and "in" in value:
|
|
# Membership filter using "In" operator
|
|
in_values = cast(Sequence[Any], value["in"])
|
|
filter_list.append((key, "In", in_values))
|
|
else:
|
|
# Simple equality filter using "Eq" operator
|
|
filter_list.append((key, "Eq", cast(Any, value)))
|
|
|
|
if not filter_list:
|
|
return None
|
|
|
|
if len(filter_list) == 1:
|
|
return filter_list[0]
|
|
|
|
# Combine multiple filters with AND
|
|
and_filter: AndFilter = ("And", filter_list)
|
|
return and_filter
|
|
|
|
async def delete_many(self, namespace: str, ids: list[str]) -> None:
|
|
"""
|
|
Delete multiple vectors from Turbopuffer.
|
|
|
|
Args:
|
|
namespace: The namespace containing the vectors
|
|
ids: List of vector identifiers to delete
|
|
"""
|
|
if not ids:
|
|
return
|
|
|
|
ns = self._get_namespace(namespace)
|
|
|
|
try:
|
|
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")
|
|
except InternalServerError as exc:
|
|
logger.warning(
|
|
"Turbopuffer unavailable for delete from namespace %s (%s after retries)",
|
|
namespace,
|
|
exc.status_code,
|
|
)
|
|
raise VectorStoreError(
|
|
f"Turbopuffer unavailable while deleting vectors in namespace {namespace}"
|
|
) from exc
|
|
except Exception:
|
|
logger.exception(
|
|
f"Failed to delete {len(ids)} vectors from namespace {namespace}"
|
|
)
|
|
raise
|
|
|
|
async def delete_namespace(self, namespace: str) -> None:
|
|
"""
|
|
Delete an entire namespace and all its vectors from Turbopuffer.
|
|
|
|
Args:
|
|
namespace: The namespace to delete
|
|
"""
|
|
ns = self._get_namespace(namespace)
|
|
|
|
try:
|
|
await 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
|
|
|
|
async def close(self) -> None:
|
|
"""Close the Turbopuffer client and release resources."""
|
|
await self.tpuf.close()
|
|
logger.debug("Turbopuffer client closed")
|
|
|
|
async def probe_namespace_dim(self, namespace: str) -> int | None:
|
|
"""Inspect a Turbopuffer namespace schema to recover the vector dim.
|
|
|
|
Turbopuffer namespaces are lazy-created; ``namespace.exists()`` returns
|
|
False before the first write. The schema response maps attribute name
|
|
to ``AttributeSchemaConfig``; the vector field's ``type`` string is
|
|
a bracket-prefixed dim with a width suffix, e.g. ``"[768]f32"``,
|
|
``"[1536]f16"``, ``"[256]i8"``.
|
|
|
|
Returns ``None`` only when the namespace does not exist yet
|
|
(NotFoundError or ``exists() == False``). When the namespace
|
|
exists but its schema lacks a parseable ``vector`` attribute,
|
|
raises ``VectorStoreError`` — silently bucketing that as "missing"
|
|
would let a corrupt namespace pass the startup validator.
|
|
"""
|
|
ns = self._get_namespace(namespace)
|
|
try:
|
|
if not await ns.exists():
|
|
return None
|
|
except NotFoundError:
|
|
return None
|
|
|
|
try:
|
|
schema = await ns.schema()
|
|
except NotFoundError:
|
|
return None
|
|
|
|
vector_attr = schema.get("vector")
|
|
if vector_attr is None:
|
|
raise VectorStoreError(
|
|
f"Turbopuffer namespace {namespace!r} exists but its schema"
|
|
+ " has no 'vector' attribute; cannot probe dim."
|
|
)
|
|
type_str = str(vector_attr.type)
|
|
match = re.search(r"\[(\d+)\]", type_str)
|
|
if match is None:
|
|
raise VectorStoreError(
|
|
f"Turbopuffer namespace {namespace!r} has an unparseable"
|
|
+ f" vector type {type_str!r}; expected `[<dim>]<width>`"
|
|
+ " (e.g. `[768]f32`). SDK format may have changed."
|
|
)
|
|
return int(match.group(1))
|