Pythonlong_term_memory.py
Copy
1
from openai import OpenAI
2
from pinecone import Pinecone, ServerlessSpec
3
import hashlib, time, json
4
 
5
openai_client = OpenAI()
6
pc = Pinecone()
7
 
8
INDEX_NAME = "agent-memory"
9
NAMESPACE = "facts"
10
 
11
# Create index if it doesn't exist
12
if INDEX_NAME not in pc.list_indexes().names():
13
    pc.create_index(
14
        name=INDEX_NAME, dimension=1536, metric="cosine",
15
        spec=ServerlessSpec(cloud="aws", region="us-east-1")
16
    )
17
 
18
index = pc.Index(INDEX_NAME)
19
 
20
def embed(text: str) -> list[float]:
21
    resp = openai_client.embeddings.create(model="text-embedding-3-small", input=text)
22
    return resp.data[0].embedding
23
 
24
def store(content: str, category: str = "general") -> str:
25
    fact_id = "mem_" + hashlib.md5(content.encode()).hexdigest()[:8]
26
    index.upsert(
27
        vectors=[{"id": fact_id, "values": embed(content), "metadata": {"content": content, "category": category, "ts": time.time()}}],
28
        namespace=NAMESPACE
29
    )
30
    return fact_id
31
 
32
def recall(query: str, top_k: int = 3, threshold: float = 0.75) -> list[str]:
33
    results = index.query(vector=embed(query), top_k=top_k, include_metadata=True, namespace=NAMESPACE)
34
    return [m.metadata["content"] for m in results.matches if m.score >= threshold]
35
 
36
# Test it
37
store("User Alice is a Python developer working on a RAG application", "user_profile")
38
store("Alice prefers dark mode and vim keybindings", "user_preference")
39
store("Alice's current project uses Pinecone and FastAPI", "project_context")
40
 
41
time.sleep(1)  # index propagation
42
print(recall("What programming tools does this user prefer?"))
43