mem0ai--mem0/mem0/vector_stores/faiss.py

import json
import logging
import os
import pickle
import uuid
import warnings
from pathlib import Path
from typing import Any, Dict, List, Optional

import numpy as np
from pydantic import BaseModel

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 mem0.vector_stores.base import VectorStoreBase

logger = logging.getLogger(__name__)


class SafeUnpickler(pickle.Unpickler):
    """
    Restricted unpickler that only allows safe built-in types.

    This prevents arbitrary code execution via pickle deserialization by only
    allowing a whitelist of safe types (dict, list, str, int, float, bool, tuple, None).
    """

    # Only allow builtins module
    SAFE_MODULES = frozenset({"builtins", "__builtin__"})
    # Only allow safe basic types
    SAFE_NAMES = frozenset({"dict", "list", "str", "int", "float", "bool", "tuple", "set", "frozenset", "NoneType"})

    def find_class(self, module: str, name: str) -> Any:
        """Override find_class to only allow safe types."""
        if module in self.SAFE_MODULES and name in self.SAFE_NAMES:
            import builtins

            if hasattr(builtins, name):
                return getattr(builtins, name)
            # NoneType special case
            if name == "NoneType":
                return type(None)
        raise pickle.UnpicklingError(
            f"Unsafe pickle: attempted to load '{module}.{name}'. "
            f"Only basic Python types are allowed for security reasons."
        )


def _safe_pickle_load(file_path: str) -> Any:
    """
    Safely load a pickle file using restricted unpickler.

    Args:
        file_path: Path to the pickle file.

    Returns:
        The deserialized object (only basic Python types allowed).

    Raises:
        pickle.UnpicklingError: If the pickle contains unsafe types.
    """
    with open(file_path, "rb") as f:
        return SafeUnpickler(f).load()


