madDegen commited on
Commit
66c8911
·
verified ·
1 Parent(s): f5d3e1c

consolidate: Evo ChromaDB vector store wrapper

Browse files
Files changed (1) hide show
  1. evo/chromadb_store.py +32 -0
evo/chromadb_store.py ADDED
@@ -0,0 +1,32 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ """
2
+ Agent Q3 [Evo] — ChromaDB Store
3
+ Persistent vector store. 384-dim nomic-embed-text embeddings.
4
+ """
5
+ import chromadb, os
6
+
7
+ CHROMA_PATH = os.getenv("CHROMA_PATH", "./chroma_data")
8
+
9
+ class ChromaStore:
10
+ def __init__(self, collection: str = "agent_q3_evo"):
11
+ self.client = chromadb.PersistentClient(path=CHROMA_PATH)
12
+ self.collection = self.client.get_or_create_collection(
13
+ name=collection,
14
+ metadata={"hnsw:space": "cosine"}
15
+ )
16
+
17
+ def add(self, ids: list, embeddings: list, documents: list, metadatas: list):
18
+ self.collection.add(ids=ids, embeddings=embeddings, documents=documents, metadatas=metadatas)
19
+
20
+ def query(self, embedding: list, n_results: int = 5) -> dict:
21
+ return self.collection.query(query_embeddings=[embedding], n_results=n_results)
22
+
23
+ def count(self) -> int:
24
+ return self.collection.count()
25
+
26
+ def export_jsonl(self, path: str):
27
+ import json
28
+ results = self.collection.get(include=["documents","metadatas"])
29
+ with open(path, "w") as f:
30
+ for doc, meta in zip(results["documents"], results["metadatas"]):
31
+ f.write(json.dumps({"text": doc, "meta": meta}) + "\n")
32
+ print(f"Exported {self.count()} docs to {path}")