474 lines
15 KiB
Python
474 lines
15 KiB
Python
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
|
|
|
|
try:
|
|
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 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)
|