Update app.py

app.py

@@ -2,13 +2,15 @@ import gradio as gr
 import requests
 import mimetypes
 import json, os
+import asyncio
+import aiohttp
 
 LLM_API = os.environ.get("LLM_API")
 LLM_URL = os.environ.get("LLM_URL")
 
 USER_ID = "HuggingFace Space"  # Placeholder user ID
 
-def send_chat_message(LLM_URL, LLM_API, category, file_id):
+async def send_chat_message(LLM_URL, LLM_API, category, file_id):
     payload = {
         "inputs": {},
         "query": category,
@@ -24,59 +26,55 @@ def send_chat_message(LLM_URL, LLM_API, category, file_id):
         ]
     }
     print("Sending chat message payload:", payload)  # Debug information
-    response = requests.post(
-        f"{LLM_URL}/chat-messages",
-        headers={"Authorization": f"Bearer {LLM_API}"},
-        json=payload,
-        stream=True  # Enable streaming
-    )
-    print("Request URL:", f"{LLM_URL}/chat-messages")
-    print("Response status code:", response.status_code)
-    if response.status_code == 404:
-        return "Error: Endpoint not found (404)"
-    
-    # Handle the stream of events
-    last_thought = None
-    try:
-        for line in response.iter_lines(decode_unicode=True):
-            if line:
-                try:
-                    data = json.loads(line.split("data: ")[1])
-                    if data.get("event") == "agent_thought":
-                        last_thought = data.get("thought")
-                except (IndexError, json.JSONDecodeError):
-                    continue
-    except requests.exceptions.JSONDecodeError:
-        return "Error: Invalid JSON response"
-    
-    if last_thought:
-        # Structure the thought text
-        return last_thought.strip()
-    else:
-        return "Error: No thought found in the response"
+    async with aiohttp.ClientSession() as session:
+        async with session.post(
+            f"{LLM_URL}/chat-messages",
+            headers={"Authorization": f"Bearer {LLM_API}"},
+            json=payload
+        ) as response:
+            print("Request URL:", f"{LLM_URL}/chat-messages")
+            print("Response status code:", response.status)
+            if response.status == 404:
+                return "Error: Endpoint not found (404)"
+            
+            last_thought = None
+            async for line in response.content:
+                if line:
+                    try:
+                        data = json.loads(line.split(b"data: ")[1].decode("utf-8"))
+                        if data.get("event") == "agent_thought":
+                            last_thought = data.get("thought")
+                    except (IndexError, json.JSONDecodeError):
+                        continue
+            
+            if last_thought:
+                return last_thought.strip()
+            else:
+                return "Error: No thought found in the response"
 
-def upload_file(LLM_URL, LLM_API, file_path, user_id):
+async def upload_file(LLM_URL, LLM_API, file_path, user_id):
     if not os.path.exists(file_path):
         return f"Error: File {file_path} not found"
     mime_type, _ = mimetypes.guess_type(file_path)
     with open(file_path, 'rb') as f:
-        response = requests.post(
-            f"{LLM_URL}/files/upload",
-            headers={"Authorization": f"Bearer {LLM_API}"},
-            files={"file": (file_path, f, mime_type)},
-            data={"user": user_id}
-        )
-    print("Upload response status code:", response.status_code)  # Debug information
-    if response.status_code == 404:
-        return "Error: Endpoint not found (404)"
-    print("Raw upload response text:", response.text)  # Debug information
-    try:
-        return response.json()
-    except requests.exceptions.JSONDecodeError:
-        return "Error: Invalid JSON response"
+        async with aiohttp.ClientSession() as session:
+            async with session.post(
+                f"{LLM_URL}/files/upload",
+                headers={"Authorization": f"Bearer {LLM_API}"},
+                data={"file": (file_path, f, mime_type), "user": user_id}
+            ) as response:
+                print("Upload response status code:", response.status)  # Debug information
+                if response.status == 404:
+                    return "Error: Endpoint not found (404)"
+                response_text = await response.text()
+                print("Raw upload response text:", response_text)  # Debug information
+                try:
+                    return json.loads(response_text)
+                except json.JSONDecodeError:
+                    return "Error: Invalid JSON response"
 
-def handle_input(file_path, category):
-    upload_response = upload_file(LLM_URL, LLM_API, file_path, USER_ID)
+async def handle_input(file_path, category):
+    upload_response = await upload_file(LLM_URL, LLM_API, file_path, USER_ID)
     print("Upload response:", upload_response)  # Debug information
     if "error" in upload_response:
         return upload_response
@@ -84,7 +82,7 @@ def handle_input(file_path, category):
     if not file_id:
         return "Error: No file ID returned from upload"
     
-    chat_response = send_chat_message(LLM_URL, LLM_API, category, file_id)
+    chat_response = await send_chat_message(LLM_URL, LLM_API, category, file_id)
     print("Chat response:", chat_response)  # Debug information
     return chat_response