emdedd/ChromaEmbedding.py

from langchain_community.vectorstores.chroma import Chroma

from emdedd.Embedding import Embedding


class ChromaEmbedding(Embedding):
    db: Chroma

    def __init__(self, embedding, path, collection, collection_metadata = None):
        self.db = Chroma(
            embedding_function=embedding,
            persist_directory=path,
            collection_name=collection,
            collection_metadata=collection_metadata
        )

    def embedd(self, chunks, metadata: list[dict] = None):
        self.__store_embeddings(chunks, metadata)

    def __store_embeddings(self, chunks, metadata: list[dict] = None):
        self.db.add_texts(
            texts=chunks,
            metadatas=metadata
        )
        self.db.persist()

    def get_vector_store(self):
        return self.db