import logging
import os
import pickle
import uuid
import warnings
from pathlib import Path
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 selfmemory.vector_stores.base import VectorStoreBase
logger = logging.getLogger(__name__)
class OutputData(BaseModel):
id: str | None # memory id
score: float | None # distance
payload: dict | None # metadata
class FAISS(VectorStoreBase):
def __init__(
self,
collection_name: str,
path: str | None = 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: list[dict] | None = None,
ids: list[str] | None = 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, strict=False)):
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: dict | None = 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: list[float] | None = None,
payload: dict | None = 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: dict | None = 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)
