chatbot/vectorstore.py

"""
Modulo que permite gestionar la vector store.
"""
from langchain_chroma import Chroma
import requests, zipfile, io, os

class ChromaDB:
    """
    Clase para gestionar una base ChromaDB
    """

    def __init__(self, embedding_model) -> None:
        if not os.path.exists("chroma_db"):
            print("Descargando base de conocimiento...")
            zip_file_url = "https://drive.google.com/uc?export=download&id=" + os.environ["GDRIVE_ID"]
            r = requests.get(zip_file_url)
            z = zipfile.ZipFile(io.BytesIO(r.content))
            z.extractall()
            print("OK")

        self.db = Chroma(
            collection_name="res_1650",
            embedding_function=embedding_model,
            persist_directory="./chroma_db",
        )