Update app.py

app.py

@@ -1,35 +1,84 @@
 import gradio as gr
 import requests
-import json
+import uuid
+import os
+from datetime import datetime
 
-API_URL = "https://dssa9yu3o7.execute-api.ap-south-1.amazonaws.com/dev/generatePresignedUrl"
+# API Gateway endpoints (secured, NOT exposed in UI)
+PRESIGN_API_URL = os.environ.get("PRESIGN_API_URL")      # e.g., Lambda@API Gateway for presigned S3 URL
+INFERENCE_API_URL = os.environ.get("INFERENCE_API_URL")  # e.g., Lambda B endpoint for LLM QA
 
-def upload_pdf(file):
-    filename = file.name
-    content_type = "application/pdf"
+# Session state
+user_id = "johndoe"  # This could also be a login-generated ID
+session_id = datetime.now().strftime("%Y%m%d_%H%M%S") + "_" + str(uuid.uuid4())[:8]
+
+upload_success = False
 
-    # Step 1: Get pre-signed URL
-    response = requests.post(API_URL, json={"filename": filename, "contentType": content_type})
+def upload_pdf(file):
+    global upload_success
+    filename = os.path.basename(file.name)
+    
+    # Step 1: Get presigned URL from API
+    response = requests.post(PRESIGN_API_URL, json={
+        "filename": filename,
+        "user_id": user_id,
+        "session_id": session_id
+    })
+    
     if response.status_code != 200:
-        return "Failed to get upload URL"
+        return "❌ Failed to get presigned URL"
 
-    upload_url = response.json().get("uploadURL")
+    data = response.json()
+    presigned_url = data["url"]
+    content_type = data.get("content_type", "application/pdf")
 
     # Step 2: Upload to S3
     with open(file.name, "rb") as f:
-        upload_response = requests.put(upload_url, data=f, headers={"Content-Type": content_type})
+        upload_resp = requests.put(presigned_url, data=f, headers={"Content-Type": content_type})
 
-    if upload_response.status_code == 200:
-        return "✅ Upload successful!"
+    if upload_resp.status_code == 200:
+        upload_success = True
+        return f"✅ Uploaded and indexing triggered! You can now ask questions."
     else:
-        return f"❌ Upload failed: {upload_response.status_code}"
+        return "❌ Upload failed!"
+
+def chat_with_doc(message, history):
+    if not upload_success:
+        return "Please upload a PDF first.", history
+
+    if message.lower() == "exit":
+        return "🛑 Session ended. Thank you!", []
+
+    payload = {
+        "question": message,
+        "user_id": user_id,
+        "session_id": session_id,
+        "chat_history": history
+    }
+
+    resp = requests.post(INFERENCE_API_URL, json=payload)
+    
+    if resp.status_code != 200:
+        return "❌ Error calling LLM. Try again later.", history
+
+    answer = resp.json()["answer"]
+    history.append((message, answer))
+    return "", history
+
+with gr.Blocks() as demo:
+    gr.Markdown("## 📄 Upload your PDF and Ask Questions")
+    
+    with gr.Row():
+        file_input = gr.File(label="Upload PDF", file_types=[".pdf"])
+        upload_button = gr.Button("Upload")
+        upload_output = gr.Textbox(label="Upload Status")
+
+    upload_button.click(upload_pdf, inputs=file_input, outputs=upload_output)
+
+    chatbot = gr.Chatbot(label="Document Chat")
+    question = gr.Textbox(label="Ask a question (or type 'exit' to end)", placeholder="Ask something about the document...")
+    ask_button = gr.Button("Send")
 
-ui = gr.Interface(
-    fn=upload_pdf,
-    inputs=gr.File(file_types=[".pdf"]),
-    outputs="text",
-    title="Upload PDF to S3",
-    description="Select a PDF file to upload to AWS S3 via pre-signed URL"
-)
+    ask_button.click(chat_with_doc, inputs=[question, chatbot], outputs=[question, chatbot])
 
-ui.launch()
+demo.launch()