""" LanceDB vector store implementation. This module provides a LanceDB-based implementation of the VectorStore interface. """ import asyncio import logging import re from collections.abc import Sequence from typing import Any, cast import lancedb import pyarrow as pa from lancedb import AsyncConnection, AsyncTable from src.config import settings from src.exceptions import VectorStoreError from . import VectorQueryResult, VectorRecord, VectorStore logger = logging.getLogger(__name__) # Pattern for valid SQL identifiers (alphanumeric + underscore, not starting with digit) _VALID_IDENTIFIER_PATTERN = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$") # Schema for LanceDB tables # id: string, vector: fixed_size_list of float32 (dimension from embedding settings) # Additional metadata columns are added dynamically # pyright: reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownParameterType=false class LanceDBVectorStore(VectorStore): """ LanceDB implementation of the VectorStore interface. Uses LanceDB's async embedded mode for local vector storage. Each namespace corresponds to a LanceDB table. """ _db: AsyncConnection | None = None _db_path: str _db_lock: asyncio.Lock def __init__(self) -> None: """Initialize the LanceDB vector store.""" super().__init__() self._db_path = settings.VECTOR_STORE.LANCEDB_PATH self._db = None self._db_lock = asyncio.Lock() async def _get_db(self) -> AsyncConnection: """Get or create the async database connection (asyncio-safe).""" if self._db is not None: return self._db async with self._db_lock: # Double-check after acquiring lock 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.""" db = await self._get_db() table_names = await db.table_names() if namespace in table_names: return await db.open_table(namespace) return None async def _get_or_create_table( self, namespace: str, ) -> AsyncTable: """ Get existing table or create if not exists. Args: namespace: Table name (namespace) Returns: LanceDB async table """ db = await self._get_db() table_names = await db.table_names() if namespace in table_names: return await db.open_table(namespace) # Create empty table with base schema # Handle race condition: another worker may have created the table # between our check and create_table call fields: list[pa.Field] = [ pa.field("id", pa.string()), pa.field( "vector", pa.list_(pa.float32(), settings.EMBEDDING.VECTOR_DIMENSIONS) ), ] fields.extend(self._metadata_fields_for_namespace(namespace)) schema = pa.schema(fields) try: table = await db.create_table(namespace, schema=schema) # pyright: ignore[reportUnknownArgumentType] return table except Exception: # Table may have been created by another worker, try to open it return await db.open_table(namespace) def _metadata_fields_for_namespace(self, namespace: str) -> list[pa.Field]: """ Infer standard metadata columns based on namespace structure. Namespaces use format: {prefix}.{type}.{hash} - Documents: {prefix}.doc.{hash} - Messages: {prefix}.msg.{hash} """ parts = namespace.split(".") if len(parts) < 3: return [] # Second-to-last part indicates the type (doc or msg) ns_type = parts[-2] if ns_type == "msg": return [ pa.field("message_id", pa.string(), nullable=True), pa.field("session_name", pa.string(), nullable=True), pa.field("peer_name", pa.string(), nullable=True), ] if ns_type == "doc": return [ pa.field("workspace_name", pa.string(), nullable=True), pa.field("observer", pa.string(), nullable=True), pa.field("observed", pa.string(), nullable=True), pa.field("session_name", pa.string(), nullable=True), pa.field("level", pa.string(), nullable=True), ] return [] def _row_to_dict(self, vector: VectorRecord) -> dict[str, Any]: """Convert a VectorRecord to a dict for LanceDB.""" row: dict[str, Any] = { "id": vector.id, "vector": vector.embedding, } # Add metadata fields if vector.metadata: reserved_keys = {"id", "vector", "_distance"} for key in vector.metadata: if key not in reserved_keys: row[key] = vector.metadata[key] return row async def upsert_many( self, namespace: str, vectors: list[VectorRecord], ) -> None: """ Upsert multiple vectors into LanceDB. Args: namespace: The namespace (table) to store the vectors in vectors: List of VectorRecord objects to upsert """ if not vectors: return try: rows = [self._row_to_dict(v) for v in vectors] table = await self._get_or_create_table(namespace) # Use merge_insert for upsert behavior 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}") return except Exception as e: logger.exception( f"Failed to upsert {len(vectors)} vectors to namespace {namespace}" ) raise VectorStoreError( f"Failed to upsert {len(vectors)} vectors to namespace {namespace}" ) from e async def query( self, namespace: str, embedding: list[float], *, top_k: int = 10, filters: dict[str, Any] | None = None, max_distance: float | None = None, ) -> list[VectorQueryResult]: """ Query for similar vectors in LanceDB. Args: namespace: The namespace (table) 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) Returns: List of VectorQueryResult objects, ordered by similarity (most similar first) """ 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 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) # Execute query # LanceDB async API returns list of dicts with incomplete type annotations results = cast(list[dict[str, Any]], await query.to_list()) # Convert to VectorQueryResult objects query_results: list[VectorQueryResult] = [] 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) metadata: dict[str, Any] = { k: v for k, v in row.items() if k not in ("id", "vector", "_distance") } query_results.append( VectorQueryResult( id=str(row["id"]), score=dist, metadata=metadata, ) ) logger.debug( f"Query returned {len(query_results)} results from namespace {namespace}" ) return query_results except Exception: logger.exception(f"Failed to query namespace {namespace}") raise def _build_where_clause(self, filters: dict[str, Any]) -> str | None: """ Convert a filter dict to SQL WHERE clause syntax. Supports filter formats: - {"key": "value"} -> key = 'value' - {"key": {"in": ["a", "b"]}} -> key IN ('a', 'b') Args: filters: Dictionary of attribute -> value filters Returns: SQL WHERE clause string or None if no filters Raises: ValueError: If a filter key is not a valid SQL identifier """ if not filters: return None conditions: list[str] = [] for key, value in filters.items(): # Validate key is a safe SQL identifier to prevent injection if not _VALID_IDENTIFIER_PATTERN.match(key): raise ValueError(f"Invalid filter key: {key!r}") # Check if value is a dict with "in" operator if isinstance(value, dict) and "in" in value: # IN clause for list membership in_values = cast(Sequence[Any], value["in"]) if in_values: escaped_values = [ f"'{str(v).replace(chr(39), chr(39) + chr(39))}'" if isinstance(v, str) else str(v) for v in in_values ] conditions.append(f"{key} IN ({', '.join(escaped_values)})") # Handle string values with proper quoting elif isinstance(value, str): # Escape single quotes in the value escaped_value = value.replace("'", "''") conditions.append(f"{key} = '{escaped_value}'") elif isinstance(value, bool): conditions.append(f"{key} = {str(value).lower()}") elif value is None: conditions.append(f"{key} IS NULL") else: conditions.append(f"{key} = {value}") return " AND ".join(conditions) if conditions else None async def delete_many(self, namespace: str, ids: list[str]) -> None: """ Delete multiple vectors from LanceDB. Args: namespace: The namespace (table) containing the vectors ids: List of vector identifiers to delete """ if not ids: return table = await self._get_table(namespace) if table is None: logger.debug(f"Table {namespace} does not exist, nothing to delete") return try: # 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) await table.delete(f"id IN ({in_clause})") logger.debug(f"Deleted {len(ids)} vectors from namespace {namespace}") 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 (table) and all its vectors from LanceDB. Args: namespace: The namespace (table) to delete """ try: 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: logger.exception(f"Failed to delete namespace {namespace}") raise async def close(self) -> None: """Close the LanceDB connection and release resources.""" if self._db is not None: # AsyncConnection provides an explicit close() method (synchronous) # which we invoke to ensure proper cleanup of resources self._db.close() self._db = None logger.debug("LanceDB connection closed")