hugging-science/Proxy-My-inference-api

1import os2import json3import time4import requests5from flask import Flask, request, jsonify, Response, stream_with_context6 7app = Flask(__name__)8 9ENDPOINT_URL = "https://o1zwshreete15x04.us-east-1.aws.endpoints.huggingface.cloud"10HF_TOKEN     = os.getenv("HF_TOKEN", "")11CHUNK_SIZE   = int(os.getenv("CHUNK_SIZE", "2"))    # chars per stream chunk12CHUNK_DELAY  = float(os.getenv("CHUNK_DELAY", "0.012"))  # seconds between chunks13 14 15def _headers():16    return {17        "Accept":        "application/json",18        "Authorization": f"Bearer {HF_TOKEN}",19        "Content-Type":  "application/json",20    }21 22 23@app.route("/health", methods=["GET"])24def health():25    """26    Lightweight health check — just pings the endpoint with a tiny prompt.27    """28    if not ENDPOINT_URL or not HF_TOKEN:29        return jsonify({30            "status":  "unhealthy",31            "error":   "ENDPOINT_URL or HF_TOKEN env var not set"32        }), 50333 34    try:35        resp = requests.post(36            ENDPOINT_URL,37            headers=_headers(),38            json={39                "inputs": [{"role": "user", "content": "Hi"}],40                "parameters": {41                    "max_new_tokens": 5,42                    "temperature":    0.1,43                    "do_sample":      False,44                }45            },46            timeout=20,47        )48        if resp.status_code == 200:49            return jsonify({"status": "healthy", "code": 200})50        return jsonify({51            "status": "unhealthy",52            "code":   resp.status_code,53            "error":  resp.text[:300],54        }), 50355 56    except requests.exceptions.Timeout:57        return jsonify({"status": "unhealthy", "error": "Endpoint timed out"}), 50358    except Exception as e:59        return jsonify({"status": "unhealthy", "error": str(e)}), 50360 61 62@app.route("/chat", methods=["POST"])63def chat():64    """65    Accepts:66    {67        "messages":    [{role, content}, ...],68        "max_tokens":  512,69        "temperature": 0.7,70        "top_p":       0.9,71        "do_sample":   true,72        "stream":      true      <-- if true, SSE stream back to caller73    }74    """75    data = request.get_json(silent=True) or {}76 77    messages    = data.get("messages", [])78    max_tokens  = int(data.get("max_tokens", 512))79    temperature = float(data.get("temperature", 0.7))80    top_p       = float(data.get("top_p", 0.9))81    do_sample   = bool(data.get("do_sample", temperature > 0))82    stream      = bool(data.get("stream", True))83 84    if not messages:85        return jsonify({"error": "messages array required"}), 40086 87    # Build payload for your custom handler88    payload = {89        "inputs": [{"role": m["role"], "content": m["content"]} for m in messages],90        "parameters": {91            "max_new_tokens": max_tokens,92            "temperature":    temperature,93            "top_p":          top_p,94            "do_sample":      do_sample,95        }96    }97 98    # ── Non-streaming ─────────────────────────────────────────99    if not stream:100        try:101            resp = requests.post(102                ENDPOINT_URL,103                headers=_headers(),104                json=payload,105                timeout=90,106            )107            resp.raise_for_status()108            result = resp.json()109            text = _extract_text(result)110            return jsonify({"generated_text": text, "ok": True})111        except Exception as e:112            return jsonify({"error": str(e)}), 500113 114    # ── Streaming ─────────────────────────────────────────────115    def generate():116        try:117            resp = requests.post(118                ENDPOINT_URL,119                headers=_headers(),120                json=payload,121                timeout=90,122            )123            resp.raise_for_status()124            result    = resp.json()125            full_text = _extract_text(result)126 127            if not full_text:128                yield _sse({"error": "Empty response from model"})129                return130 131            # ── Smooth streaming in small character chunks ────132            # We buffer into "word-aware" chunks so words don't133            # get split mid-character in a jarring way.134            buffer = ""135            for char in full_text:136                buffer += char137 138                # Flush on punctuation/spaces for natural rhythm139                should_flush = (140                    len(buffer) >= CHUNK_SIZE or141                    char in (' ', '\n', '.', ',', '!', '?', ':', ';')142                )143 144                if should_flush and buffer:145                    yield _sse({"token": buffer})146                    buffer = ""147                    time.sleep(CHUNK_DELAY)148 149            # Flush any remaining buffer150            if buffer:151                yield _sse({"token": buffer})152 153            yield "data: [DONE]\n\n"154 155        except requests.exceptions.Timeout:156            yield _sse({"error": "Endpoint timed out — try again"})157        except requests.exceptions.HTTPError as e:158            yield _sse({"error": f"Endpoint error {e.response.status_code}: {e.response.text[:200]}"})159        except Exception as e:160            yield _sse({"error": str(e)})161 162    return Response(163        stream_with_context(generate()),164        content_type="text/event-stream",165        headers={166            "Cache-Control":     "no-cache",167            "X-Accel-Buffering": "no",168            "Connection":        "keep-alive",169            "Access-Control-Allow-Origin": "*",170        }171    )172 173 174@app.route("/", methods=["GET"])175def root():176    return jsonify({177        "name":    "SmilyAI Proxy",178        "status":  "running",179        "routes":  ["/health", "/chat"],180        "model":   "SmilyAI 1.2B ChatML",181    })182 183 184def _extract_text(result):185    """Pull generated_text out of whatever shape the endpoint returns."""186    if isinstance(result, list) and len(result) > 0:187        return result[0].get("generated_text", "")188    if isinstance(result, dict):189        return result.get("generated_text", "")190    return str(result)191 192 193def _sse(obj):194    """Format a dict as an SSE data line."""195    return f"data: {json.dumps(obj)}\n\n"196 197 198if __name__ == "__main__":199    app.run(host="0.0.0.0", port=7860, debug=False)