_utils/langchain_utils/retriever.py

from _utils.langchain_utils.vector_stores import vector_store


class Retriever:
    def chroma_retriever(
        self, lista_de_documents, search_type="similarity", search_kwargs={"k": 1}
    ):
        retriever = vector_store.chroma(lista_de_documents).as_retriever(
            search_type=search_type,
            search_kwargs=search_kwargs,
        )
        return retriever


retriever = Retriever()