Update app.py

Fix UI/UX and check if NEBIUS_API_KEY is the issue

app.py

@@ -2,96 +2,153 @@ import gradio as gr
 import PyPDF2
 import os
 from openai import OpenAI
+import sys # Import sys to write to stderr
 
-# Function to extract text from PDF
-def extract_text_from_pdf(pdf_file):
-    reader = PyPDF2.PdfReader(pdf_file)
-    text = ""
-    for page in reader.pages:
-        if page == reader.pages[0]:
-            continue
-        text += page.extract_text() + "\n"
-    return text
-
-# Mock LLM response function
-# Replace this with your actual LLM integration
+# --- Configuration & Client Initialization ---
+# It's good practice to check for the API key at startup.
+# This provides a clear error if the secret isn't set in Hugging Face Spaces.
+NEBIUS_API_KEY = os.getenv("NEBIUS_API_KEY")
+if not NEBIUS_API_KEY:
+    # Use a more descriptive error message.
+    raise ValueError("API Key not found. Please set the NEBIUS_API_KEY secret in your Hugging Face Space settings.")
 
+# Initialize the OpenAI client with your custom endpoint.
+# Ensure this base_url is correct and publicly accessible.
 client = OpenAI(
     base_url="https://api.studio.nebius.com/v1/",
-    api_key=os.getenv("NEBIUS_API_KEY")
+    api_key=NEBIUS_API_KEY
 )
 
-def answer_question(pdf_text, question, history):
-    # You may want to truncate pdf_text if it's too long for the model
-    context = pdf_text[:16000]  # adjust as needed
-    messages = [
-        {"role": "system", "content": '''You are a helpful assistant who specializes in body composition, diet, and exercise. 
-         Answer questions based on the provided document. Encourage the user to seek a professional if they have serious concerns whenever appropriate.'''},
-        {"role": "user", "content": f"Document: {context}"},
-        {"role": "user", "content": f"Question: {question}"}
-    ]
+# --- Core Functions ---
+
+def extract_text_from_pdf(pdf_file):
+    """Extracts text from an uploaded PDF file object."""
+    # pdf_file is a temporary file object from Gradio.
+    if not pdf_file:
+        return ""
+    try:
+        reader = PyPDF2.PdfReader(pdf_file.name)
+        text = ""
+        # Extract text from all pages except the first one.
+        for i, page in enumerate(reader.pages):
+            if i == 0:
+                continue
+            page_text = page.extract_text()
+            if page_text:
+                text += page_text + "\n"
+        return text
+    except Exception as e:
+        print(f"Error reading PDF: {e}", file=sys.stderr)
+        return ""
+
+
+def get_llm_answer(pdf_text, question, history):
+    """
+    Sends the context, history, and question to the LLM and returns the answer.
+    """
+    # Truncate the PDF text to avoid exceeding the model's context limit.
+    context = pdf_text[:16000]
+
+    # The system prompt guides the model's behavior.
+    system_prompt = '''You are a helpful assistant who specializes in body composition, diet, and exercise.
+    Answer questions based on the provided document. Encourage the user to seek a professional
+    if they have serious concerns whenever appropriate.'''
+
+    # Construct the message payload for the API.
+    messages = [{"role": "system", "content": system_prompt}]
+    
+    # Add document context if available.
+    if context:
+        messages.append({"role": "user", "content": f"Use the following document to answer my question:\n\n{context}"})
+
+    # Add the conversation history.
     if history:
-        for msg in history:
-            if msg["role"] in ["user", "assistant"]:
-                messages.append(msg)
-    # Add the new user question
-    messages.append({"role": "user", "content": question}) 
-
-    response = client.chat.completions.create(
-        model="meta-llama/Meta-Llama-3.1-70B-Instruct",
-        temperature=0.6,
-        top_p=0.95,
-        messages=messages
-    )
-    return response.choices[0].message.content
+        messages.extend(history)
+
+    # Add the latest user question.
+    messages.append({"role": "user", "content": question})
 
-# Gradio chatbot logic
+    try:
+        response = client.chat.completions.create(
+            model="meta-llama/Meta-Llama-3.1-70B-Instruct",
+            temperature=0.6,
+            top_p=0.95,
+            messages=messages
+        )
+        return response.choices[0].message.content
+    except Exception as e:
+        print(f"Error calling OpenAI API: {e}", file=sys.stderr)
+        # Return a user-friendly error message.
+        return "Sorry, I encountered an error while trying to generate a response. Please check the logs."
+
+
+# --- Gradio Interface Logic ---
+
+# Use a class to manage state (the extracted PDF text).
 class PDFChatbot:
     def __init__(self):
         self.pdf_text = None
