Create main.py

main.py

@@ -0,0 +1,128 @@
+import os
+import requests
+from typing import List, Optional, Dict, Any
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+import gradio as gr
+
+# ===== Config =====
+HF_TOKEN = os.getenv("HF_TOKEN", "")
+ROUTER_URL = "https://router.huggingface.co/v1/chat/completions"
+
+DEFAULT_MODEL = os.getenv("MODEL_NAME", "Qwen/Qwen2.5-1.5B-Instruct")
+
+# ===== FastAPI (OpenAI-ish) =====
+api = FastAPI(title="My Hosted-Model Proxy (API + UI)")
+
+class ChatMessage(BaseModel):
+    role: str
+    content: str
+
+class ChatReq(BaseModel):
+    model: str = Field(default=DEFAULT_MODEL)
+    messages: List[ChatMessage]
+    max_tokens: int = Field(default=256, ge=1, le=2048)
+    temperature: float = Field(default=0.7, ge=0.0, le=2.0)
+    stream: bool = False
+
+@api.get("/health")
+def health():
+    return {"status": "ok", "model_default": DEFAULT_MODEL, "proxying": True}
+
+@api.get("/v1/models")
+def models():
+    return {"object": "list", "data": [{"id": DEFAULT_MODEL, "object": "model", "owned_by": "hf-router"}]}
+
+@api.post("/v1/chat/completions")
+def chat(req: ChatReq):
+    if not HF_TOKEN:
+        raise HTTPException(500, "HF_TOKEN missing. Add it in Space Settings → Secrets as HF_TOKEN.")
+
+    payload = req.model_dump()
+    # Convert Pydantic messages to plain dicts
+    payload["messages"] = [m.model_dump() for m in req.messages]
+
+    r = requests.post(
+        ROUTER_URL,
+        headers={
+            "Authorization": f"Bearer {HF_TOKEN}",
+            "Content-Type": "application/json",
+        },
+        json=payload,
+        timeout=120,
+    )
+
+    # Router errors are still JSON; keep them readable
+    try:
+        out = r.json()
+    except Exception:
+        raise HTTPException(502, f"Router returned non-JSON: {r.text[:500]}")
+
+    if r.status_code >= 400:
+        raise HTTPException(r.status_code, out)
+
+    return out
+
+# ===== Gradio UI =====
+def ui_chat(user_message: str, history: list, model: str, temperature: float, max_tokens: int):
+    """
+    history: list of [user, assistant] pairs from Gradio
+    """
+    messages: List[Dict[str, str]] = []
+    # Optional system prompt to make it feel natural
+    messages.append({"role": "system", "content": "Be natural, concise, and helpful."})
+
+    for u, a in history:
+        messages.append({"role": "user", "content": u})
+        messages.append({"role": "assistant", "content": a})
+
+    messages.append({"role": "user", "content": user_message})
+
+    req = {
+        "model": model,
+        "messages": messages,
+        "temperature": float(temperature),
+        "max_tokens": int(max_tokens),
+        "stream": False,
+    }
+
+    if not HF_TOKEN:
+        return "HF_TOKEN is missing. Add it in Space Settings → Secrets as HF_TOKEN."
+
+    r = requests.post(
+        ROUTER_URL,
+        headers={"Authorization": f"Bearer {HF_TOKEN}", "Content-Type": "application/json"},
+        json=req,
+        timeout=120,
+    )
+    data = r.json()
+    if r.status_code >= 400:
+        return f"Error: {data}"
+
+    return data["choices"][0]["message"]["content"]
+
+with gr.Blocks(title="My AI (Hosted Model) — UI + API") as demo:
+    gr.Markdown("## 💬 Chat UI (Hosted model via HF Router)\nThis Space also exposes an **OpenAI-style API** at `/v1/chat/completions`.")
+
+    with gr.Row():
+        model_dd = gr.Textbox(value=DEFAULT_MODEL, label="Model", interactive=True)
+    with gr.Row():
+        temp = gr.Slider(0, 2, value=0.7, step=0.1, label="Temperature")
+        mx = gr.Slider(32, 1024, value=256, step=16, label="Max tokens")
+
+    gr.ChatInterface(
+        fn=lambda msg, hist: ui_chat(msg, hist, model_dd.value, temp.value, mx.value),
+        title="Chat",
+        description="Type a message. This calls the hosted model and returns the assistant response.",
+    )
+
+    gr.Markdown(
+        "**API endpoints:**\n"
+        "- `GET /health`\n"
+        "- `GET /v1/models`\n"
+        "- `POST /v1/chat/completions`\n"
+    )
+
+# Mount Gradio onto FastAPI at /
+app = gr.mount_gradio_app(api, demo, path="/")