Update app.py

app.py

@@ -3,7 +3,7 @@ import base64
 import requests
 from pathlib import Path
 from typing import Optional, List, Dict, Any
-from flask import Flask, request, Response
+from flask import Flask, request, jsonify
 
 app = Flask(__name__)
 
@@ -24,6 +24,12 @@ def encode_image_to_base64_from_url(image_url: str, timeout: int = 30) -> str:
     return base64.b64encode(resp.content).decode("utf-8")
 
 def normalize_messages_to_prompt(messages: List[Dict[str, str]]) -> str:
+    # Convert structured chat history into a single prompt string
+    # Expected roles: system, user, assistant
+    # Output example:
+    # System: ...
+    # User: ...
+    # Assistant: ...
     lines = []
     for msg in messages:
         role = str(msg.get("role", "")).strip().lower()
@@ -37,20 +43,25 @@ def normalize_messages_to_prompt(messages: List[Dict[str, str]]) -> str:
         else:
             prefix = "User"
         lines.append(f"{prefix}: {content}")
+    # Optionally add a final cue for the assistant to respond
     return "\n".join(lines).strip()
 
 def extract_output(obj: dict) -> str:
+    # Try common shapes:
+    # 1) { data: { outputs: [ { output: "..." } ] } }
     try:
         out = obj["data"]["outputs"][0]["output"]
         if isinstance(out, str):
             return out
     except Exception:
         pass
+    # 2) Flat { "output": "..." }
     if isinstance(obj, dict) and isinstance(obj.get("output"), str):
         return obj["output"]
+    # Fallback: empty
     return ""
 
-def call_roboflow_playground(image_b64: str, prompt: str, model: str = DEFAULT_MODEL, timeout: int = 60) -> dict:
+def call_roboflow_playground(image_b64: str, prompt: str, model: str = DEFAULT_MODEL, timeout: int = 60) -> Dict[str, Any]:
     payload = {
         "inputs": {
             "image": {"type": "base64", "value": image_b64},
@@ -73,66 +84,60 @@ def call_roboflow_playground(image_b64: str, prompt: str, model: str = DEFAULT_M
 @app.route("/infer", methods=["POST"])
 def infer():
     try:
-        data = request.get_json(force=True) or {}
+        data = request.get_json(force=True, silent=False) or {}
 
-        # Image input
+        # Image sources: image_base64, image_url, or image_path (optional local)
         image_b64 = (data.get("image_base64") or "").strip()
         image_url = (data.get("image_url") or "").strip()
         image_path = (data.get("image_path") or "").strip()
 
         if not image_b64:
             if image_url:
-                image_b64 = encode_image_to_base64_from_url(image_url)
+                try:
+                    image_b64 = encode_image_to_base64_from_url(image_url)
+                except Exception as e:
+                    return jsonify({"error": f"Failed to fetch/encode image_url: {str(e)}"}), 400
             elif image_path:
-                image_b64 = encode_image_to_base64_from_path(image_path)
+                try:
+                    image_b64 = encode_image_to_base64_from_path(image_path)
+                except Exception as e:
+                    return jsonify({"error": f"Failed to read/encode image_path: {str(e)}"}), 400
             else:
-                return Response("Missing image. Provide image_base64, image_url, or image_path.", status=400)
+                return jsonify({"error": "Provide one of: image_base64, image_url, or image_path"}), 400
 
-        # Messages -> prompt
+        # Messages to prompt
         messages = data.get("messages", [])
-        if not isinstance(messages, list):
-            return Response("messages must be a list of {role, content}.", status=400)
+        if not isinstance(messages, list) or not all(isinstance(m, dict) for m in messages):
+            return jsonify({"error": "messages must be a list of {role, content}"}), 400
 
         prompt = normalize_messages_to_prompt(messages)
         if not prompt:
-            return Response("messages is empty or has no content.", status=400)
+            return jsonify({"error": "messages is empty or has no valid content"}), 400
 
+        # Model (optional override)
         model = (data.get("model") or DEFAULT_MODEL).strip() or DEFAULT_MODEL
 
-        upstream = call_roboflow_playground(image_b64=image_b64, prompt=prompt, model=model)
-        output_text = extract_output(upstream).strip()
+        # Call external API
+        raw_response = call_roboflow_playground(image_b64=image_b64, prompt=prompt, model=model)
 
-        if not output_text:
-            return Response("No output from model.", status=502)
+        # Extract simplified output
+        output_text = extract_output(raw_response)
 
-        # Return only the AI response as plain text
-        return Response(output_text, mimetype="text/plain", status=200)
+        return jsonify({
+            "output": output_text,
+            "raw_response": raw_response
+        }), 200
 
     except requests.HTTPError as e:
         status = getattr(e.response, "status_code", 502)
         text = getattr(e.response, "text", "")
-        return Response(f"Upstream HTTP error ({status}): {text}", status=502)
+        return jsonify({"error": "HTTP error from upstream", "status": status, "response_text": text}), 502
     except Exception as e:
-        return Response(str(e), status=500)
+        return jsonify({"error": str(e)}), 500
 
 def main():
+    # Run Flask on port 7860
     app.run(host="0.0.0.0", port=7860, debug=False)
 
 if __name__ == "__main__":
-    main()
-
-Client test (super simple):
-import requests
-API_URL = "https://corvo-ai-xx-mistral.hf.space/infer"
-
-payload = {
-    "image_url": "https://images.unsplash.com/photo-1518791841217-8f162f1e1131",
-    "messages": [
-        {"role": "system", "content": "You are a helpful vision assistant."},
-        {"role": "user", "content": "What do you see in this image?"}
-    ]
-}
-
-r = requests.post(API_URL, json=payload, timeout=60)
-print(r.status_code)
-print(r.text)
+    main()