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

mem0/vector_stores/redis.py

@@ -0,0 +1,295 @@
+import json
+import logging
+from datetime import datetime
+from functools import reduce
+
+import numpy as np
+import pytz
+import redis
+from redis.commands.search.query import Query
+from redisvl.index import SearchIndex
+from redisvl.query import VectorQuery
+from redisvl.query.filter import Tag
+
+from mem0.memory.utils import extract_json
+from mem0.vector_stores.base import VectorStoreBase
+
+logger = logging.getLogger(__name__)
+
+# TODO: Improve as these are not the best fields for the Redis's perspective. Might do away with them.
+DEFAULT_FIELDS = [
+    {"name": "memory_id", "type": "tag"},
+    {"name": "hash", "type": "tag"},
+    {"name": "agent_id", "type": "tag"},
+    {"name": "run_id", "type": "tag"},
+    {"name": "user_id", "type": "tag"},
+    {"name": "memory", "type": "text"},
+    {"name": "metadata", "type": "text"},
+    # TODO: Although it is numeric but also accepts string
+    {"name": "created_at", "type": "numeric"},
+    {"name": "updated_at", "type": "numeric"},
+    {
+        "name": "embedding",
+        "type": "vector",
+        "attrs": {"distance_metric": "cosine", "algorithm": "flat", "datatype": "float32"},
+    },
+]
+
+excluded_keys = {"user_id", "agent_id", "run_id", "hash", "data", "created_at", "updated_at"}
+
+
+class MemoryResult:
+    def __init__(self, id: str, payload: dict, score: float = None):
+        self.id = id
+        self.payload = payload
+        self.score = score
+
+
+class RedisDB(VectorStoreBase):
+    def __init__(
+        self,
+        redis_url: str,
+        collection_name: str,
+        embedding_model_dims: int,
+    ):
+        """
+        Initialize the Redis vector store.
+
+        Args:
+            redis_url (str): Redis URL.
+            collection_name (str): Collection name.
+            embedding_model_dims (int): Embedding model dimensions.
+        """
+        self.embedding_model_dims = embedding_model_dims
+        index_schema = {
+            "name": collection_name,
+            "prefix": f"mem0:{collection_name}",
+        }
+
+        fields = DEFAULT_FIELDS.copy()
+        fields[-1]["attrs"]["dims"] = embedding_model_dims
+
+        self.schema = {"index": index_schema, "fields": fields}
+
+        self.client = redis.Redis.from_url(redis_url)
+        self.index = SearchIndex.from_dict(self.schema)
+        self.index.set_client(self.client)
+        self.index.create(overwrite=True)
+
+    def create_col(self, name=None, vector_size=None, distance=None):
+        """
+        Create a new collection (index) in Redis.
+
+        Args:
+            name (str, optional): Name for the collection. Defaults to None, which uses the current collection_name.
+            vector_size (int, optional): Size of the vector embeddings. Defaults to None, which uses the current embedding_model_dims.
+            distance (str, optional): Distance metric to use. Defaults to None, which uses 'cosine'.
+
+        Returns:
+            The created index object.
+        """
+        # Use provided parameters or fall back to instance attributes
+        collection_name = name or self.schema["index"]["name"]
+        embedding_dims = vector_size or self.embedding_model_dims
+        distance_metric = distance or "cosine"
+
+        # Create a new schema with the specified parameters
+        index_schema = {
+            "name": collection_name,
+            "prefix": f"mem0:{collection_name}",
+        }
+
+        # Copy the default fields and update the vector field with the specified dimensions
+        fields = DEFAULT_FIELDS.copy()
+        fields[-1]["attrs"]["dims"] = embedding_dims
+        fields[-1]["attrs"]["distance_metric"] = distance_metric
+
+        # Create the schema
+        schema = {"index": index_schema, "fields": fields}
+
+        # Create the index
+        index = SearchIndex.from_dict(schema)
+        index.set_client(self.client)
+        index.create(overwrite=True)
+
+        # Update instance attributes if creating a new collection
+        if name:
+            self.schema = schema
+            self.index = index
+
+        return index
+
+    def insert(self, vectors: list, payloads: list = None, ids: list = None):
+        data = []
+        for vector, payload, id in zip(vectors, payloads, ids):
+            # Start with required fields
+            entry = {
+                "memory_id": id,
+                "hash": payload["hash"],
+                "memory": payload["data"],
+                "created_at": int(datetime.fromisoformat(payload["created_at"]).timestamp()),
+                "embedding": np.array(vector, dtype=np.float32).tobytes(),
+            }
+
+            # Conditionally add optional fields
+            for field in ["agent_id", "run_id", "user_id"]:
+                if field in payload:
+                    entry[field] = payload[field]
+
+            # Add metadata excluding specific keys
+            entry["metadata"] = json.dumps({k: v for k, v in payload.items() if k not in excluded_keys})
+
+            data.append(entry)
+        self.index.load(data, id_field="memory_id")
+
+    def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None):
+        conditions = [Tag(key) == value for key, value in filters.items() if value is not None]
+        filter = reduce(lambda x, y: x & y, conditions)
+
+        v = VectorQuery(
+            vector=np.array(vectors, dtype=np.float32).tobytes(),
+            vector_field_name="embedding",
+            return_fields=["memory_id", "hash", "agent_id", "run_id", "user_id", "memory", "metadata", "created_at"],
+            filter_expression=filter,
+            num_results=limit,
+        )
+
+        results = self.index.query(v)
+
+        return [
+            MemoryResult(
+                id=result["memory_id"],
+                score=result["vector_distance"],
+                payload={
+                    "hash": result["hash"],
+                    "data": result["memory"],
+                    "created_at": datetime.fromtimestamp(
+                        int(result["created_at"]), tz=pytz.timezone("US/Pacific")
+                    ).isoformat(timespec="microseconds"),
+                    **(
+                        {
+                            "updated_at": datetime.fromtimestamp(
+                                int(result["updated_at"]), tz=pytz.timezone("US/Pacific")
+                            ).isoformat(timespec="microseconds")
+                        }
+                        if "updated_at" in result
+                        else {}
+                    ),
+                    **{field: result[field] for field in ["agent_id", "run_id", "user_id"] if field in result},
+                    **{k: v for k, v in json.loads(extract_json(result["metadata"])).items()},
+                },
+            )
+            for result in results
+        ]
+
+    def delete(self, vector_id):
+        self.index.drop_keys(f"{self.schema['index']['prefix']}:{vector_id}")
+
+    def update(self, vector_id=None, vector=None, payload=None):
+        data = {
+            "memory_id": vector_id,
+            "hash": payload["hash"],
+            "memory": payload["data"],
+            "created_at": int(datetime.fromisoformat(payload["created_at"]).timestamp()),
+            "updated_at": int(datetime.fromisoformat(payload["updated_at"]).timestamp()),
+            "embedding": np.array(vector, dtype=np.float32).tobytes(),
+        }
+
+        for field in ["agent_id", "run_id", "user_id"]:
+            if field in payload:
+                data[field] = payload[field]
+
+        data["metadata"] = json.dumps({k: v for k, v in payload.items() if k not in excluded_keys})
+        self.index.load(data=[data], keys=[f"{self.schema['index']['prefix']}:{vector_id}"], id_field="memory_id")
+
+    def get(self, vector_id):
+        result = self.index.fetch(vector_id)
+        payload = {
+            "hash": result["hash"],
+            "data": result["memory"],
+            "created_at": datetime.fromtimestamp(int(result["created_at"]), tz=pytz.timezone("US/Pacific")).isoformat(
+                timespec="microseconds"
+            ),
+            **(
+                {
+                    "updated_at": datetime.fromtimestamp(
+                        int(result["updated_at"]), tz=pytz.timezone("US/Pacific")
+                    ).isoformat(timespec="microseconds")
+                }
+                if "updated_at" in result
+                else {}
+            ),
+            **{field: result[field] for field in ["agent_id", "run_id", "user_id"] if field in result},
+            **{k: v for k, v in json.loads(extract_json(result["metadata"])).items()},
+        }
+
+        return MemoryResult(id=result["memory_id"], payload=payload)
+
+    def list_cols(self):
+        return self.index.listall()
+
+    def delete_col(self):
+        self.index.delete()
+
+    def col_info(self, name):
+        return self.index.info()
+
+    def reset(self):
+        """
+        Reset the index by deleting and recreating it.
+        """
+        collection_name = self.schema["index"]["name"]
+        logger.warning(f"Resetting index {collection_name}...")
+        self.delete_col()
+
+        self.index = SearchIndex.from_dict(self.schema)
+        self.index.set_client(self.client)
+        self.index.create(overwrite=True)
+
+        # or use
+        # self.create_col(collection_name, self.embedding_model_dims)
+
+        # Recreate the index with the same parameters
+        self.create_col(collection_name, self.embedding_model_dims)
+
+    def list(self, filters: dict = None, limit: int = None) -> list:
+        """
+        List all recent created memories from the vector store.
+        """
+        conditions = [Tag(key) == value for key, value in filters.items() if value is not None]
+        filter = reduce(lambda x, y: x & y, conditions)
+        query = Query(str(filter)).sort_by("created_at", asc=False)
+        if limit is not None:
+            query = Query(str(filter)).sort_by("created_at", asc=False).paging(0, limit)
+
+        results = self.index.search(query)
+        return [
+            [
+                MemoryResult(
+                    id=result["memory_id"],
+                    payload={
+                        "hash": result["hash"],
+                        "data": result["memory"],
+                        "created_at": datetime.fromtimestamp(
+                            int(result["created_at"]), tz=pytz.timezone("US/Pacific")
+                        ).isoformat(timespec="microseconds"),
+                        **(
+                            {
+                                "updated_at": datetime.fromtimestamp(
+                                    int(result["updated_at"]), tz=pytz.timezone("US/Pacific")
+                                ).isoformat(timespec="microseconds")
+                            }
+                            if result.__dict__.get("updated_at")
+                            else {}
+                        ),
+                        **{
+                            field: result[field]
+                            for field in ["agent_id", "run_id", "user_id"]
+                            if field in result.__dict__
+                        },
+                        **{k: v for k, v in json.loads(extract_json(result["metadata"])).items()},
+                    },
+                )
+                for result in results.docs
+            ]
+        ]

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 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], 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)