chroma.py

import chromadb


class VectorStore:
    def __init__(self, collection_name) -> None:
        self.client = chromadb.PersistentClient(path="/tmp/embeddings.vector")
        self.collection_name = collection_name
        self.collection = self.client.get_or_create_collection(collection_name)

    def embed_documents(
        self, docs: list, metadatas: list, ids: list, refresh: bool = False
    ) -> chromadb.Collection:
        if refresh:
            print("Refreshing collection")
            self.client.delete_collection(name=self.collection.name)
            self.collection = self.client.get_or_create_collection(self.collection_name)
        elif len(ids) == self.collection.count():
            print("Documents already embedded")
            return
        print(f"Embedding {len(ids)} documents..")
        self.collection.upsert(documents=docs, metadatas=metadatas, ids=ids)

    def query_collection(self, query, num_results=10):
        results = self.collection.query(
            query_texts=[query],
            n_results=num_results,
        )
        return results