File size: 716 Bytes
169c9db |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
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)
|