import logging import uuid from typing import List, Optional from pydantic import BaseModel try: import vecs except ImportError: raise ImportError("The 'vecs' library is required. Please install it using 'pip install vecs'.") from mem0.configs.vector_stores.supabase import IndexMeasure, IndexMethod from mem0.vector_stores.base import VectorStoreBase logger = logging.getLogger(__name__) class OutputData(BaseModel): id: Optional[str] score: Optional[float] payload: Optional[dict] class Supabase(VectorStoreBase): def __init__( self, connection_string: str, collection_name: str, embedding_model_dims: int, index_method: IndexMethod = IndexMethod.AUTO, index_measure: IndexMeasure = IndexMeasure.COSINE, ): """ Initialize the Supabase vector store using vecs. Args: connection_string (str): PostgreSQL connection string collection_name (str): Collection name embedding_model_dims (int): Dimension of the embedding vector index_method (IndexMethod): Index method to use. Defaults to AUTO. index_measure (IndexMeasure): Distance measure to use. Defaults to COSINE. """ self.db = vecs.create_client(connection_string) self.collection_name = collection_name self.embedding_model_dims = embedding_model_dims self.index_method = index_method self.index_measure = index_measure collections = self.list_cols() if collection_name not in collections: self.create_col(embedding_model_dims) def _preprocess_filters(self, filters: Optional[dict] = None) -> Optional[dict]: """ Preprocess filters to be compatible with vecs. Args: filters (Dict, optional): Filters to preprocess. Multiple filters will be combined with AND logic. """ if filters is None: return None if len(filters) == 1: # For single filter, keep the simple format key, value = next(iter(filters.items())) return {key: {"$eq": value}} # For multiple filters, use $and clause return {"$and": [{key: {"$eq": value}} for key, value in filters.items()]} def create_col(self, embedding_model_dims: Optional[int] = None) -> None: """ Create a new collection with vector support. Will also initialize vector search index. Args: embedding_model_dims (int, optional): Dimension of the embedding vector. If not provided, uses the dimension specified in initialization. """ dims = embedding_model_dims or self.embedding_model_dims if not dims: raise ValueError( "embedding_model_dims must be provided either during initialization or when creating collection" ) logger.info(f"Creating new collection: {self.collection_name}") try: self.collection = self.db.get_or_create_collection(name=self.collection_name, dimension=dims) self.collection.create_index(method=self.index_method.value, measure=self.index_measure.value) logger.info(f"Successfully created collection {self.collection_name} with dimension {dims}") except Exception as e: logger.error(f"Failed to create collection: {str(e)}") raise def insert( self, vectors: List[List[float]], payloads: Optional[List[dict]] = None, ids: Optional[List[str]] = None ): """ Insert vectors into the collection. Args: vectors (List[List[float]]): List of vectors to insert payloads (List[Dict], optional): List of payloads corresponding to vectors ids (List[str], optional): List of IDs corresponding to vectors """ logger.info(f"Inserting {len(vectors)} vectors into collection {self.collection_name}") if not ids: ids = [str(uuid.uuid4()) for _ in vectors] if not payloads: payloads = [{} for _ in vectors] records = [(id, vector, payload) for id, vector, payload in zip(ids, vectors, payloads)] self.collection.upsert(records) def search( self, query: str, vectors: List[float], top_k: int = 5, filters: Optional[dict] = None ) -> List[OutputData]: """ Search for similar vectors. Args: query (str): Query. vectors (List[float]): Query vector. top_k (int, optional): Number of results to return. Defaults to 5. filters (Dict, optional): Filters to apply to the search. Defaults to None. Returns: List[OutputData]: Search results """ filters = self._preprocess_filters(filters) results = self.collection.query( data=vectors, limit=top_k, filters=filters, include_metadata=True, include_value=True ) return [OutputData(id=str(result[0]), score=max(0.0, 1.0 - float(result[1])), payload=result[2]) for result in results] def delete(self, vector_id: str): """ Delete a vector by ID. Args: vector_id (str): ID of the vector to delete """ self.collection.delete([(vector_id,)]) def update(self, vector_id: str, vector: Optional[List[float]] = None, payload: Optional[dict] = None): """ Update a vector and/or its payload. Args: vector_id (str): ID of the vector to update vector (List[float], optional): Updated vector payload (Dict, optional): Updated payload """ if vector is None: # If only updating metadata, we need to get the existing vector existing = self.get(vector_id) if existing and existing.payload: vector = existing.payload.get("vector", []) if vector: self.collection.upsert([(vector_id, vector, payload or {})]) def get(self, vector_id: str) -> Optional[OutputData]: """ Retrieve a vector by ID. Args: vector_id (str): ID of the vector to retrieve Returns: Optional[OutputData]: Retrieved vector data or None if not found """ result = self.collection.fetch([(vector_id,)]) if not result: return None record = result[0] return OutputData(id=str(record.id), score=None, payload=record.metadata) def list_cols(self) -> List[str]: """ List all collections. Returns: List[str]: List of collection names """ return self.db.list_collections() def delete_col(self): """Delete the collection.""" self.db.delete_collection(self.collection_name) def col_info(self) -> dict: """ Get information about the collection. Returns: Dict: Collection information including name and configuration """ info = self.collection.describe() return { "name": info.name, "count": info.vectors, "dimension": info.dimension, "index": {"method": info.index_method, "metric": info.distance_metric}, } def list(self, filters: Optional[dict] = None, top_k: int = 100) -> List[OutputData]: """ List vectors in the collection. Args: filters (Dict, optional): Filters to apply top_k (int, optional): Maximum number of results to return. Defaults to 100. Returns: List[OutputData]: List of vectors """ filters = self._preprocess_filters(filters) query = [0] * self.embedding_model_dims ids = self.collection.query( data=query, limit=top_k, filters=filters, include_metadata=True, include_value=False ) ids = [id[0] for id in ids] records = self.collection.fetch(ids=ids) return [[OutputData(id=str(record[0]), score=None, payload=record[2]) for record in records]] def reset(self): """Reset the index by deleting and recreating it.""" logger.warning(f"Resetting index {self.collection_name}...") self.delete_col() self.create_col(self.embedding_model_dims)