agentrun-mem0ai 0.0.11__py3-none-any.whl

agentrun_mem0/vector_stores/faiss.py

@@ -0,0 +1,479 @@
+import logging
+import os
+import pickle
+import uuid
+from pathlib import Path
+from typing import Dict, List, Optional
+
+import numpy as np
+from pydantic import BaseModel
+
+import warnings
+
+try:
+    # Suppress SWIG deprecation warnings from FAISS
+    warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*SwigPy.*")
+    warnings.filterwarnings("ignore", category=DeprecationWarning, message=".*swigvarlink.*")
+    
+    logging.getLogger("faiss").setLevel(logging.WARNING)
+    logging.getLogger("faiss.loader").setLevel(logging.WARNING)
+
+    import faiss
+except ImportError:
+    raise ImportError(
+        "Could not import faiss python package. "
+        "Please install it with `pip install faiss-gpu` (for CUDA supported GPU) "
+        "or `pip install faiss-cpu` (depending on Python version)."
+    )
+
+from agentrun_mem0.vector_stores.base import VectorStoreBase
+
+logger = logging.getLogger(__name__)
+
+
+class OutputData(BaseModel):
+    id: Optional[str]  # memory id
+    score: Optional[float]  # distance
+    payload: Optional[Dict]  # metadata
+
+
+class FAISS(VectorStoreBase):
+    def __init__(
+        self,
+        collection_name: str,
+        path: Optional[str] = None,
+        distance_strategy: str = "euclidean",
+        normalize_L2: bool = False,
+        embedding_model_dims: int = 1536,
+    ):
+        """
+        Initialize the FAISS vector store.
+
+        Args:
+            collection_name (str): Name of the collection.
+            path (str, optional): Path for local FAISS database. Defaults to None.
+            distance_strategy (str, optional): Distance strategy to use. Options: 'euclidean', 'inner_product', 'cosine'.
+                Defaults to "euclidean".
+            normalize_L2 (bool, optional): Whether to normalize L2 vectors. Only applicable for euclidean distance.
+                Defaults to False.
+        """
+        self.collection_name = collection_name
+        self.path = path or f"/tmp/faiss/{collection_name}"
+        self.distance_strategy = distance_strategy
+        self.normalize_L2 = normalize_L2
+        self.embedding_model_dims = embedding_model_dims
+
+        # Initialize storage structures
+        self.index = None
+        self.docstore = {}
+        self.index_to_id = {}
+
+        # Create directory if it doesn't exist
+        if self.path:
+            os.makedirs(os.path.dirname(self.path), exist_ok=True)
+
+            # Try to load existing index if available
+            index_path = f"{self.path}/{collection_name}.faiss"
+            docstore_path = f"{self.path}/{collection_name}.pkl"
+            if os.path.exists(index_path) and os.path.exists(docstore_path):
+                self._load(index_path, docstore_path)
+            else:
+                self.create_col(collection_name)
+
+    def _load(self, index_path: str, docstore_path: str):
+        """
+        Load FAISS index and docstore from disk.
+
+        Args:
+            index_path (str): Path to FAISS index file.
+            docstore_path (str): Path to docstore pickle file.
+        """
+        try:
+            self.index = faiss.read_index(index_path)
+            with open(docstore_path, "rb") as f:
+                self.docstore, self.index_to_id = pickle.load(f)
+            logger.info(f"Loaded FAISS index from {index_path} with {self.index.ntotal} vectors")
+        except Exception as e:
+            logger.warning(f"Failed to load FAISS index: {e}")
+
+            self.docstore = {}
+            self.index_to_id = {}
+
+    def _save(self):
+        """Save FAISS index and docstore to disk."""
+        if not self.path or not self.index:
+            return
+
+        try:
+            os.makedirs(self.path, exist_ok=True)
+            index_path = f"{self.path}/{self.collection_name}.faiss"
+            docstore_path = f"{self.path}/{self.collection_name}.pkl"
+
+            faiss.write_index(self.index, index_path)
+            with open(docstore_path, "wb") as f:
+                pickle.dump((self.docstore, self.index_to_id), f)
+        except Exception as e:
+            logger.warning(f"Failed to save FAISS index: {e}")
+
+    def _parse_output(self, scores, ids, limit=None) -> List[OutputData]:
+        """
+        Parse the output data.
+
+        Args:
+            scores: Similarity scores from FAISS.
+            ids: Indices from FAISS.
+            limit: Maximum number of results to return.
+
+        Returns:
+            List[OutputData]: Parsed output data.
+        """
+        if limit is None:
+            limit = len(ids)
+
+        results = []
+        for i in range(min(len(ids), limit)):
+            if ids[i] == -1:  # FAISS returns -1 for empty results
+                continue
+
+            index_id = int(ids[i])
+            vector_id = self.index_to_id.get(index_id)
+            if vector_id is None:
+                continue
+
+            payload = self.docstore.get(vector_id)
+            if payload is None:
+                continue
+
+            payload_copy = payload.copy()
+
+            score = float(scores[i])
+            entry = OutputData(
+                id=vector_id,
+                score=score,
+                payload=payload_copy,
+            )
+            results.append(entry)
+
+        return results
+
+    def create_col(self, name: str, distance: str = None):
+        """
+        Create a new collection.
+
+        Args:
+            name (str): Name of the collection.
+            distance (str, optional): Distance metric to use. Overrides the distance_strategy
+                passed during initialization. Defaults to None.
+
+        Returns:
+            self: The FAISS instance.
+        """
+        distance_strategy = distance or self.distance_strategy
+
+        # Create index based on distance strategy
+        if distance_strategy.lower() == "inner_product" or distance_strategy.lower() == "cosine":
+            self.index = faiss.IndexFlatIP(self.embedding_model_dims)
+        else:
+            self.index = faiss.IndexFlatL2(self.embedding_model_dims)
+
+        self.collection_name = name
+
+        self._save()
+
+        return self
+
+    def insert(
+        self,
+        vectors: List[list],
+        payloads: Optional[List[Dict]] = None,
+        ids: Optional[List[str]] = None,
+    ):
+        """
+        Insert vectors into a collection.
+
+        Args:
+            vectors (List[list]): List of vectors to insert.
+            payloads (Optional[List[Dict]], optional): List of payloads corresponding to vectors. Defaults to None.
+            ids (Optional[List[str]], optional): List of IDs corresponding to vectors. Defaults to None.
+        """
+        if self.index is None:
+            raise ValueError("Collection not initialized. Call create_col first.")
+
+        if ids is None:
+            ids = [str(uuid.uuid4()) for _ in range(len(vectors))]
+
+        if payloads is None:
+            payloads = [{} for _ in range(len(vectors))]
+
+        if len(vectors) != len(ids) or len(vectors) != len(payloads):
+            raise ValueError("Vectors, payloads, and IDs must have the same length")
+
+        vectors_np = np.array(vectors, dtype=np.float32)
+
+        if self.normalize_L2 and self.distance_strategy.lower() == "euclidean":
+            faiss.normalize_L2(vectors_np)
+
+        self.index.add(vectors_np)
+
+        starting_idx = len(self.index_to_id)
+        for i, (vector_id, payload) in enumerate(zip(ids, payloads)):
+            self.docstore[vector_id] = payload.copy()
+            self.index_to_id[starting_idx + i] = vector_id
+
+        self._save()
+
+        logger.info(f"Inserted {len(vectors)} vectors into collection {self.collection_name}")
+
+    def search(
+        self, query: str, vectors: List[list], limit: int = 5, filters: Optional[Dict] = None
+    ) -> List[OutputData]:
+        """
+        Search for similar vectors.
+
+        Args:
+            query (str): Query (not used, kept for API compatibility).
+            vectors (List[list]): List of vectors to search.
+            limit (int, optional): Number of results to return. Defaults to 5.
+            filters (Optional[Dict], optional): Filters to apply to the search. Defaults to None.
+
+        Returns:
+            List[OutputData]: Search results.
+        """
+        if self.index is None:
+            raise ValueError("Collection not initialized. Call create_col first.")
+
+        query_vectors = np.array(vectors, dtype=np.float32)
+
+        if len(query_vectors.shape) == 1:
+            query_vectors = query_vectors.reshape(1, -1)
+
+        if self.normalize_L2 and self.distance_strategy.lower() == "euclidean":
+            faiss.normalize_L2(query_vectors)
+
+        fetch_k = limit * 2 if filters else limit
+        scores, indices = self.index.search(query_vectors, fetch_k)
+
+        results = self._parse_output(scores[0], indices[0], limit)
+
+        if filters:
+            filtered_results = []
+            for result in results:
+                if self._apply_filters(result.payload, filters):
+                    filtered_results.append(result)
+                    if len(filtered_results) >= limit:
+                        break
+            results = filtered_results[:limit]
+
+        return results
+
+    def _apply_filters(self, payload: Dict, filters: Dict) -> bool:
+        """
+        Apply filters to a payload.
+
+        Args:
+            payload (Dict): Payload to filter.
+            filters (Dict): Filters to apply.
+
+        Returns:
+            bool: True if payload passes filters, False otherwise.
+        """
+        if not filters or not payload:
+            return True
+
+        for key, value in filters.items():
+            if key not in payload:
+                return False
+
+            if isinstance(value, list):
+                if payload[key] not in value:
+                    return False
+            elif payload[key] != value:
+                return False
+
+        return True
+
+    def delete(self, vector_id: str):
+        """
+        Delete a vector by ID.
+
+        Args:
+            vector_id (str): ID of the vector to delete.
+        """
+        if self.index is None:
+            raise ValueError("Collection not initialized. Call create_col first.")
+
+        index_to_delete = None
+        for idx, vid in self.index_to_id.items():
+            if vid == vector_id:
+                index_to_delete = idx
+                break
+
+        if index_to_delete is not None:
+            self.docstore.pop(vector_id, None)
+            self.index_to_id.pop(index_to_delete, None)
+
+            self._save()
+
+            logger.info(f"Deleted vector {vector_id} from collection {self.collection_name}")
+        else:
+            logger.warning(f"Vector {vector_id} not found in collection {self.collection_name}")
+
+    def update(
+        self,
+        vector_id: str,
+        vector: Optional[List[float]] = None,
+        payload: Optional[Dict] = None,
+    ):
+        """
+        Update a vector and its payload.
+
+        Args:
+            vector_id (str): ID of the vector to update.
+            vector (Optional[List[float]], optional): Updated vector. Defaults to None.
+            payload (Optional[Dict], optional): Updated payload. Defaults to None.
+        """
+        if self.index is None:
+            raise ValueError("Collection not initialized. Call create_col first.")
+
+        if vector_id not in self.docstore:
+            raise ValueError(f"Vector {vector_id} not found")
+
+        current_payload = self.docstore[vector_id].copy()
+
+        if payload is not None:
+            self.docstore[vector_id] = payload.copy()
+            current_payload = self.docstore[vector_id].copy()
+
+        if vector is not None:
+            self.delete(vector_id)
+            self.insert([vector], [current_payload], [vector_id])
+        else:
+            self._save()
+
+        logger.info(f"Updated vector {vector_id} in collection {self.collection_name}")
+
+    def get(self, vector_id: str) -> OutputData:
+        """
+        Retrieve a vector by ID.
+
+        Args:
+            vector_id (str): ID of the vector to retrieve.
+
+        Returns:
+            OutputData: Retrieved vector.
+        """
+        if self.index is None:
+            raise ValueError("Collection not initialized. Call create_col first.")
+
+        if vector_id not in self.docstore:
+            return None
+
+        payload = self.docstore[vector_id].copy()
+
+        return OutputData(
+            id=vector_id,
+            score=None,
+            payload=payload,
+        )
+
+    def list_cols(self) -> List[str]:
+        """
+        List all collections.
+
+        Returns:
+            List[str]: List of collection names.
+        """
+        if not self.path:
+            return [self.collection_name] if self.index else []
+
+        try:
+            collections = []
+            path = Path(self.path).parent
+            for file in path.glob("*.faiss"):
+                collections.append(file.stem)
+            return collections
+        except Exception as e:
+            logger.warning(f"Failed to list collections: {e}")
+            return [self.collection_name] if self.index else []
+
+    def delete_col(self):
+        """
+        Delete a collection.
+        """
+        if self.path:
+            try:
+                index_path = f"{self.path}/{self.collection_name}.faiss"
+                docstore_path = f"{self.path}/{self.collection_name}.pkl"
+
+                if os.path.exists(index_path):
+                    os.remove(index_path)
+                if os.path.exists(docstore_path):
+                    os.remove(docstore_path)
+
+                logger.info(f"Deleted collection {self.collection_name}")
+            except Exception as e:
+                logger.warning(f"Failed to delete collection: {e}")
+
+        self.index = None
+        self.docstore = {}
+        self.index_to_id = {}
+
+    def col_info(self) -> Dict:
+        """
+        Get information about a collection.
+
+        Returns:
+            Dict: Collection information.
+        """
+        if self.index is None:
+            return {"name": self.collection_name, "count": 0}
+
+        return {
+            "name": self.collection_name,
+            "count": self.index.ntotal,
+            "dimension": self.index.d,
+            "distance": self.distance_strategy,
+        }
+
+    def list(self, filters: Optional[Dict] = None, limit: int = 100) -> List[OutputData]:
+        """
+        List all vectors in a collection.
+
+        Args:
+            filters (Optional[Dict], optional): Filters to apply to the list. Defaults to None.
+            limit (int, optional): Number of vectors to return. Defaults to 100.
+
+        Returns:
+            List[OutputData]: List of vectors.
+        """
+        if self.index is None:
+            return []
+
+        results = []
+        count = 0
+
+        for vector_id, payload in self.docstore.items():
+            if filters and not self._apply_filters(payload, filters):
+                continue
+
+            payload_copy = payload.copy()
+
+            results.append(
+                OutputData(
+                    id=vector_id,
+                    score=None,
+                    payload=payload_copy,
+                )
+            )
+
+            count += 1
+            if count >= limit:
+                break
+
+        return [results]
+
+    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.collection_name)

