Update main.py

main.py

@@ -1,271 +1,292 @@
 """
-Multi-model AI gateway endpoint — HF Spaces compatible.
-Authorization via 'connect' API key header.
-Streaming always enabled. Function calling supported.
-Uses httpx async — no openai SDK network issues.
+OpenAI-compatible /v1 API Gateway
+Proxies to NVIDIA NIM API with streaming always enabled,
+function calling support, and per-model system prompts.
+
+Deploy on Hugging Face Spaces (Docker).
+Authorization: Bearer connect
 """
 
 import json
-from typing import AsyncGenerator, Optional
+import time
+import uuid
+import asyncio
+from typing import Any, AsyncGenerator
 
 import httpx
-from fastapi import FastAPI, HTTPException, Header, Request
-from fastapi.responses import StreamingResponse
+from fastapi import FastAPI, HTTPException, Request
 from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel, Field
 
-from system_prompts import get_system_prompt
-
-# ── Config ──────────────────────────────────────────────────────────────────
+from system_prompts import SYSTEM_PROMPTS, MODEL_MAP, REVERSE_MODEL_MAP, EXTRA_BODY_MODELS
 
-CONNECT_KEY = "connect"
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
 
 NVIDIA_BASE_URL = "https://integrate.api.nvidia.com/v1"
 NVIDIA_API_KEY  = "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"
-NVIDIA_CHAT_URL = f"{NVIDIA_BASE_URL}/chat/completions"
-
-# Model registry: display-name → real model id + optional extra body
-MODELS = {
-    "Bielik-11b": {
-        "model_id":   "speakleash/bielik-11b-v2.6-instruct",
-        "extra_body": {
-            "chat_template_kwargs": {
-                "enable_thinking": False,
-                "clear_thinking":  True,
-            }
-        },
-    },
-    "GLM-4.7": {
-        "model_id":   "z-ai/glm4.7",
-        "extra_body": {
-            "chat_template_kwargs": {
-                "enable_thinking": False,
-                "clear_thinking":  True,
-            }
-        },
-    },
-    "Mistral-Small-4": {
-        "model_id":   "mistralai/mistral-small-4-119b-2603",
-        "extra_body": {},
-    },
-    "DeepSeek-V3.1": {
-        "model_id":   "deepseek-ai/deepseek-v3.1",
-        "extra_body": {},
-    },
-    "Kimi-K2": {
-        "model_id":   "moonshotai/kimi-k2-instruct",
-        "extra_body": {},
-    },
-}
-
-# ── FastAPI ──────────────────────────────────────────────────────────────────
+GATEWAY_API_KEY = "connect"
+
+# ---------------------------------------------------------------------------
+# App
+# ---------------------------------------------------------------------------
 
 app = FastAPI(
-    title="Multi-Model AI Gateway",
+    title="AI Gateway",
+    description="OpenAI-compatible gateway to NVIDIA NIM models",
     version="1.0.0",
-    description="Streaming endpoint for Bielik-11b, GLM-4.7, Mistral-Small-4, DeepSeek-V3.1, Kimi-K2",
 )
 
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
+    allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
-# ── Auth ─────────────────────────────────────────────────────────────────────
+# ---------------------------------------------------------------------------
+# Auth
+# ---------------------------------------------------------------------------
 
-def verify_key(authorization: Optional[str]) -> None:
-    if not authorization:
-        raise HTTPException(status_code=401, detail="Missing Authorization header")
-    scheme, _, token = authorization.partition(" ")
-    if scheme.lower() != "bearer" or token != CONNECT_KEY:
-        raise HTTPException(status_code=403, detail="Invalid API key")
+def verify_api_key(request: Request) -> None:
+    auth = request.headers.get("Authorization", "")
+    if not auth.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Missing Bearer token")
+    token = auth.removeprefix("Bearer ").strip()
+    if token != GATEWAY_API_KEY:
+        raise HTTPException(status_code=401, detail="Invalid API key")
 
-# ── Schemas ───────────────────────────────────────────────────────────────────
+# ---------------------------------------------------------------------------
+# Pydantic models (OpenAI-compatible)
+# ---------------------------------------------------------------------------
 
