Update app.py

app.py

@@ -1,72 +1,259 @@
-from flask import Flask, request, Response, jsonify
-from flask_cors import CORS
-import requests
+# ─────────────────────────────────────────────────────────────────────────────
+# Mira Proxy Space  •  Docker SDK  •  No GPU needed
+#
+# Routes:
+#   POST /v1/chat/completions                →  mira-1-large  (default)
+#   POST /v1/{model_key}/chat/completions    →  any registered model
+#   GET  /v1/models                          →  list available models
+#   GET  /health                             →  proxy + upstream health
+#   GET  /                                   →  info page
+#
+# Env secrets (set in Space settings):
+#   HF_TOKEN   — used both to call upstream spaces AND to optionally
+#                gate inbound requests (if AUTH_REQUIRED=true)
+#   AUTH_REQUIRED — set to "true" to require Bearer token on inbound calls
+# ─────────────────────────────────────────────────────────────────────────────
+
+from __future__ import annotations
+
 import os
+import time
+import logging
+from functools import wraps
+
+import requests
+from flask import Flask, request, Response, jsonify, abort
+from flask_cors import CORS
+
+# ── Logging ───────────────────────────────────────────────────────────────────
+logging.basicConfig(
+    level=logging.INFO,
+    format="%(asctime)s [%(levelname)s] %(message)s",
+)
+log = logging.getLogger(__name__)
 
+# ── App ───────────────────────────────────────────────────────────────────────
 app = Flask(__name__)
 CORS(app, origins=["*"])
 
-HF_TOKEN = os.environ.get("HF_TOKEN", "")
+# ── Config ────────────────────────────────────────────────────────────────────
+HF_TOKEN      = os.environ.get("HF_TOKEN", "")
+AUTH_REQUIRED = os.environ.get("AUTH_REQUIRED", "false").lower() == "true"
 
