Update app.py

app.py

@@ -1,29 +1,16 @@
 import os
-from openai import AsyncOpenAI  # importing openai for API usage
-import chainlit as cl  # importing chainlit for our app
-from chainlit.prompt import Prompt, PromptMessage  # importing prompt tools
-from chainlit.playground.providers import ChatOpenAI  # importing ChatOpenAI tools
-from dotenv import load_dotenv
-import asyncio
-import textract  # importing textract for document processing
+import chainlit as cl  # Chainlit import
+from openai import AsyncOpenAI  # OpenAI API
+import textract  # For extracting text from documents
 
 # Load OpenAI API key from environment variables
 api_key = os.getenv("OPENAI_API_KEY")
 
-# Templates for ChatOpenAI interaction
-system_template = """You are a helpful assistant who always speaks in a pleasant tone and answers based on the provided document!
-"""
-
-user_template = """{input}
-Use the document content to respond to the user query step by step.
-"""
-
 # Function to extract text from uploaded documents
-async def extract_text_from_file(file):
-    # Extract text from the uploaded file using textract
-    return textract.process(file).decode('utf-8')
+async def extract_text_from_file(file_path):
+    return textract.process(file_path).decode('utf-8')
 
-# Function to start the chat session
+# Chat initialization
 @cl.on_chat_start  
 async def start_chat():
     settings = {
@@ -34,72 +21,51 @@ async def start_chat():
         "frequency_penalty": 0,
         "presence_penalty": 0,
     }
-
     cl.user_session.set("settings", settings)
 
-    # Welcome message with file upload option
     await cl.Message(content="Welcome! Please upload a document to begin.").send()
 
-# Function to handle uploaded files
-@cl.on_file_upload
-async def on_file_upload(files: list):
-    # Handle the uploaded files, assuming it's the first file
-    if files:
-        file = files[0]
-        file_content = await extract_text_from_file(file.path)
-
-        # Save document content in session for later use
-        cl.user_session.set("document_content", file_content)
-        
-        # Inform the user that the document was successfully uploaded
-        await cl.Message(content=f"Document '{file.name}' uploaded successfully! You can now ask questions based on the document content.").send()
-
-# Function to handle user messages
-@cl.on_message  
+# This handles incoming files manually (instead of on_file_upload)
+@cl.on_message
 async def main(message: cl.Message):
-    settings = cl.user_session.get("settings")
-    document_content = cl.user_session.get("document_content", "")
-
-    if not document_content:
-        # If no document is uploaded, prompt the user to upload one
-        await cl.Message(content="Please upload a document first.").send()
-        return
+    if message.attachments:  # Check if the message contains file attachments
+        uploaded_file = message.attachments[0]
+        file_path = uploaded_file['path']
 
-    client = AsyncOpenAI()
-
-    prompt = Prompt(
-        provider=ChatOpenAI.id,
-        messages=[
-            PromptMessage(
-                role="system",
-                template=system_template,
-                formatted=system_template,
-            ),
-            PromptMessage(
-                role="user",
-                template=user_template,
-                formatted=user_template.format(input=message.content),
-            ),
-        ],
-        inputs={"input": message.content, "document": document_content},
-        settings=settings,
-    )
-
-    msg = cl.Message(content="")
-
-    # Call OpenAI
-    async for stream_resp in await client.chat.completions.create(
-        messages=[m.to_openai() for m in prompt.messages], stream=True, **settings
-    ):
-        token = stream_resp.choices[0].delta.content
-        if not token:
-            token = ""
-        await msg.stream_token(token)
+        # Extract text from the uploaded document
+        file_content = await extract_text_from_file(file_path)
+        
+        # Store document content in user session
+        cl.user_session.set("document_content", file_content)
 
-    # Update the prompt object with the completion
-    prompt.completion = msg.content
-    msg.prompt = prompt
+        await cl.Message(content=f"Document '{uploaded_file['name']}' uploaded successfully! You can now ask questions based on the document content.").send()
+    else:
+        document_content = cl.user_session.get("document_content", "")
+        
+        if not document_content:
+            await cl.Message(content="Please upload a document first.").send()
+            return
 
-    # Send and close the message stream
-    await msg.send()
+        settings = cl.user_session.get("settings")
+        client = AsyncOpenAI()
 
+        # Create the prompt for OpenAI based on user message and document content
+        prompt = f"Document Content: {document_content}\n\nUser Query: {message.content}"
+        
+        msg = cl.Message(content="")
+
+        # Send prompt to OpenAI and stream response
+        async for stream_resp in await client.chat.completions.create(
+            model=settings["model"],
+            messages=[{"role": "system", "content": "Answer based on the provided document."},
+                      {"role": "user", "content": prompt}],
+            stream=True,
+            **settings,
+        ):
+            token = stream_resp.choices[0].delta.content
+            if not token:
+                token = ""
+            await msg.stream_token(token)
+
+        # Send final response
+        await msg.send()