Spaces:
Sleeping
Sleeping
File size: 3,018 Bytes
b48421e b307602 b48421e b307602 b48421e b307602 b48421e b307602 b48421e b307602 | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 | """
core/memory.py — Per-team ChromaDB memory.
Each user/team gets an isolated ChromaDB collection: team_{user_id}
Zero LLM API calls — embeddings are local (SentenceTransformer).
"""
import uuid
from typing import List
CHROMA_PATH = "./chroma_data"
EMBEDDER_MODEL = "all-MiniLM-L6-v2"
# Lazy singletons — loaded on first request, not at startup
_client = None
_embedder = None
def _get_client():
global _client
if _client is None:
import chromadb
_client = chromadb.PersistentClient(path=CHROMA_PATH)
return _client
def _get_embedder():
global _embedder
if _embedder is None:
from sentence_transformers import SentenceTransformer
_embedder = SentenceTransformer(EMBEDDER_MODEL)
return _embedder
def _collection(team_id: str):
"""Get or create a per-team ChromaDB collection."""
safe_id = team_id.replace("-", "_")
return _get_client().get_or_create_collection(
name=f"team_{safe_id}",
metadata={"heuristic": "cosine"},
)
def store_comment(team_id: str, code: str, comment: dict) -> None:
col = _collection(team_id)
text = f"CODE:\n{code}\nCOMMENT:\n{comment['comment']}"
col.add(
ids=[str(uuid.uuid4())],
embeddings=[_get_embedder().encode(text).tolist()],
documents=[text],
metadatas=[{
"line": comment["line"],
"comment": comment["comment"],
"severity": comment["severity"],
"confidence": float(comment["confidence"]),
}],
)
def retrieve_similar(team_id: str, code: str, top_k: int = 3) -> List[dict]:
col = _collection(team_id)
if col.count() == 0:
return []
results = col.query(
query_embeddings=[_get_embedder().encode(f"CODE:\n{code}").tolist()],
n_results=min(top_k, col.count()),
)
memories = [
{
"comment": results["metadatas"][0][i]["comment"],
"severity": results["metadatas"][0][i]["severity"],
"confidence": results["metadatas"][0][i]["confidence"],
"distance": results["distances"][0][i],
}
for i in range(len(results["ids"][0]))
]
return sorted(memories, key=lambda x: x["distance"])
def get_all_memories(team_id: str, limit: int = 50) -> List[dict]:
col = _collection(team_id)
if col.count() == 0:
return []
results = col.get(limit=limit, include=["metadatas"])
return [
{
"id": results["ids"][i],
"comment": results["metadatas"][i]["comment"],
"severity": results["metadatas"][i]["severity"],
"confidence": results["metadatas"][i]["confidence"],
}
for i in range(len(results["ids"]))
]
def delete_memory(team_id: str, memory_id: str) -> None:
_collection(team_id).delete(ids=[memory_id])
def memory_count(team_id: str) -> int:
try:
return _collection(team_id).count()
except Exception:
return 0
|