agentrun_mem0/vector_stores/langchain.py

@@ -0,0 +1,180 @@
+import logging
+from typing import Dict, List, Optional
+
+from pydantic import BaseModel
+
+try:
+    from langchain_community.vectorstores import VectorStore
+except ImportError:
+    raise ImportError(
+        "The 'langchain_community' library is required. Please install it using 'pip install langchain_community'."
+    )
+
+from agentrun_mem0.vector_stores.base import VectorStoreBase
+
+logger = logging.getLogger(__name__)
+
+
+class OutputData(BaseModel):
+    id: Optional[str]  # memory id
+    score: Optional[float]  # distance
+    payload: Optional[Dict]  # metadata
+
+
+class Langchain(VectorStoreBase):
+    def __init__(self, client: VectorStore, collection_name: str = "mem0"):
+        self.client = client
+        self.collection_name = collection_name
+
+    def _parse_output(self, data: Dict) -> List[OutputData]:
+        """
+        Parse the output data.
+
+        Args:
+            data (Dict): Output data or list of Document objects.
+
+        Returns:
+            List[OutputData]: Parsed output data.
+        """
+        # Check if input is a list of Document objects
+        if isinstance(data, list) and all(hasattr(doc, "metadata") for doc in data if hasattr(doc, "__dict__")):
+            result = []
+            for doc in data:
+                entry = OutputData(
+                    id=getattr(doc, "id", None),
+                    score=None,  # Document objects typically don't include scores
+                    payload=getattr(doc, "metadata", {}),
+                )
+                result.append(entry)
+            return result
+
+        # Original format handling
+        keys = ["ids", "distances", "metadatas"]
+        values = []
+
+        for key in keys:
+            value = data.get(key, [])
+            if isinstance(value, list) and value and isinstance(value[0], list):
+                value = value[0]
+            values.append(value)
+
+        ids, distances, metadatas = values
+        max_length = max(len(v) for v in values if isinstance(v, list) and v is not None)
+
+        result = []
+        for i in range(max_length):
+            entry = OutputData(
+                id=ids[i] if isinstance(ids, list) and ids and i < len(ids) else None,
+                score=(distances[i] if isinstance(distances, list) and distances and i < len(distances) else None),
+                payload=(metadatas[i] if isinstance(metadatas, list) and metadatas and i < len(metadatas) else None),
+            )
+            result.append(entry)
+
+        return result
+
+    def create_col(self, name, vector_size=None, distance=None):
+        self.collection_name = name
+        return self.client
+
+    def insert(
+        self, vectors: List[List[float]], payloads: Optional[List[Dict]] = None, ids: Optional[List[str]] = None
+    ):
+        """
+        Insert vectors into the LangChain vectorstore.
+        """
+        # Check if client has add_embeddings method
+        if hasattr(self.client, "add_embeddings"):
+            # Some LangChain vectorstores have a direct add_embeddings method
+            self.client.add_embeddings(embeddings=vectors, metadatas=payloads, ids=ids)
+        else:
+            # Fallback to add_texts method
+            texts = [payload.get("data", "") for payload in payloads] if payloads else [""] * len(vectors)
+            self.client.add_texts(texts=texts, metadatas=payloads, ids=ids)
+
+    def search(self, query: str, vectors: List[List[float]], limit: int = 5, filters: Optional[Dict] = None):
+        """
+        Search for similar vectors in LangChain.
+        """
+        # For each vector, perform a similarity search
+        if filters:
+            results = self.client.similarity_search_by_vector(embedding=vectors, k=limit, filter=filters)
+        else:
+            results = self.client.similarity_search_by_vector(embedding=vectors, k=limit)
+
+        final_results = self._parse_output(results)
+        return final_results
+
+    def delete(self, vector_id):
+        """
+        Delete a vector by ID.
+        """
+        self.client.delete(ids=[vector_id])
+
+    def update(self, vector_id, vector=None, payload=None):
+        """
+        Update a vector and its payload.
+        """
+        self.delete(vector_id)
+        self.insert(vector, payload, [vector_id])
+
+    def get(self, vector_id):
+        """
+        Retrieve a vector by ID.
+        """
+        docs = self.client.get_by_ids([vector_id])
+        if docs and len(docs) > 0:
+            doc = docs[0]
+            return self._parse_output([doc])[0]
+        return None
+
+    def list_cols(self):
+        """
+        List all collections.
+        """
+        # LangChain doesn't have collections
+        return [self.collection_name]
+
+    def delete_col(self):
+        """
+        Delete a collection.
+        """
+        logger.warning("Deleting collection")
+        if hasattr(self.client, "delete_collection"):
+            self.client.delete_collection()
+        elif hasattr(self.client, "reset_collection"):
+            self.client.reset_collection()
+        else:
+            self.client.delete(ids=None)
+
+    def col_info(self):
+        """
+        Get information about a collection.
+        """
+        return {"name": self.collection_name}
+
+    def list(self, filters=None, limit=None):
+        """
+        List all vectors in a collection.
+        """
+        try:
+            if hasattr(self.client, "_collection") and hasattr(self.client._collection, "get"):
+                # Convert mem0 filters to Chroma where clause if needed
+                where_clause = None
+                if filters:
+                    # Handle all filters, not just user_id
+                    where_clause = filters
+
+                result = self.client._collection.get(where=where_clause, limit=limit)
+
+                # Convert the result to the expected format
+                if result and isinstance(result, dict):
+                    return [self._parse_output(result)]
+                return []
+        except Exception as e:
+            logger.error(f"Error listing vectors from Chroma: {e}")
+            return []
+
+    def reset(self):
+        """Reset the index by deleting and recreating it."""
+        logger.warning(f"Resetting collection: {self.collection_name}")
+        self.delete_col()