atolat30/pythonic-rag-fastapi-react

1import os2from typing import List3from chainlit.types import AskFileResponse4from aimakerspace.text_utils import CharacterTextSplitter, TextFileLoader, PDFLoader5from aimakerspace.openai_utils.prompts import (6    UserRolePrompt,7    SystemRolePrompt,8    AssistantRolePrompt,9)10from aimakerspace.openai_utils.embedding import EmbeddingModel11from aimakerspace.vectordatabase import VectorDatabase12from aimakerspace.openai_utils.chatmodel import ChatOpenAI13import chainlit as cl14 15system_template = """\16You are a helpful AI assistant that answers questions based on the provided context. 17Your task is to:181. Carefully read and understand the context192. Answer the user's question using ONLY the information from the context203. If the answer cannot be found in the context, say "I cannot find the answer in the provided context"214. If you find partial information, share what you found and indicate if more information might be needed22 23Remember: Only use information from the provided context to answer questions."""24system_role_prompt = SystemRolePrompt(system_template)25 26user_prompt_template = """\27Context:28{context}29 30Based on the above context, please answer the following question. If the answer cannot be found in the context, say "I cannot find the answer in the provided context." If you find partial information, share what you found and indicate if more information might be needed.31 32Question:33{question}34 35Please provide a clear and concise answer based ONLY on the information in the context above."""36user_role_prompt = UserRolePrompt(user_prompt_template)37 38class RetrievalAugmentedQAPipeline:39    def __init__(self, llm: ChatOpenAI(), vector_db_retriever: VectorDatabase) -> None:40        self.llm = llm41        self.vector_db_retriever = vector_db_retriever42 43    async def arun_pipeline(self, user_query: str):44        # Get more contexts with a broader search45        print("\nSearching for relevant contexts...")46        context_list = self.vector_db_retriever.search_by_text(user_query, k=5)  # Increased from 3 to 547        48        print("\nRetrieved contexts:")49        for i, (context, score) in enumerate(context_list):50            print(f"\nContext {i+1} (score: {score:.3f}):")51            print(context[:500] + "..." if len(context) > 500 else context)  # Show more context52        53        # Limit total context length to approximately 3000 tokens (12000 characters)54        context_prompt = ""55        total_length = 056        max_length = 12000  # Reduced from 24000 to 1200057        58        # Sort contexts by score before truncating59        sorted_contexts = sorted(context_list, key=lambda x: x[1], reverse=True)60        61        for context, score in sorted_contexts:62            if total_length + len(context) > max_length:63                print(f"\nSkipping context with score {score:.3f} due to length limit")64                continue65            context_prompt += context + "\n"66            total_length += len(context)67            68        print(f"\nUsing {len(context_prompt.split())} words of context")69 70        formatted_system_prompt = system_role_prompt.create_message()71        formatted_user_prompt = user_role_prompt.create_message(question=user_query, context=context_prompt)72 73        print("\nFinal messages being sent to the model:")74        print("\nSystem prompt:")75        print(formatted_system_prompt)76        print("\nUser prompt:")77        print(formatted_user_prompt)78 79        async def generate_response():80            async for chunk in self.llm.astream([formatted_system_prompt, formatted_user_prompt]):81                yield chunk82 83        return {"response": generate_response(), "context": context_list}84 85text_splitter = CharacterTextSplitter()86 87 88def process_file(file: AskFileResponse):89    import tempfile90    import shutil91    92    print(f"Processing file: {file.name}")93    94    # Create a temporary file with the correct extension95    suffix = f".{file.name.split('.')[-1]}"96    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as temp_file:97        # Copy the uploaded file content to the temporary file98        shutil.copyfile(file.path, temp_file.name)99        print(f"Created temporary file at: {temp_file.name}")100        101        # Create appropriate loader102        if file.name.lower().endswith('.pdf'):103            loader = PDFLoader(temp_file.name)104        else:105            loader = TextFileLoader(temp_file.name)106            107        try:108            # Load and process the documents109            documents = loader.load_documents()110            texts = text_splitter.split_texts(documents)111            return texts112        finally:113            # Clean up the temporary file114            try:115                os.unlink(temp_file.name)116            except Exception as e:117                print(f"Error cleaning up temporary file: {e}")118 119 120@cl.on_chat_start121async def on_chat_start():122    files = None123 124    # Wait for the user to upload a file125    while files == None:126        files = await cl.AskFileMessage(127            content="Please upload a Text or PDF file to begin!",128            accept=["text/plain", "application/pdf"],129            max_size_mb=2,130            timeout=180,131        ).send()132 133    file = files[0]134    print(f"Received file: {file.name} ({file.type})")135 136    msg = cl.Message(137        content=f"Processing `{file.name}`..."138    )139    await msg.send()140 141    # load the file142    try:143        texts = process_file(file)144        print(f"Successfully processed file. Generated {len(texts)} text chunks")145        print("Sample of first chunk:", texts[0][:200] if texts else "No texts generated")146    except Exception as e:147        print(f"Error processing file: {str(e)}")148        await cl.Message(content=f"Error processing file: {str(e)}").send()149        return150 151    # Create a dict vector store152    try:153        vector_db = VectorDatabase()154        vector_db = await vector_db.abuild_from_list(texts)155        print("Successfully created vector database")156    except Exception as e:157        print(f"Error creating vector database: {str(e)}")158        await cl.Message(content=f"Error creating vector database: {str(e)}").send()159        return160    161    try:162        chat_openai = ChatOpenAI()163        print("Successfully initialized ChatOpenAI")164    except Exception as e:165        print(f"Error initializing ChatOpenAI: {str(e)}")166        await cl.Message(content=f"Error initializing ChatOpenAI: {str(e)}").send()167        return168 169    # Create a chain170    retrieval_augmented_qa_pipeline = RetrievalAugmentedQAPipeline(171        vector_db_retriever=vector_db,172        llm=chat_openai173    )174    175    # Let the user know that the system is ready176    msg.content = f"Processing `{file.name}` done. You can now ask questions!"177    await msg.update()178 179    cl.user_session.set("chain", retrieval_augmented_qa_pipeline)180 181 182@cl.on_message183async def main(message):184    chain = cl.user_session.get("chain")185    if not chain:186        await cl.Message(content="Error: Chat session not initialized. Please try uploading the file again.").send()187        return188 189    msg = cl.Message(content="")190    try:191        result = await chain.arun_pipeline(message.content)192        print(f"Retrieved {len(result['context'])} relevant contexts")193        194        async for stream_resp in result["response"]:195            await msg.stream_token(stream_resp)196 197        await msg.send()198    except Exception as e:199        print(f"Error in chat pipeline: {str(e)}")200        await cl.Message(content=f"Error processing your question: {str(e)}").send()