+        self.pdf_filename = None
 
     def upload_pdf(self, pdf_file):
         if pdf_file is None:
-            return "No PDF uploaded."
+            return "Status: No PDF uploaded."
+        
         self.pdf_text = extract_text_from_pdf(pdf_file)
-        return "PDF uploaded and processed. You can now ask questions."
-
-    def chat(self, message, history):
-        if self.pdf_text is None:
-            return "Please upload a PDF first.", history
-        # Ensure history is a list of dicts with 'role' and 'content'
-        if history is None:
-            history = []
-        # Only keep user/assistant messages for context
-        context_history = [msg for msg in history if msg["role"] in ["user", "assistant"]]
-        answer = answer_question(self.pdf_text, message, context_history)
-        print("User:", message)
-        print("LLM Answer:", answer)
-        history = history + [
-            {"role": "user", "content": message},
-            {"role": "assistant", "content": answer}
-        ]
-        return "", history
+        self.pdf_filename = os.path.basename(pdf_file.name)
 
-pdf_bot = PDFChatbot()
+        if not self.pdf_text:
+            return f"Status: Could not extract text from {self.pdf_filename}. It might be empty, scanned, or protected."
+            
+        return f"Status: Successfully processed {self.pdf_filename}. You can now ask questions."
 
-with gr.Blocks() as demo:
-    gr.Markdown("Body Composition Agent\nUpload a PDF of your recent body spec or any other document you use to track your body composition and ask questions about its content.")
-    pdf_file = gr.File(label="Upload PDF", file_types=[".pdf"])
-    upload_btn = gr.Button("Process PDF")
-    upload_status = gr.Textbox(label="Status", interactive=False)
-    chatbot = gr.Chatbot(type="messages")
-    msg = gr.Textbox(label="Your question:")
-    submit_btn = gr.Button("Submit")
+    def chat(self, user_message, history):
+        if not self.pdf_text:
+            # Add an instruction to the chatbot window if no PDF is uploaded.
+            history.append([user_message, "Please upload a PDF document first."])
+            return "", history
 
-    def handle_upload(pdf):
-        return pdf_bot.upload_pdf(pdf)
+        # Get the answer from the LLM.
+        answer = get_llm_answer(self.pdf_text, user_message, history)
+        
+        # Append the user message and the assistant's answer to the history.
+        history.append([user_message, answer])
+        
+        # Return an empty string to clear the input textbox and the updated history.
+        return "", history
+
+# Instantiate the bot.
+pdf_bot = PDFChatbot()
 
-    def handle_chat(message, history):
-        return pdf_bot.chat(message, history)
+# Build the Gradio UI.
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# Body Composition Agent\nUpload a document about your body composition and ask questions about its content.")
+    
+    with gr.Row():
+        with gr.Column(scale=1):
+            pdf_file = gr.File(label="Upload PDF", file_types=[".pdf"])
+            upload_btn = gr.Button("Process PDF", variant="primary")
+            upload_status = gr.Textbox(label="Status", interactive=False, value="Status: Waiting for PDF...")
+        
+        with gr.Column(scale=2):
+            chatbot = gr.Chatbot(label="Chat History", bubble_full_width=False, height=500)
+            msg_textbox = gr.Textbox(label="Your Question:", interactive=True, placeholder="Type your question here...")
+            # Clear button is useful for starting a new conversation.
+            clear_btn = gr.ClearButton([msg_textbox, chatbot], value="Clear Chat")
 
-    upload_btn.click(handle_upload, inputs=pdf_file, outputs=upload_status)
-    submit_btn.click(handle_chat, inputs=[msg, chatbot], outputs=[msg, chatbot])
-    msg.submit(handle_chat, inputs=[msg, chatbot], outputs=[msg, chatbot])
+    # Wire up the event listeners.
+    upload_btn.click(
+        fn=pdf_bot.upload_pdf,
+        inputs=[pdf_file],
+        outputs=[upload_status]
+    )
+    
+    # Allow submitting questions with Enter key or the button.
+    msg_textbox.submit(
+        fn=pdf_bot.chat,
+        inputs=[msg_textbox, chatbot],
+        outputs=[msg_textbox, chatbot]
+    )
 
-demo.launch()
+# Launch the app.
+demo.launch(debug=True) # Use debug=True to see errors in the console.