Spaces:
Sleeping
Sleeping
File size: 2,116 Bytes
3523b6d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
import json, os, time
import numpy as np
from sentence_transformers import SentenceTransformer
def _cosine_sim(a: np.ndarray, b: np.ndarray) -> float:
an = np.linalg.norm(a)
bn = np.linalg.norm(b)
if an == 0 or bn == 0:
return 0.0
return float(np.dot(a, b) / (an * bn))
class SemanticMemory:
"""
Per-agent vector memory: stores (text, category, embedding, ts)
Fast, JSON-backed, CPU-only.
"""
def __init__(self, agent_id: str, file_path: str = None, model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
self.agent_id = agent_id
self.file_path = file_path or f"semantic_memory_{agent_id}.json"
self.model = SentenceTransformer(model_name)
self._init_store()
def _init_store(self):
if not os.path.exists(self.file_path):
with open(self.file_path, "w") as f:
json.dump({"entries": []}, f)
def _load(self):
with open(self.file_path, "r") as f:
return json.load(f)
def _save(self, obj):
with open(self.file_path, "w") as f:
json.dump(obj, f, indent=2)
def add(self, text: str, category: str):
store = self._load()
emb = self.model.encode([text])[0].tolist() # list[float]
store["entries"].append({
"text": text,
"category": category,
"embedding": emb,
"ts": int(time.time())
})
self._save(store)
def query(self, query_text: str, category: str = None, top_k: int = 5):
store = self._load()
entries = store.get("entries", [])
if category:
entries = [e for e in entries if e.get("category") == category]
if not entries:
return []
q = self.model.encode([query_text])[0]
scored = []
for e in entries:
v = np.array(e["embedding"], dtype=np.float32)
s = _cosine_sim(q, v)
scored.append((s, e))
scored.sort(key=lambda x: x[0], reverse=True)
return [e for (s, e) in scored[:top_k] if s > 0.35] # simple confidence cutoff |