Update app.py

app.py

@@ -1,10 +1,9 @@
-# app.py (minimal SSE + CORS)
 from flask import Flask, request, Response, jsonify, stream_with_context
 from flask_cors import CORS
 import os, json, httpx
 
 app = Flask(__name__)
-CORS(app)  # simple: allow all origins for testing
+CORS(app)
 
 OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_NEW")
 
@@ -31,38 +30,54 @@ def _payload(q: str):
 
 def stream_openrouter(q: str):
     if not OPENROUTER_API_KEY:
-        yield "data: Missing OPENROUTER_API_NEW\n\n"; return
-    with httpx.Client(timeout=None) as client:
-        with client.stream("POST", "https://openrouter.ai/api/v1/chat/completions",
-                           headers=_headers(), json=_payload(q)) as r:
-            r.raise_for_status()
-            for raw in r.iter_lines():
-                if not raw: 
-                    continue
-                data = raw[5:].strip() if raw.startswith(b"data:") else raw.strip()
-                if data == b"[DONE]":
-                    yield "data: [DONE]\n\n"
-                    break
-                try:
-                    obj = json.loads(data.decode("utf-8"))
-                    delta = obj.get("choices", [{}])[0].get("delta", {}).get("content")
-                    if delta:
-                        # Send plain text chunks (still framed as SSE data)
-                        yield f"data: {json.dumps(delta)}\n\n"
-                except Exception:
-                    pass
+        yield "data: Error: Missing API key\n\n"
+        return
+        
+    try:
+        with httpx.Client(timeout=30.0) as client:
+            with client.stream("POST", 
+                             "https://openrouter.ai/api/v1/chat/completions",
+                             headers=_headers(),
+                             json=_payload(q)) as response:
+                
+                response.raise_for_status()
+                
+                for line in response.iter_lines():
+                    if line.startswith('data: '):
+                        data = line[6:]  # Remove 'data: ' prefix
+                        if data == '[DONE]':
+                            break
+                        try:
+                            json_data = json.loads(data)
+                            if 'choices' in json_data:
+                                content = json_data['choices'][0].get('delta', {}).get('content')
+                                if content:
+                                    yield f"data: {json.dumps(content)}\n\n"
+                        except json.JSONDecodeError:
+                            continue
+                
+                yield "data: [DONE]\n\n"
+                
+    except httpx.HTTPError as e:
+        yield f"data: Error: {str(e)}\n\n"
+    except Exception as e:
+        yield f"data: Unexpected error: {str(e)}\n\n"
 
 @app.post("/ask")
 def ask():
     payload = request.get_json(silent=True) or {}
-    q = (payload.get("question") or payload.get("userPrompt") or "").strip()
+    q = payload.get("question", "").strip()
     if not q:
         q = "Say hello in one short sentence."
-    resp = Response(stream_with_context(stream_openrouter(q)), mimetype="text/event-stream")
-    resp.headers["Cache-Control"] = "no-cache"
-    resp.headers["X-Accel-Buffering"] = "no"
-    resp.headers["Connection"] = "keep-alive"
-    return resp
+        
+    return Response(
+        stream_with_context(stream_openrouter(q)),
+        mimetype="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "X-Accel-Buffering": "no"
+        }
+    )
 
 if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=7860, debug=False, threaded=True)
+    app.run(host="0.0.0.0", port=7860, debug=False)