-# Map each model key to its real HF endpoint. Add new models here only.
-ENDPOINTS = {
-    "mira-1-large": "https://ixvuours3f75nti4.us-east-1.aws.endpoints.huggingface.cloud/v1/",
-    "mira-1-xl": "https://uk68x3qqblj2iz7i.us-east-1.aws.endpoints.huggingface.cloud/v1/",
+# ── Endpoint registry ─────────────────────────────────────────────────────────
+# Each value is the Space root URL (no trailing path).
+# The proxy appends /v1/chat/completions when forwarding.
+# Update these slugs to match your actual Space URLs.
+ENDPOINTS: dict[str, str] = {
+    "mira-1-large": "https://ml-intern-explorers-inf-end-1.hf.space",
+    "mira-1-xl":    "https://ml-intern-explorers-inf-end-2.hf.space",
 }
 
-def proxy_request(endpoint_url, data):
-    headers = {
+DEFAULT_MODEL = "mira-1-large"
+
+# Model metadata for /v1/models response (OpenAI-compatible)
+MODEL_META = {
+    "mira-1-large": {"description": "Mira 1 Large",  "context_window": 8192},
+    "mira-1-xl":    {"description": "Mira 1 XL",     "context_window": 16384},
+}
+
+# ── Auth guard (optional) ─────────────────────────────────────────────────────
+def require_auth(f):
+    @wraps(f)
+    def decorated(*args, **kwargs):
+        if not AUTH_REQUIRED:
+            return f(*args, **kwargs)
+        auth = request.headers.get("Authorization", "")
+        if not auth.startswith("Bearer ") or auth.split(" ", 1)[1] != HF_TOKEN:
+            return jsonify({"error": "Unauthorized"}), 401
+        return f(*args, **kwargs)
+    return decorated
+
+# ── Upstream headers ──────────────────────────────────────────────────────────
+def _upstream_headers() -> dict:
+    return {
         "Authorization": f"Bearer {HF_TOKEN}",
-        "Content-Type": "application/json"
+        "Content-Type":  "application/json",
+        "Accept":        "application/json",
     }
 
-    if data.get('stream', False):
-        resp = requests.post(
-            f"{endpoint_url}chat/completions",
+# ── Core proxy logic ──────────────────────────────────────────────────────────
+def proxy_request(model_key: str, data: dict) -> Response:
+    base = ENDPOINTS[model_key].rstrip("/")
+    target = f"{base}/v1/chat/completions"
+
+    # Ensure the model field in the payload matches what the Space expects
+    data.setdefault("model", model_key)
+
+    log.info(f"→ {model_key}  stream={data.get('stream', False)}")
+
+    if data.get("stream", False):
+        # ── Streaming: forward SSE chunks as they arrive ──────────────────────
+        upstream = requests.post(
+            target,
             json=data,
-            headers=headers,
+            headers=_upstream_headers(),
             stream=True,
-            timeout=300
+            timeout=(10, 300),   # (connect, read)
         )
+        upstream.raise_for_status()
 
         def generate():
-            for chunk in resp.iter_content(chunk_size=None, decode_unicode=False):
-                if chunk:
-                    yield chunk
-
-        return Response(generate(), content_type='text/event-stream')
-    else:
-        resp = requests.post(
-            f"{endpoint_url}chat/completions",
-            json=data,
-            headers=headers,
-            timeout=60
+            try:
+                for chunk in upstream.iter_content(
+                    chunk_size=None, decode_unicode=False
+                ):
+                    if chunk:
+                        yield chunk
+            except Exception as exc:
+                log.error(f"Streaming error for {model_key}: {exc}")
+                yield b"data: [DONE]\n\n"
+
+        return Response(
+            generate(),
+            status=upstream.status_code,
+            content_type="text/event-stream",
+            headers={
+                "Cache-Control":    "no-cache",
+                "X-Accel-Buffering":"no",          # disable nginx buffering
+                "Connection":       "keep-alive",
+            },
         )
-        return jsonify(resp.json())
 
-# Original route — kept as-is so existing Mira-1-Large calls with no path change still work
-@app.route('/v1/chat/completions', methods=['POST'])
+    # ── Non-streaming ─────────────────────────────────────────────────────────
+    upstream = requests.post(
+        target,
+        json=data,
+        headers=_upstream_headers(),
+        timeout=(10, 60),
+    )
+    upstream.raise_for_status()
+    return jsonify(upstream.json())
+
+
+# ── Routes ────────────────────────────────────────────────────────────────────
+
+@app.route("/", methods=["GET"])
+def index():
+    """Human-readable info page."""
+    return jsonify({
+        "service":       "Mira Proxy",
+        "version":       "1.0.0",
+        "models":        list(ENDPOINTS.keys()),
+        "default_model": DEFAULT_MODEL,
+        "endpoints": {
+            "chat":   "POST /v1/chat/completions",
+            "models": "GET  /v1/models",
+            "health": "GET  /health",
+        },
+    })
+
+
+@app.route("/v1/models", methods=["GET"])
+@require_auth
+def list_models():
+    """OpenAI-compatible model listing."""
+    now = int(time.time())
+    return jsonify({
+        "object": "list",
+        "data": [
+            {
+                "id":       key,
+                "object":   "model",
+                "created":  now,
+                "owned_by": "mira",
+                **MODEL_META.get(key, {}),
+            }
+            for key in ENDPOINTS
+        ],
+    })
+
+
+@app.route("/v1/chat/completions", methods=["POST"])
+@require_auth
 def chat_completions_default():
-    data = request.json
+    """
+    Default route — always routes to mira-1-large.
+    Existing callers need zero changes.
+    """
+    data = request.json or {}
+    # If caller already specifies a known model in the payload, honour it.
+    requested = data.get("model", DEFAULT_MODEL)
+    model_key = requested if requested in ENDPOINTS else DEFAULT_MODEL
     try:
-        return proxy_request(ENDPOINTS["mira-1-large"], data)
+        return proxy_request(model_key, data)
+    except requests.HTTPError as e:
+        log.error(f"Upstream HTTP error: {e}")
+        return jsonify({"error": str(e)}), e.response.status_code
+    except requests.Timeout:
+        return jsonify({"error": "upstream timeout"}), 504
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        log.error(f"Proxy error: {e}")
+        return jsonify({"error": str(e)}), 502
 
-# New: per-model route, e.g. /v1/mira-1-xl/chat/completions
-@app.route('/v1/<model_key>/chat/completions', methods=['POST'])
-def chat_completions_by_model(model_key):
+
+@app.route("/v1/<model_key>/chat/completions", methods=["POST"])
+@require_auth
+def chat_completions_by_model(model_key: str):
+    """
+    Per-model route.
+    e.g. POST /v1/mira-1-xl/chat/completions
+    """
     if model_key not in ENDPOINTS:
-        return jsonify({"error": f"Unknown model '{model_key}'"}), 400
-    data = request.json
+        return jsonify({
+            "error":            f"Unknown model '{model_key}'",
+            "available_models": list(ENDPOINTS.keys()),
+        }), 404
+
+    data = request.json or {}
     try:
-        return proxy_request(ENDPOINTS[model_key], data)
+        return proxy_request(model_key, data)
+    except requests.HTTPError as e:
+        log.error(f"Upstream HTTP error ({model_key}): {e}")
+        return jsonify({"error": str(e)}), e.response.status_code
+    except requests.Timeout:
+        return jsonify({"error": "upstream timeout"}), 504
     except Exception as e:
-        return jsonify({"error": str(e)}), 500
+        log.error(f"Proxy error ({model_key}): {e}")
+        return jsonify({"error": str(e)}), 502
 
-@app.route('/health', methods=['GET'])
+
+@app.route("/health", methods=["GET"])
 def health():
-    return jsonify({"status": "ok", "endpoint": "mira-proxy", "models": list(ENDPOINTS.keys())})
+    """
+    Proxy liveness + optional upstream health checks.
+    Pass ?upstream=true to also ping each Space's /health endpoint.
+    """
+    result: dict = {
+        "status":   "ok",
+        "service":  "mira-proxy",
+        "models":   list(ENDPOINTS.keys()),
+        "upstream": {},
+    }
+
+    if request.args.get("upstream", "false").lower() == "true":
+        for key, base in ENDPOINTS.items():
+            url = f"{base.rstrip('/')}/health"
+            try:
+                r = requests.get(
+                    url,
+                    headers={"Authorization": f"Bearer {HF_TOKEN}"},
+                    timeout=8,
+                )
+                result["upstream"][key] = {
+                    "status":      "ok" if r.ok else "error",
+                    "http_status": r.status_code,
+                    "body":        r.json() if r.ok else r.text[:200],
+                }
+            except requests.Timeout:
+                result["upstream"][key] = {"status": "timeout"}
+            except Exception as exc:
+                result["upstream"][key] = {"status": "error", "detail": str(exc)}
+
+    return jsonify(result)
+
 
-if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=7860)
+# ── Dev entrypoint ────────────────────────────────────────────────────────────
+if __name__ == "__main__":
+    # gunicorn is used in production (see Dockerfile CMD)
+    app.run(host="0.0.0.0", port=7860, debug=False)