itsprasun/pythonic-rag-fastapi-react

1import os2from typing import List3from aimakerspace.text_utils import CharacterTextSplitter, TextFileLoader, PDFLoader4from aimakerspace.openai_utils.prompts import (5    UserRolePrompt,6    SystemRolePrompt,7    AssistantRolePrompt,8)9from aimakerspace.openai_utils.embedding import EmbeddingModel10from aimakerspace.vectordatabase import VectorDatabase11from aimakerspace.openai_utils.chatmodel import ChatOpenAI12 13system_template = """\14Use the following context to answer a users question. If you cannot find the answer in the context, say you don't know the answer."""15system_role_prompt = SystemRolePrompt(system_template)16 17user_prompt_template = """\18Context:19{context}20 21Question:22{question}23"""24user_role_prompt = UserRolePrompt(user_prompt_template)25 26class RetrievalAugmentedQAPipeline:27    def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase) -> None:28        self.llm = llm29        self.vector_db_retriever = vector_db_retriever30 31    async def arun_pipeline(self, user_query: str):32        context_list = self.vector_db_retriever.search_by_text(user_query, k=4)33 34        context_prompt = ""35        for context in context_list:36            context_prompt += context[0] + "\n"37 38        formatted_system_prompt = system_role_prompt.create_message()39        formatted_user_prompt = user_role_prompt.create_message(question=user_query, context=context_prompt)40 41        async def generate_response():42            async for chunk in self.llm.astream([formatted_system_prompt, formatted_user_prompt]):43                yield chunk44 45        return {"response": generate_response(), "context": context_list}46 47text_splitter = CharacterTextSplitter()48 49def process_file(file_path: str, file_name: str):50    print(f"Processing file: {file_name}")51 52    # Create appropriate loader based on file extension53    if file_name.lower().endswith('.pdf'):54        loader = PDFLoader(file_path)55    else:56        loader = TextFileLoader(file_path)57 58    # Load and process the documents59    documents = loader.load_documents()60    texts = text_splitter.split_texts(documents)61    return texts