from setup.easy_imports import RunnablePassthrough, create_retrieval_chain

class Chain:

    def create_prompt_model_chain(self, prompt, model):
        return prompt | model

    def create_prompt_model_retriever_chain(self, prompt, model, retriever):
        chain = prompt | model
        return create_retrieval_chain(retriever, chain)

    def invoke_retrieval_chain(self, chain, busca_no_vetor):
        resposta = chain.invoke({"input": busca_no_vetor})

        class Resposta:
            def __init__(self, resposta):
                self.final_answer = resposta["answer"].content
                self.complete_obj = resposta

        return Resposta(resposta)


chain = Chain()