|
|
from sentence_transformers import SentenceTransformer
|
|
|
import chromadb
|
|
|
|
|
|
|
|
|
embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
|
|
|
|
|
|
|
|
|
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)
|
|
|
|