.

app.py

@@ -2,8 +2,7 @@ import gradio as gr
 import requests
 import os
 import json
-import re
-import html
+import sseclient
 
 # Set up the API endpoint and key
 API_URL = os.getenv("RUNPOD_API_URL")
@@ -14,34 +13,11 @@ headers = {
     "Content-Type": "application/json"
 }
 
-def escape_html(text):
-    return html.escape(text)
+# Fixed system prompt
+SYSTEM_PROMPT = "You an advanced artificial intelligence system, capable of <thinking> <reflection> and you output a brief and small to the point <output>."
 
-def format_response_for_display(text):
-    # Escape HTML entities
-    text = escape_html(text)
-    
-    # Format <thinking> and <reflection> tags (case-insensitive)
-    text = re.sub(r'(?i)(<thinking>)(.*?)(</thinking>)', 
-                  r'<span style="font-family: monospace; color: blue;">\1</span><i>\2</i><span style="font-family: monospace; color: blue;">\3</span>', 
-                  text, flags=re.DOTALL)
-    text = re.sub(r'(?i)(<reflection>)(.*?)(</reflection>)', 
-                  r'<span style="font-family: monospace; color: green;">\1</span><i>\2</i><span style="font-family: monospace; color: green;">\3</span>', 
-                  text, flags=re.DOTALL)
-    
-    # Remove <output> tags but keep content
-    text = re.sub(r'(?i)<output>(.*?)</output>', r'\1', text, flags=re.DOTALL)
-    
-    # Replace newlines with <br> tags
-    text = text.replace('\n', '<br>')
-    
-    return text
-
-def respond(message, history, system_message, max_tokens, temperature, top_p):
-    if system_message is not None:
-        messages = [{"role": "system", "content": system_message}]
-    else:
-        messages = []
+def stream_response(message, history, max_tokens, temperature, top_p):
+    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
     
     for human, assistant in history:
         messages.append({"role": "user", "content": human})
@@ -54,55 +30,44 @@ def respond(message, history, system_message, max_tokens, temperature, top_p):
         "messages": messages,
         "max_tokens": max_tokens,
         "temperature": temperature,
-        "top_p": top_p
+        "top_p": top_p,
+        "stream": True
     }
     
     try:
-        print(f"Sending request to API: {API_URL}")
-        print(f"Request data: {json.dumps(data, indent=2)}")
-        
-        response = requests.post(API_URL, headers=headers, json=data)
-        
-        print(f"Raw API Response: {response.text}")
-        print(f"Response Status Code: {response.status_code}")
-        
+        response = requests.post(API_URL, headers=headers, json=data, stream=True)
         response.raise_for_status()
+        client = sseclient.SSEClient(response)
         
-        response_json = response.json()
-        print(f"Formatted API Response: {json.dumps(response_json, indent=2)}")
+        full_response = ""
+        for event in client.events():
+            if event.data != "[DONE]":
+                try:
+                    chunk = json.loads(event.data)
+                    if 'choices' in chunk and len(chunk['choices']) > 0:
+                        content = chunk['choices'][0]['delta'].get('content', '')
+                        full_response += content
+                        # Replace < and > with their HTML entities
+                        display_content = content.replace('<', '&lt;').replace('>', '&gt;')
+                        yield display_content
+                except json.JSONDecodeError:
+                    print(f"Failed to decode JSON: {event.data}")
         
-        if 'choices' in response_json and len(response_json['choices']) > 0:
-            content = response_json['choices'][0]['message']['content']
-            formatted_content = format_response_for_display(content)
-            print(f"Formatted content for display: {formatted_content}")  # For debugging
-            return formatted_content
-        else:
-            return f"Error: Unexpected response format. Full response: {response_json}"
-    
     except requests.exceptions.RequestException as e:
-        print(f"Request Exception: {str(e)}")
-        return f"Error: {str(e)}"
-    except ValueError as e:
-        print(f"JSON Parsing Error: {str(e)}")
-        return f"Error: Invalid JSON response. {str(e)}"
-    except KeyError as e:
-        print(f"Key Error: {str(e)}")
-        return f"Error: Unexpected response structure. Missing key: {str(e)}"
+        yield f"Error: {str(e)}"
     except Exception as e:
-        print(f"Unexpected Error: {str(e)}")
-        return f"Unexpected error: {str(e)}"
+        yield f"Unexpected error: {str(e)}"
 
 demo = gr.ChatInterface(
-    respond,
+    stream_response,
     additional_inputs=[
-        gr.Textbox(label="System message"),
         gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max tokens"),
         gr.Slider(minimum=0.1, maximum=2.0, value=0.7, step=0.1, label="Temperature"),
         gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)"),
     ],
-    css=".message-wrap { white-space: pre-wrap; }"
 )
 
 if __name__ == "__main__":
     print(f"Starting application with API URL: {API_URL}")
+    print(f"Using system prompt: {SYSTEM_PROMPT}")
     demo.launch()