from aimakerspace.vectordatabase import VectorDatabase

class RetrievalAugmentedQAPipeline:
    def __init__(self, llm, vector_db_retriever: VectorDatabase, 
                 system_role_prompt, user_role_prompt
        ) -> None:
        self.llm = llm
        self.vector_db_retriever = vector_db_retriever
        self.system_role_prompt = system_role_prompt
        self.user_role_prompt = user_role_prompt

    async def arun_pipeline(self, user_query: str):
        context_list = self.vector_db_retriever.search_by_text(user_query, k=4)

        context_prompt = ""
        for context in context_list:
            context_prompt += context[0] + "\n"

        formatted_system_prompt = self.system_role_prompt.create_message()

        formatted_user_prompt = self.user_role_prompt.create_message(question=user_query, context=context_prompt)

        async def generate_response():
            async for chunk in self.llm.astream([formatted_system_prompt, formatted_user_prompt]):
                yield chunk

        return {"response": generate_response(), "context": context_list}