-class Message(BaseModel):
-    role: str
-    content: str | list
+class FunctionParameters(BaseModel):
+    type: str = "object"
+    properties: dict[str, Any] = {}
+    required: list[str] = []
 
-class ToolFunction(BaseModel):
+class FunctionDef(BaseModel):
     name: str
-    description: Optional[str] = None
-    parameters: Optional[dict] = None
+    description: str | None = None
+    parameters: FunctionParameters | None = None
 
 class Tool(BaseModel):
     type: str = "function"
-    function: ToolFunction
+    function: FunctionDef
+
+class ToolChoice(BaseModel):
+    type: str = "function"
+    function: dict[str, str] | None = None
+
+class Message(BaseModel):
+    role: str
+    content: str | list[Any] | None = None
+    name: str | None = None
+    tool_calls: list[Any] | None = None
+    tool_call_id: str | None = None
 
-class ChatRequest(BaseModel):
-    model: str = Field(..., description="Bielik-11b | GLM-4.7 | Mistral-Small-4 | DeepSeek-V3.1 | Kimi-K2")
+class ChatCompletionRequest(BaseModel):
+    model: str
     messages: list[Message]
-    tools: Optional[list[Tool]] = None
-    tool_choice: Optional[str | dict] = None
-    temperature: Optional[float] = None
-    max_tokens: Optional[int] = None
-    top_p: Optional[float] = None
-    presence_penalty: Optional[float] = None
-    frequency_penalty: Optional[float] = None
-    inject_system_prompt: bool = Field(default=True)
-
-# ── Core stream helper ────────────────────────────────────────────────────────
-
-def _build_payload(model_name: str, messages: list[dict], tools, tool_choice, kwargs: dict) -> dict:
-    cfg = MODELS[model_name]
-    payload: dict = {
-        "model":    cfg["model_id"],
+    temperature: float | None = None
+    top_p: float | None = None
+    max_tokens: int | None = None
+    tools: list[Tool] | None = None
+    tool_choice: str | ToolChoice | None = None
+    # stream is ALWAYS True – ignored if provided, always forced to True
+    stream: bool = True
+    stop: list[str] | str | None = None
+    presence_penalty: float | None = None
+    frequency_penalty: float | None = None
+    seed: int | None = None
+    n: int | None = None
+    logprobs: bool | None = None
+    top_logprobs: int | None = None
+    user: str | None = None
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def resolve_model(requested: str) -> str:
+    """Map display name or raw NVIDIA model ID to NVIDIA model ID."""
+    if requested in MODEL_MAP:
+        return MODEL_MAP[requested]
+    if requested in REVERSE_MODEL_MAP:
+        return requested  # already a raw ID
+    raise HTTPException(
+        status_code=400,
+        detail=f"Unknown model '{requested}'. Available: {list(MODEL_MAP.keys())}",
+    )
+
+def get_display_name(nvidia_id: str) -> str:
+    return REVERSE_MODEL_MAP.get(nvidia_id, nvidia_id)
+
+def inject_system_prompt(messages: list[Message], display_name: str) -> list[dict]:
+    """Inject per-model system prompt if not already present."""
+    prompt = SYSTEM_PROMPTS.get(display_name)
+    serialized = [m.model_dump(exclude_none=True) for m in messages]
+
+    if prompt:
+        has_system = any(m["role"] == "system" for m in serialized)
+        if not has_system:
+            serialized = [{"role": "system", "content": prompt}] + serialized
+
+    return serialized
+
+def build_nvidia_payload(req: ChatCompletionRequest, nvidia_model: str) -> dict:
+    display = get_display_name(nvidia_model)
+    messages = inject_system_prompt(req.messages, display)
+
+    payload: dict[str, Any] = {
+        "model": nvidia_model,
         "messages": messages,
-        "stream":   True,
-        **kwargs,
+        "stream": True,  # ALWAYS TRUE
     }
-    if tools:
-        payload["tools"] = tools
-    if tool_choice is not None:
-        payload["tool_choice"] = tool_choice
-    # merge extra_body fields at top level (NVIDIA NIM style)
-    if cfg["extra_body"]:
-        payload.update(cfg["extra_body"])
+
+    # Optional params
+    if req.temperature is not None:
+        payload["temperature"] = req.temperature
+    if req.top_p is not None:
+        payload["top_p"] = req.top_p
+    if req.max_tokens is not None:
+        payload["max_tokens"] = req.max_tokens
+    if req.stop is not None:
+        payload["stop"] = req.stop
+    if req.presence_penalty is not None:
+        payload["presence_penalty"] = req.presence_penalty
+    if req.frequency_penalty is not None:
+        payload["frequency_penalty"] = req.frequency_penalty
+    if req.seed is not None:
+        payload["seed"] = req.seed
+    if req.n is not None:
+        payload["n"] = req.n
+    if req.user is not None:
+        payload["user"] = req.user
+
+    # Function calling / tools
+    if req.tools:
+        payload["tools"] = [t.model_dump(exclude_none=True) for t in req.tools]
+    if req.tool_choice is not None:
+        if isinstance(req.tool_choice, str):
+            payload["tool_choice"] = req.tool_choice
+        else:
+            payload["tool_choice"] = req.tool_choice.model_dump(exclude_none=True)
+
+    # Extra body for specific models (e.g. GLM-4.7 thinking params)
+    extra = EXTRA_BODY_MODELS.get(nvidia_model, {})
+    payload.update(extra)
+
     return payload
 
+# ---------------------------------------------------------------------------
+# SSE streaming proxy
+# ---------------------------------------------------------------------------
 
-async def stream_nvidia(
-    model_name: str,
-    messages: list[dict],
-    tools,
-    tool_choice,
-    kwargs: dict,
-) -> AsyncGenerator[bytes, None]:
-    payload = _build_payload(model_name, messages, tools, tool_choice, kwargs)
+async def stream_nvidia(payload: dict) -> AsyncGenerator[bytes, None]:
     headers = {
         "Authorization": f"Bearer {NVIDIA_API_KEY}",
-        "Content-Type":  "application/json",
-        "Accept":        "text/event-stream",
+        "Content-Type": "application/json",
+        "Accept": "text/event-stream",
     }
 
-    async with httpx.AsyncClient(timeout=httpx.Timeout(120.0, connect=10.0)) as client:
+    async with httpx.AsyncClient(timeout=300) as client:
         async with client.stream(
             "POST",
-            NVIDIA_CHAT_URL,
+            f"{NVIDIA_BASE_URL}/chat/completions",
             headers=headers,
             json=payload,
         ) as response:
             if response.status_code != 200:
                 body = await response.aread()
-                error_msg = body.decode(errors="replace")
-                yield f"data: {json.dumps({'error': error_msg, 'status': response.status_code})}\n\n".encode()
+                error_detail = body.decode(errors="replace")
+                error_chunk = {
+                    "error": {
+                        "message": f"Upstream error {response.status_code}: {error_detail}",
+                        "type": "upstream_error",
+                        "code": response.status_code,
+                    }
+                }
+                yield f"data: {json.dumps(error_chunk)}\n\n".encode()
+                yield b"data: [DONE]\n\n"
                 return
 
             async for line in response.aiter_lines():
-                if line:
+                if line.startswith("data: "):
                     yield f"{line}\n\n".encode()
+                    if line == "data: [DONE]":
+                        return
+                elif line.strip():
+                    # Pass through any unexpected lines
+                    yield f"data: {line}\n\n".encode()
 
-# ── Shared logic ──────────────────────────────────────────────────────────────
-
-def prepare_messages(model_name: str, raw_messages: list[dict], inject: bool) -> list[dict]:
-    messages = list(raw_messages)
-    if inject:
-        system_prompt = get_system_prompt(model_name)
-        if not messages or messages[0].get("role") != "system":
-            messages.insert(0, {"role": "system", "content": system_prompt})
-    return messages
-
-
-def extract_kwargs(source, fields: tuple) -> dict:
-    kwargs = {}
-    for field in fields:
-        if isinstance(source, dict):
-            val = source.get(field)
-        else:
-            val = getattr(source, field, None)
-        if val is not None:
-            kwargs[field] = val
-    return kwargs
-
-
-OPTIONAL_FIELDS = ("temperature", "max_tokens", "top_p", "presence_penalty", "frequency_penalty")
-
-SSE_HEADERS = {
-    "Cache-Control":     "no-cache",
-    "X-Accel-Buffering": "no",
-}
-
-# ── Endpoints ─────────────────────────────────────────────────────────────────
+# ---------------------------------------------------------------------------
+# Routes
+# ---------------------------------------------------------------------------
 
 @app.get("/")
 async def root():
-    return {
-        "service": "Multi-Model AI Gateway",
-        "models":  list(MODELS.keys()),
-        "auth":    "Bearer connect",
-        "docs":    "/docs",
-    }
-
-
-@app.get("/models")
-async def list_models(authorization: Optional[str] = Header(default=None)):
-    verify_key(authorization)
-    return {
-        name: {"model_id": cfg["model_id"]}
-        for name, cfg in MODELS.items()
-    }
+    return {"status": "ok", "service": "AI Gateway", "version": "1.0.0"}
+
+@app.get("/v1/models")
+async def list_models(request: Request):
+    verify_api_key(request)
+    now = int(time.time())
+    models = []
+    for display_name in MODEL_MAP:
+        models.append({
+            "id": display_name,
+            "object": "model",
+            "created": now,
+            "owned_by": "ai-gateway",
+        })
+    return {"object": "list", "data": models}
 
+@app.post("/v1/chat/completions")
+async def chat_completions(request: Request, req: ChatCompletionRequest):
+    verify_api_key(request)
 
-@app.post("/chat")
-async def chat(
-    request: ChatRequest,
-    authorization: Optional[str] = Header(default=None),
-):
-    verify_key(authorization)
-
-    if request.model not in MODELS:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Unknown model '{request.model}'. Available: {list(MODELS.keys())}",
-        )
-
-    messages = prepare_messages(
-        request.model,
-        [m.model_dump() for m in request.messages],
-        request.inject_system_prompt,
-    )
-    kwargs = extract_kwargs(request, OPTIONAL_FIELDS)
-    tools  = [t.model_dump() for t in request.tools] if request.tools else None
+    nvidia_model = resolve_model(req.model)
+    payload = build_nvidia_payload(req, nvidia_model)
 
     return StreamingResponse(
-        stream_nvidia(request.model, messages, tools, request.tool_choice, kwargs),
+        stream_nvidia(payload),
         media_type="text/event-stream",
-        headers=SSE_HEADERS,
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no",
+        },
     )
 
