evo/chromadb_store.py

"""
Agent Q3 [Evo] — ChromaDB Store
Persistent vector store. 384-dim nomic-embed-text embeddings.
"""
import chromadb, os

CHROMA_PATH = os.getenv("CHROMA_PATH", "./chroma_data")

class ChromaStore:
    def __init__(self, collection: str = "agent_q3_evo"):
        self.client     = chromadb.PersistentClient(path=CHROMA_PATH)
        self.collection = self.client.get_or_create_collection(
            name=collection,
            metadata={"hnsw:space": "cosine"}
        )

    def add(self, ids: list, embeddings: list, documents: list, metadatas: list):
        self.collection.add(ids=ids, embeddings=embeddings, documents=documents, metadatas=metadatas)

    def query(self, embedding: list, n_results: int = 5) -> dict:
        return self.collection.query(query_embeddings=[embedding], n_results=n_results)

    def count(self) -> int:
        return self.collection.count()

    def export_jsonl(self, path: str):
        import json
        results = self.collection.get(include=["documents","metadatas"])
        with open(path, "w") as f:
            for doc, meta in zip(results["documents"], results["metadatas"]):
                f.write(json.dumps({"text": doc, "meta": meta}) + "\n")
        print(f"Exported {self.count()} docs to {path}")