Recall.Ai / utils /memory_manager.py
Omkar1872's picture
Upload 5 files
169c9db verified
raw
history blame contribute delete
716 Bytes
from sentence_transformers import SentenceTransformer
import chromadb
# Initialize embedding model
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
# Initialize ChromaDB
client = chromadb.Client()
collection = client.create_collection("memory")
def save_to_memory(text, metadata=None):
embedding = embedding_model.encode(text).tolist()
collection.add(documents=[text], embeddings=[embedding], metadatas=[metadata or {}])
def retrieve_memory(query, n_results=3):
embedding = embedding_model.encode(query).tolist()
results = collection.query(query_embeddings=[embedding], n_results=n_results)
return results
def delete_memory(ids):
collection.delete(ids=ids)