Create app.py

app.py

@@ -0,0 +1,476 @@
+import os
+import json
+import time
+import httpx
+import uvicorn
+import gradio as gr
+from fastapi import FastAPI, Request, Header, HTTPException
+from fastapi.responses import JSONResponse, StreamingResponse
+
+# ---------------------------------------------------------------------
+# Configuration
+# ---------------------------------------------------------------------
+BYTEZ_CHAT_URL = "https://api.bytez.com/models/v2/openai/v1/chat/completions"
+BYTEZ_MODELS_URL = "https://api.bytez.com/models/v2/list/models"
+BYTEZ_IMAGE_URL = "https://api.bytez.com/models/v2/openai/dall-e-3"
+BYTEZ_AUTH = os.getenv("BYTEZ_API_KEY")          # your Bytez key
+LOCAL_API_KEY = os.getenv("LOCAL_API_KEY")      # optional local guard
+BYTEZ_AUTH_2 = os.getenv("BYTEZ_API_KEY_2")
+# ---------------------------------------------------------------------
+# FastAPI app
+# ---------------------------------------------------------------------
+api = FastAPI(title="Bytez → OpenAI Proxy (v1 + v2)")
+
+def check_key(auth: str | None):
+    """Validate the Bearer token (optional local key)."""
+    if not auth or not auth.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Missing or invalid API key")
+    user_key = auth.split("Bearer ")[1].strip()
+    if LOCAL_API_KEY and user_key != LOCAL_API_KEY:
+        raise HTTPException(status_code=403, detail="Unauthorized API key")
+
+# ---------------------------------------------------------------------
+# Root / health
+# ---------------------------------------------------------------------
+@api.get("/")
+def root():
+    return {"status": "ok", "message": "Bytez proxy (v1+v2) running"}
+
+# ---------------------------------------------------------------------
+# --------------------------  /v1  ------------------------------------
+# ---------------------------------------------------------------------
+
+@api.get("/v1/models")
+async def v1_models(authorization: str = Header(None)):
+    check_key(authorization)
+    if not BYTEZ_AUTH:
+        raise HTTPException(status_code=500, detail="Server BYTEZ_API_KEY not configured")
+
+    async with httpx.AsyncClient(timeout=30) as c:
+        r = await c.get(BYTEZ_MODELS_URL, headers={"Authorization": BYTEZ_AUTH})
+
+    try:
+        data = r.json()
+    except json.JSONDecodeError:
+        raise HTTPException(status_code=502, detail="Upstream returned invalid JSON")
+
+    # Transform Bytez → OpenAI list
+    models_list = [
+        {"id": m.get("id") or m.get("name"), "object": "model"}
+        for m in (data if isinstance(data, list) else data.get("data", []))
+    ]
+    return JSONResponse(
+        {"object": "list", "data": models_list},
+        headers={"Access-Control-Allow-Origin": "*"}
+    )
+
+
+@api.post("/v1/chat/completions")
+async def v1_chat(request: Request, authorization: str = Header(None)):
+    """Exactly the same implementation you already had – untouched."""
+    check_key(authorization)
+    if not BYTEZ_AUTH:
+        raise HTTPException(status_code=500, detail="Server BYTEZ_API_KEY not configured")
+
+    payload = await request.json()
+    stream = payload.get("stream", False)
+    headers = {"Authorization": BYTEZ_AUTH, "Content-Type": "application/json"}
+
+    # ---------- streaming helper ----------
+    async def v1_event_stream():
+        async with httpx.AsyncClient(timeout=120) as client:
+            async with client.stream("POST", BYTEZ_CHAT_URL, headers=headers, json=payload) as upstream:
+                async for line in upstream.aiter_lines():
+                    line = line.strip()
+                    if not line:
+                        continue
+                    json_str = line[6:] if line.startswith("data: ") else line
+                    try:
+                        chunk = json.loads(json_str)
+                    except json.JSONDecodeError:
+                        continue
+
+                    if json_str == "[DONE]":
+                        yield "data: [DONE]\n\n"
+                        break
+
+                    # ----- adapt Bytez chunk to OpenAI -----
+                    content = ""
+                    if "token" in chunk:
+                        content = chunk["token"]
+                    elif "choices" in chunk and chunk["choices"]:
+                        delta = chunk["choices"][0].get("delta", {})
+                        content = delta.get("content", "")
+                    elif "text" in chunk:
+                        content = chunk["text"]
+                    else:
+                        content = str(chunk)
+
+                    openai_chunk = {
+                        "id": "chatcmpl-proxy-stream",
+                        "object": "chat.completion.chunk",
+                        "created": int(time.time()),
+                        "model": payload.get("model", "unknown"),
+                        "choices": [
+                            {
+                                "index": 0,
+                                "delta": {"role": "assistant", "content": content},
+                                "finish_reason": None,
+                            }
+                        ],
+                    }
+                    yield f"data: {json.dumps(openai_chunk)}\n\n"
+        yield "data: [DONE]\n\n"
+
+    # ---------- non-stream ----------
+    if not stream:
+        async with httpx.AsyncClient(timeout=120) as c:
+            r = await c.post(BYTEZ_CHAT_URL, headers=headers, json=payload)
+        try:
+            data = r.json()
+        except json.JSONDecodeError:
+            raise HTTPException(status_code=502, detail="Upstream returned invalid JSON")
+
+        if "choices" not in data:
+            content = data.get("output") or data.get("response") or data.get("message") or str(data)
+            data = {
+                "id": "chatcmpl-proxy",
+                "object": "chat.completion",
+                "choices": [{"index": 0, "message": {"role": "assistant", "content": content}}],
+            }
+        return JSONResponse(data, headers={"Access-Control-Allow-Origin": "*"})
+
+    return StreamingResponse(
+        v1_event_stream(),
+        media_type="text/event-stream",
+        headers={"Access-Control-Allow-Origin": "*"},
+    )
+
+# ---------------------------------------------------------------------
+# --------------------- /v1/images/generations (FIXED) ---------------
+# ---------------------------------------------------------------------
+# ---------------------------------------------------------------------
+# --------------------- /v1/images/generations  -----------------------
+# ---------------------------------------------------------------------
+@api.post("/v1/images/generations")
+async def v1_images_generations(request: Request, authorization: str = Header(None)):
+    """
+    OpenAI-compatible image generation endpoint that forwards to:
+      https://api.bytez.com/models/v2/openai/dall-e-3
+
+    It accepts the usual OpenAI body:
+      {
+        "model": "dall-e-3",           # REQUIRED by clients, but ignored
+        "prompt": "text",              # REQUIRED
+        "n": 1,
+        "size": "1024x1024",
+        "response_format": "url" | "b64_json",
+        "quality": "standard" | "hd",
+        "style": "vivid" | "natural"
+      }
+    and converts it to the Bytez format from your curl:
+      { "text": "<prompt>" }
+    """
+    check_key(authorization)
+
+    # Use SAME key as chat unless you *really* have a separate one
+    if not BYTEZ_AUTH:
+        raise HTTPException(status_code=500, detail="Server BYTEZ_API_KEY not configured")
+
+    # ---------------- parse request ----------------
+    try:
+        payload = await request.json()
+    except json.JSONDecodeError:
+        raise HTTPException(status_code=400, detail="Invalid JSON")
+
+    prompt = payload.get("prompt") or payload.get("text")
+    if not prompt or not str(prompt).strip():
+        raise HTTPException(status_code=400, detail="Field 'prompt' is required")
+
+    # These are there only to satisfy OpenAI-compatible clients
+    _model = payload.get("model", "dall-e-3")  # ignored
+    n = int(payload.get("n", 1))
+    response_format = payload.get("response_format", "url")  # "url" | "b64_json"
+    size = payload.get("size", "1024x1024")
+    quality = payload.get("quality")
+    style = payload.get("style")
+
+    # ---------------- build Bytez request ----------------
+    # Minimal payload that we KNOW works from your curl:
+    bytez_payload = {
+        "text": prompt,
+    }
+
+    # Only add extras if Bytez actually supports them.
+    # Comment these out if you’re unsure; or keep them if Bytez docs say so.
+    # bytez_payload["num_outputs"] = n
+    # bytez_payload["size"] = size
+    # if quality in ["standard", "hd"]:
+    #     bytez_payload["quality"] = quality
+    # if style in ["vivid", "natural"]:
+    #     bytez_payload["style"] = style
+
+    headers = {
+        "Authorization": BYTEZ_AUTH,         # << important
+        "Content-Type": "application/json",
+    }
+
+    # ---------------- call Bytez ----------------
+    async with httpx.AsyncClient(timeout=200) as client:
+        try:
+            resp = await client.post(
+                BYTEZ_IMAGE_URL,
+                json=bytez_payload,
+                headers=headers,
+            )
+            resp.raise_for_status()
+        except httpx.HTTPStatusError as e:
+            # Surface Bytez error message directly to client
+            detail = None
+            try:
+                detail = e.response.json()
+            except Exception:
+                detail = e.response.text
+            raise HTTPException(
+                status_code=e.response.status_code,
+                detail={"upstream_error": detail},
+            )
+        except Exception as e:
+            raise HTTPException(status_code=502, detail=f"Bytez unreachable: {str(e)}")
+
+    try:
+        bytez_data = resp.json()
+    except json.JSONDecodeError:
+        raise HTTPException(status_code=502, detail="Bytez returned invalid JSON")
+
+    # ---------------- map Bytez → OpenAI image response ----------------
+    # We don't know the exact shape, so we handle several possibilities.
+
+    images = []
+
+    # Case 1: { "images": ["url_or_b64", ...] }
+    if isinstance(bytez_data, dict) and "images" in bytez_data and isinstance(bytez_data["images"], list):
+        images = bytez_data["images"]
+
+    # Case 2: { "data": [ { "url": "..." } , ... ] }
+    elif isinstance(bytez_data, dict) and "data" in bytez_data and isinstance(bytez_data["data"], list):
+        for item in bytez_data["data"]:
+            if "url" in item:
+                images.append(item["url"])
+            elif "b64_json" in item:
+                images.append(item["b64_json"])
+            else:
+                images.append(str(item))
+
+    # Case 3: single string
+    elif isinstance(bytez_data, str):
+        images = [bytez_data]
+
+    # Fallback: treat whole thing as one string
+    else:
+        images = [str(bytez_data)]
+
+    if not images:
+        raise HTTPException(status_code=500, detail={"error": "No images returned from Bytez", "raw": bytez_data})
+
+    # truncate to n images
+    images = images[:n]
+
+    openai_data = []
+    for img in images:
+        # If it looks like a data URL already
+        if isinstance(img, str) and img.startswith("data:image"):
+            # Already a data URL; extract base64 if needed
+            b64_part = img.split("base64,", 1)[-1]
+            if response_format == "b64_json":
+                openai_data.append({"b64_json": b64_part})
+            else:
+                openai_data.append({"url": img})
+
+        else:
+            # Raw URL or base64 – we can't be sure, so:
+            if response_format == "b64_json":
+                # Assume it's base64 or treat as a string anyway
+                openai_data.append({"b64_json": str(img)})
+            else:
+                # Assume it's a URL or make it one if needed
+                openai_data.append({"url": str(img)})
+
+    result = {
+        "created": int(time.time()),
+        "data": openai_data,
+    }
+    return JSONResponse(result, headers={"Access-Control-Allow-Origin": "*"})
+
+
+
+# ---------------------------------------------------------------------
+# --------------------------  /v2  ------------------------------------
+# ---------------------------------------------------------------------
+
+@api.post("/v2/chat/completions")
+async def v2_chat_completions(request: Request, authorization: str = Header(None)):
+    """
+    v2 – clean OpenAI-compatible streaming.
+    * First chunk includes role=assistant (required by Continue.dev)
+    * Later chunks send only delta.content
+    * No usage events
+    """
+
+    check_key(authorization)
+
+    if not BYTEZ_AUTH_2:
+        raise HTTPException(status_code=500, detail="Server BYTEZ_API_2 not configured")
+
+    try:
+        body = await request.body()
+        payload = json.loads(body.decode("utf-8"))
+    except json.JSONDecodeError as e:
+        raise HTTPException(status_code=400, detail=f"Invalid JSON: {e}")
+
+    stream = payload.get("stream", False)
+
+    upstream_headers = {
+        "Authorization": BYTEZ_AUTH_2,
+        "Content-Type": "application/json",
+    }
+
+    # Normal content chunk (NO ROLE)
+    def make_openai_delta(content: str):
+        return {
+            "id": f"chatcmpl-v2-{int(time.time())}",
+            "object": "chat.completion.chunk",
+            "created": int(time.time()),
+            "model": payload.get("model", "unknown"),
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {"content": content},
+                    "finish_reason": None,
+                }
+            ],
+        }
+
+    async def clean_stream():
+        # FIRST CHUNK MUST SET THE ROLE → REQUIRED by Continue.dev
+        first_chunk = {
+            "id": f"chatcmpl-v2-{int(time.time())}",
+            "object": "chat.completion.chunk",
+            "created": int(time.time()),
+            "model": payload.get("model", "unknown"),
+            "choices": [
+                {
+                    "index": 0,
+                    "delta": {"role": "assistant", "content": ""},
+                    "finish_reason": None,
+                }
+            ],
+        }
+
+        # Send first role-setting chunk
+        yield f"data: {json.dumps(first_chunk)}\n\n"
+
+        async with httpx.AsyncClient(timeout=180) as client:
+            try:
+                async with client.stream(
+                    "POST", BYTEZ_CHAT_URL, headers=upstream_headers, json=payload
+                ) as upstream:
+
+                    async for line in upstream.aiter_lines():
+                        line = line.strip()
+                        if not line:
+                            continue
+
+                        json_str = line[6:] if line.startswith("data: ") else line
+
+                        # Skip usage events
+                        if "usage" in json_str.lower():
+                            continue
+
+                        if json_str == "[DONE]":
+                            yield "data: [DONE]\n\n"
+                            return
+
+                        try:
+                            chunk = json.loads(json_str)
+                        except json.JSONDecodeError:
+                            continue
+
+                        text = ""
+                        if isinstance(chunk, dict):
+                            if "token" in chunk:
+                                text = chunk["token"]
+                            elif "choices" in chunk and chunk["choices"]:
+                                delta = chunk["choices"][0].get("delta", {})
+                                text = delta.get("content", "")
+                            elif "text" in chunk:
+                                text = chunk["text"]
+                            else:
+                                text = str(chunk)
+
+                        if text:
+                            yield f"data: {json.dumps(make_openai_delta(text))}\n\n"
+
+                yield "data: [DONE]\n\n"
+
+            except Exception as e:
+                error_chunk = make_openai_delta(f"Error: {str(e)}")
+                yield f"data: {json.dumps(error_chunk)}\n\n"
+                yield "data: [DONE]\n\n"
+
+    # Non-streaming mode
+    if not stream:
+        async with httpx.AsyncClient(timeout=120) as c:
+            r = await c.post(BYTEZ_CHAT_URL, headers=upstream_headers, json=payload)
+            r.raise_for_status()
+            data = r.json()
+
+        if "choices" not in data:
+            content = (
+                data.get("output")
+                or data.get("response")
+                or data.get("message")
+                or str(data)
+            )
+            data = {
+                "id": "chatcmpl-v2",
+                "object": "chat.completion",
+                "choices": [
+                    {"index": 0, "message": {"role": "assistant", "content": content}}
+                ],
+            }
+
+        return JSONResponse(data, headers={"Access-Control-Allow-Origin": "*"})
+
+    # Streaming mode
+    return StreamingResponse(
+        clean_stream(),
+        media_type="text/event-stream",
+        headers={
+            "Access-Control-Allow-Origin": "*",
+            "Access-Control-Allow-Headers": "*",
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no",
+        },
+    )
+
+# ---------------------------------------------------------------------
+# Minimal Gradio UI (required for HF Space to start)
+# ---------------------------------------------------------------------
+with gr.Blocks() as ui:
+    gr.Markdown(
+        "### Bytez → OpenAI Proxy (v1 + **v2**)\n"
+        "- `/v1/models`  \n"
+        "- `/v1/chat/completions` (unchanged)  \n"
+        "- **`/v2/chat/completions`** – clean streaming, no usage chunk"
+    )
+
+demo = gr.mount_gradio_app(api, ui, path="/")
+
+# This makes it work on Render, Railway, Fly.io, etc.
+app = api
+
+if __name__ == "__main__":
+    # Only for local testing with Gradio
+    uvicorn.run(demo, host="0.0.0.0", port=8000)