agentrun-mem0ai 0.0.11__py3-none-any.whl

agentrun_mem0/vector_stores/supabase.py

@@ -0,0 +1,237 @@
+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 agentrun_mem0.configs.vector_stores.supabase import IndexMeasure, IndexMethod
+from agentrun_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], limit: int = 5, filters: Optional[dict] = None
+    ) -> List[OutputData]:
+        """
+        Search for similar vectors.
+
+        Args:
+            query (str): Query.
+            vectors (List[float]): Query vector.
+            limit (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=limit, filters=filters, include_metadata=True, include_value=True
+        )
+
+        return [OutputData(id=str(result[0]), score=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 []
+
+        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, limit: int = 100) -> List[OutputData]:
+        """
+        List vectors in the collection.
+
+        Args:
+            filters (Dict, optional): Filters to apply
+            limit (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=limit, 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)

agentrun_mem0/vector_stores/upstash_vector.py

@@ -0,0 +1,293 @@
+import logging
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel
+
+from agentrun_mem0.vector_stores.base import VectorStoreBase
+
+try:
+    from upstash_vector import Index
+except ImportError:
+    raise ImportError("The 'upstash_vector' library is required. Please install it using 'pip install upstash_vector'.")
+
+
+logger = logging.getLogger(__name__)
+
+
+class OutputData(BaseModel):
+    id: Optional[str]  # memory id
+    score: Optional[float]  # is None for `get` method
+    payload: Optional[Dict]  # metadata
+
+
+class UpstashVector(VectorStoreBase):
+    def __init__(
+        self,
+        collection_name: str,
+        url: Optional[str] = None,
+        token: Optional[str] = None,
+        client: Optional[Index] = None,
+        enable_embeddings: bool = False,
+    ):
+        """
+        Initialize the UpstashVector vector store.
+
+        Args:
+            url (str, optional): URL for Upstash Vector index. Defaults to None.
+            token (int, optional): Token for Upstash Vector index. Defaults to None.
+            client (Index, optional): Existing `upstash_vector.Index` client instance. Defaults to None.
+            namespace (str, optional): Default namespace for the index. Defaults to None.
+        """
+        if client:
+            self.client = client
+        elif url and token:
+            self.client = Index(url, token)
+        else:
+            raise ValueError("Either a client or URL and token must be provided.")
+
+        self.collection_name = collection_name
+
+        self.enable_embeddings = enable_embeddings
+
+    def insert(
+        self,
+        vectors: List[list],
+        payloads: Optional[List[Dict]] = None,
+        ids: Optional[List[str]] = None,
+    ):
+        """
+        Insert vectors
+
+        Args:
+            vectors (list): List of vectors to insert.
+            payloads (list, optional): List of payloads corresponding to vectors. These will be passed as metadatas to the Upstash Vector client. Defaults to None.
+            ids (list, optional): List of IDs corresponding to vectors. Defaults to None.
+        """
+        logger.info(f"Inserting {len(vectors)} vectors into namespace {self.collection_name}")
+
+        if self.enable_embeddings:
+            if not payloads or any("data" not in m or m["data"] is None for m in payloads):
+                raise ValueError("When embeddings are enabled, all payloads must contain a 'data' field.")
+            processed_vectors = [
+                {
+                    "id": ids[i] if ids else None,
+                    "data": payloads[i]["data"],
+                    "metadata": payloads[i],
+                }
+                for i, v in enumerate(vectors)
+            ]
+        else:
+            processed_vectors = [
+                {
+                    "id": ids[i] if ids else None,
+                    "vector": vectors[i],
+                    "metadata": payloads[i] if payloads else None,
+                }
+                for i, v in enumerate(vectors)
+            ]
+
+        self.client.upsert(
+            vectors=processed_vectors,
+            namespace=self.collection_name,
+        )
+
+    def _stringify(self, x):
+        return f'"{x}"' if isinstance(x, str) else x
+
+    def search(
+        self,
+        query: str,
+        vectors: List[list],
+        limit: int = 5,
+        filters: Optional[Dict] = None,
+    ) -> List[OutputData]:
+        """
+        Search for similar vectors.
+
+        Args:
+            query (list): Query vector.
+            limit (int, optional): Number of results to return. Defaults to 5.
+            filters (Dict, optional): Filters to apply to the search.
+
+        Returns:
+            List[OutputData]: Search results.
+        """
+
+        filters_str = " AND ".join([f"{k} = {self._stringify(v)}" for k, v in filters.items()]) if filters else None
+
+        response = []
+
+        if self.enable_embeddings:
+            response = self.client.query(
+                data=query,
+                top_k=limit,
+                filter=filters_str or "",
+                include_metadata=True,
+                namespace=self.collection_name,
+            )
+        else:
+            queries = [
+                {
+                    "vector": v,
+                    "top_k": limit,
+                    "filter": filters_str or "",
+                    "include_metadata": True,
+                    "namespace": self.collection_name,
+                }
+                for v in vectors
+            ]
+            responses = self.client.query_many(queries=queries)
+            # flatten
+            response = [res for res_list in responses for res in res_list]
+
+        return [
+            OutputData(
+                id=res.id,
+                score=res.score,
+                payload=res.metadata,
+            )
+            for res in response
+        ]
+
+    def delete(self, vector_id: int):
+        """
+        Delete a vector by ID.
+
+        Args:
+            vector_id (int): ID of the vector to delete.
+        """
+        self.client.delete(
+            ids=[str(vector_id)],
+            namespace=self.collection_name,
+        )
+
+    def update(
+        self,
+        vector_id: int,
+        vector: Optional[list] = None,
+        payload: Optional[dict] = None,
+    ):
+        """
+        Update a vector and its payload.
+
+        Args:
+            vector_id (int): ID of the vector to update.
+            vector (list, optional): Updated vector. Defaults to None.
+            payload (dict, optional): Updated payload. Defaults to None.
+        """
+        self.client.update(
+            id=str(vector_id),
+            vector=vector,
+            data=payload.get("data") if payload else None,
+            metadata=payload,
+            namespace=self.collection_name,
+        )
+
+    def get(self, vector_id: int) -> Optional[OutputData]:
+        """
+        Retrieve a vector by ID.
+
+        Args:
+            vector_id (int): ID of the vector to retrieve.
+
+        Returns:
+            dict: Retrieved vector.
+        """
+        response = self.client.fetch(
+            ids=[str(vector_id)],
+            namespace=self.collection_name,
+            include_metadata=True,
+        )
+        if len(response) == 0:
+            return None
+        vector = response[0]
+        if not vector:
+            return None
+        return OutputData(id=vector.id, score=None, payload=vector.metadata)
+
+    def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[List[OutputData]]:
+        """
+        List all memories.
+        Args:
+            filters (Dict, optional): Filters to apply to the search. Defaults to None.
+            limit (int, optional): Number of results to return. Defaults to 100.
+        Returns:
+            List[OutputData]: Search results.
+        """
+        filters_str = " AND ".join([f"{k} = {self._stringify(v)}" for k, v in filters.items()]) if filters else None
+
+        info = self.client.info()
+        ns_info = info.namespaces.get(self.collection_name)
+
+        if not ns_info or ns_info.vector_count == 0:
+            return [[]]
+
+        random_vector = [1.0] * self.client.info().dimension
+
+        results, query = self.client.resumable_query(
+            vector=random_vector,
+            filter=filters_str or "",
+            include_metadata=True,
+            namespace=self.collection_name,
+            top_k=100,
+        )
+        with query:
+            while True:
+                if len(results) >= limit:
+                    break
+                res = query.fetch_next(100)
+                if not res:
+                    break
+                results.extend(res)
+
+        parsed_result = [
+            OutputData(
+                id=res.id,
+                score=res.score,
+                payload=res.metadata,
+            )
+            for res in results
+        ]
+        return [parsed_result]
+
+    def create_col(self, name, vector_size, distance):
+        """
+        Upstash Vector has namespaces instead of collections. A namespace is created when the first vector is inserted.
+
+        This method is a placeholder to maintain the interface.
+        """
+        pass
+
+    def list_cols(self) -> List[str]:
+        """
+        Lists all namespaces in the Upstash Vector index.
+        Returns:
+            List[str]: List of namespaces.
+        """
+        return self.client.list_namespaces()
+
+    def delete_col(self):
+        """
+        Delete the namespace and all vectors in it.
+        """
+        self.client.reset(namespace=self.collection_name)
+        pass
+
+    def col_info(self):
+        """
+        Return general information about the Upstash Vector index.
+
+        - Total number of vectors across all namespaces
+        - Total number of vectors waiting to be indexed across all namespaces
+        - Total size of the index on disk in bytes
+        - Vector dimension
+        - Similarity function used
+        - Per-namespace vector and pending vector counts
+        """
+        return self.client.info()
+
+    def reset(self):
+        """
+        Reset the Upstash Vector index.
+        """
+        self.delete_col()