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)