consolidate: Evo ChromaDB vector store wrapper
Browse files- 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}")
|