+# Passthrough completions (legacy)
+@app.post("/v1/completions")
+async def completions(request: Request):
+    verify_api_key(request)
+    body = await request.json()
 
-@app.post("/v1/chat/completions")
-async def openai_compat(
-    raw: Request,
-    authorization: Optional[str] = Header(default=None),
-):
-    """OpenAI-compatible drop-in. Use gateway model names as 'model'."""
-    verify_key(authorization)
-    body = await raw.json()
-
-    model_name = body.get("model", "")
-    if model_name not in MODELS:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Unknown model '{model_name}'. Available: {list(MODELS.keys())}",
-        )
-
-    messages = prepare_messages(
-        model_name,
-        body.get("messages", []),
-        body.get("inject_system_prompt", True),
-    )
-    kwargs      = extract_kwargs(body, OPTIONAL_FIELDS)
-    tools       = body.get("tools")
-    tool_choice = body.get("tool_choice")
+    model_req = body.get("model", "")
+    try:
+        nvidia_model = resolve_model(model_req)
+    except HTTPException:
+        nvidia_model = model_req
+
+    body["model"] = nvidia_model
+    body["stream"] = True
 
     return StreamingResponse(
-        stream_nvidia(model_name, messages, tools, tool_choice, kwargs),
+        stream_nvidia(body),
         media_type="text/event-stream",
-        headers=SSE_HEADERS,
-    )
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no",
+        },
+    )
+
+@app.get("/health")
+async def health():
+    return {"status": "healthy"}