mem0ai-azure-mysql 0.1.115__py3-none-any.whl

mem0/vector_stores/upstash_vector.py

@@ -0,0 +1,293 @@
+import logging
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel
+
+from 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()