File size: 1,090 Bytes
0614fbf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 |
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} |