def _validate_docstore_structure(data: Any) -> tuple:
    """
    Validate that loaded data has the expected structure.

    Args:
        data: The loaded data to validate.

    Returns:
        Tuple of (docstore, index_to_id) if valid.

    Raises:
        ValueError: If the data structure is invalid.
    """
    if not isinstance(data, tuple) or len(data) != 2:
        raise ValueError("Invalid docstore format: expected tuple of (docstore, index_to_id)")

    docstore, index_to_id = data

    if not isinstance(docstore, dict):
        raise ValueError("Invalid docstore format: docstore must be a dict")

    if not isinstance(index_to_id, dict):
        raise ValueError("Invalid docstore format: index_to_id must be a dict")

    # Validate docstore entries
    for key, value in docstore.items():
        if not isinstance(key, str):
            raise ValueError(f"Invalid docstore key type: {type(key)}, expected str")
        if not isinstance(value, dict):
            raise ValueError(f"Invalid docstore value type: {type(value)}, expected dict")

    # Validate index_to_id entries
    for key, value in index_to_id.items():
        if not isinstance(key, int):
            raise ValueError(f"Invalid index_to_id key type: {type(key)}, expected int")
        if not isinstance(value, str):
            raise ValueError(f"Invalid index_to_id value type: {type(value)}, expected str")

    return docstore, index_to_id


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"
            json_docstore_path = f"{self.path}/{collection_name}.json"
            pkl_docstore_path = f"{self.path}/{collection_name}.pkl"

            # Check for index file and either JSON (preferred) or legacy pickle docstore
            if os.path.exists(index_path) and (os.path.exists(json_docstore_path) or os.path.exists(pkl_docstore_path)):
                # _load will prefer JSON over pickle and auto-migrate
                self._load(index_path, pkl_docstore_path)
            else:
                self.create_col(collection_name)

    def _load(self, index_path: str, docstore_path: str):
        """
        Load FAISS index and docstore from disk.

        Supports both JSON (preferred) and legacy pickle formats. Pickle files are loaded
        using a restricted unpickler that only allows basic Python types to prevent
        arbitrary code execution (CVE mitigation).

        Args:
            index_path (str): Path to FAISS index file.
            docstore_path (str): Path to docstore file (.json or legacy .pkl).
        """
        try:
            self.index = faiss.read_index(index_path)

            # Determine docstore format - prefer JSON over pickle
            json_docstore_path = docstore_path.replace(".pkl", ".json")

            if os.path.exists(json_docstore_path):
                # Load from JSON (safe, preferred format)
                with open(json_docstore_path, "r", encoding="utf-8") as f:
                    data = json.load(f)
                self.docstore = data.get("docstore", {})
                # JSON keys are always strings, convert back to int
                self.index_to_id = {int(k): v for k, v in data.get("index_to_id", {}).items()}
                logger.info(f"Loaded FAISS index from {index_path} with {self.index.ntotal} vectors (JSON format)")

            elif os.path.exists(docstore_path):
                # Load from legacy pickle using safe unpickler
                # This prevents arbitrary code execution from malicious pickle files
                logger.warning(
                    f"Loading legacy pickle docstore from {docstore_path}. "
                    f"Consider migrating to JSON format for better security."
                )
                data = _safe_pickle_load(docstore_path)
                self.docstore, self.index_to_id = _validate_docstore_structure(data)
                logger.info(f"Loaded FAISS index from {index_path} with {self.index.ntotal} vectors (pickle format)")

                # Auto-migrate to JSON format
                self._save()
                logger.info(f"Migrated docstore to JSON format: {json_docstore_path}")

            else:
                raise FileNotFoundError(f"No docstore found at {docstore_path} or {json_docstore_path}")

        except pickle.UnpicklingError as e:
            logger.error(f"Security error loading FAISS docstore: {e}")
            raise ValueError(f"Failed to load FAISS docstore: potentially malicious pickle file. {e}") from e
        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 using JSON format (secure)."""
        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"
            json_docstore_path = f"{self.path}/{self.collection_name}.json"

            faiss.write_index(self.index, index_path)

            # Save docstore as JSON (safe format, no code execution risk)
            # JSON keys must be strings, so convert int keys to str
            data = {
                "docstore": self.docstore,
                "index_to_id": {str(k): v for k, v in self.index_to_id.items()},
            }
            with open(json_docstore_path, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2)

        except Exception as e:
            logger.warning(f"Failed to save FAISS index: {e}")

    def _should_normalize(self) -> bool:
        """Whether vectors must be L2-normalized before indexing/searching.

        Cosine similarity is implemented on top of an inner-product index
        (``IndexFlatIP``), which only equals cosine when the inputs are unit
        vectors — so cosine *always* requires normalization. For euclidean the
        ``normalize_L2`` flag remains an opt-in.
        """
        strategy = self.distance_strategy.lower()
        if strategy == "cosine":
            return True
        return self.normalize_L2 and strategy == "euclidean"

    def _parse_output(self, scores, ids, top_k=None) -> List[OutputData]:
        """
        Parse the output data.

        Args:
            scores: Similarity scores from FAISS.
            ids: Indices from FAISS.
            top_k: Maximum number of results to return.

        Returns:
            List[OutputData]: Parsed output data.
        """
        if top_k is None:
            top_k = len(ids)

        results = []
        for i in range(min(len(ids), top_k)):
            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()

            raw_score = float(scores[i])
            if self.distance_strategy.lower() == "euclidean":
                score = 1.0 / (1.0 + raw_score)
            else:
                score = raw_score
            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._should_normalize():
            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], top_k: 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.
            top_k (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._should_normalize():
            faiss.normalize_L2(query_vectors)

        fetch_k = top_k * 2 if filters else top_k
        scores, indices = self.index.search(query_vectors, fetch_k)

        results = self._parse_output(scores[0], indices[0], fetch_k)

        if filters:
            filtered_results = []
            for result in results:
                if self._apply_filters(result.payload, filters):
                    filtered_results.append(result)
                    if len(filtered_results) >= top_k:
                        break
            results = filtered_results[:top_k]

        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:
            # Reconstruct remaining vectors and rebuild the FAISS index
            remaining_vectors = []
            new_index_to_id = {}
            new_idx = 0
            for old_idx in sorted(self.index_to_id.keys()):
                if old_idx == index_to_delete:
                    continue
                remaining_vectors.append(self.index.reconstruct(int(old_idx)))
                new_index_to_id[new_idx] = self.index_to_id[old_idx]
                new_idx += 1

            self.index.reset()
            if remaining_vectors:
                self.index.add(np.array(remaining_vectors, dtype=np.float32))

            self.docstore.pop(vector_id, None)
            self.index_to_id = new_index_to_id

            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"
                json_docstore_path = f"{self.path}/{self.collection_name}.json"
                pkl_docstore_path = f"{self.path}/{self.collection_name}.pkl"

                if os.path.exists(index_path):
                    os.remove(index_path)
                if os.path.exists(json_docstore_path):
                    os.remove(json_docstore_path)
                # Also clean up legacy pickle files if they exist
                if os.path.exists(pkl_docstore_path):
                    os.remove(pkl_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, top_k: int = 100) -> List[OutputData]:
        """
        List all vectors in a collection.

        Args:
            filters (Optional[Dict], optional): Filters to apply to the list. Defaults to None.
            top_k (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 >= top_k:
                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)