Update main.py

main.py

@@ -2,17 +2,17 @@
 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.
 """
 
 import json
-import asyncio
 from typing import AsyncGenerator, Optional
 
+import httpx
 from fastapi import FastAPI, HTTPException, Header, Request
 from fastapi.responses import StreamingResponse
 from fastapi.middleware.cors import CORSMiddleware
 from pydantic import BaseModel, Field
-from openai import OpenAI
 
 from system_prompts import get_system_prompt
 
@@ -22,6 +22,7 @@ CONNECT_KEY = "connect"
 
 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 = {
@@ -72,12 +73,9 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-client = OpenAI(base_url=NVIDIA_BASE_URL, api_key=NVIDIA_API_KEY)
-
 # ── Auth ─────────────────────────────────────────────────────────────────────
 
 def verify_key(authorization: Optional[str]) -> None:
-    """Check Bearer token matches CONNECT_KEY."""
     if not authorization:
         raise HTTPException(status_code=401, detail="Missing Authorization header")
     scheme, _, token = authorization.partition(" ")
@@ -88,7 +86,7 @@ def verify_key(authorization: Optional[str]) -> None:
 
 class Message(BaseModel):
     role: str
-    content: str | list  # supports text or multipart
+    content: str | list
 
 class ToolFunction(BaseModel):
     name: str
@@ -100,7 +98,7 @@ class Tool(BaseModel):
     function: ToolFunction
 
 class ChatRequest(BaseModel):
-    model: str = Field(..., description="Model name: Bielik-11b | GLM-4.7 | Mistral-Small-4 | DeepSeek-V3.1 | Kimi-K2")
+    model: str = Field(..., description="Bielik-11b | GLM-4.7 | Mistral-Small-4 | DeepSeek-V3.1 | Kimi-K2")
     messages: list[Message]
     tools: Optional[list[Tool]] = None
     tool_choice: Optional[str | dict] = None
@@ -109,49 +107,88 @@ class ChatRequest(BaseModel):
     top_p: Optional[float] = None
     presence_penalty: Optional[float] = None
     frequency_penalty: Optional[float] = None
-    inject_system_prompt: bool = Field(
-        default=True,
-        description="Prepend the model-specific system prompt automatically"
-    )
+    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"],
+        "messages": messages,
+        "stream":   True,
+        **kwargs,
+    }
+    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"])
+    return payload
 
-# ── Stream helper ─────────────────────────────────────────────────────────────
 
 async def stream_nvidia(
     model_name: str,
     messages: list[dict],
-    tools: Optional[list[dict]],
+    tools,
     tool_choice,
     kwargs: dict,
-    extra_body: dict,
-) -> AsyncGenerator[str, None]:
-    """Yield SSE chunks from NVIDIA NIM in a thread-safe way."""
-
-    params = {
-        "model":      MODELS[model_name]["model_id"],
-        "messages":   messages,
-        "stream":     True,          # always True
-        **kwargs,
+) -> AsyncGenerator[bytes, None]:
+    payload = _build_payload(model_name, messages, tools, tool_choice, kwargs)
+    headers = {
+        "Authorization": f"Bearer {NVIDIA_API_KEY}",
+        "Content-Type":  "application/json",
+        "Accept":        "text/event-stream",
     }
 
-    if tools:
-        params["tools"] = tools
-    if tool_choice is not None:
-        params["tool_choice"] = tool_choice
-    if extra_body:
-        params["extra_body"] = extra_body
+    async with httpx.AsyncClient(timeout=httpx.Timeout(120.0, connect=10.0)) as client:
+        async with client.stream(
+            "POST",
+            NVIDIA_CHAT_URL,
+            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()
+                return
+
+            async for line in response.aiter_lines():
+                if line:
+                    yield f"{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
 
-    loop = asyncio.get_event_loop()
 
-    def _call():
-        return client.chat.completions.create(**params)
+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
 
-    stream = await loop.run_in_executor(None, _call)
 
-    for chunk in stream:
-        data = chunk.model_dump()
-        yield f"data: {json.dumps(data)}\n\n"
+OPTIONAL_FIELDS = ("temperature", "max_tokens", "top_p", "presence_penalty", "frequency_penalty")
 
-    yield "data: [DONE]\n\n"
+SSE_HEADERS = {
+    "Cache-Control":     "no-cache",
+    "X-Accel-Buffering": "no",
+}
 
 # ── Endpoints ─────────────────────────────────────────────────────────────────
 
@@ -160,7 +197,7 @@ async def root():
     return {
         "service": "Multi-Model AI Gateway",
         "models":  list(MODELS.keys()),
-        "auth":    "Bearer <connect-key>",
+        "auth":    "Bearer connect",
         "docs":    "/docs",
     }
 
@@ -169,10 +206,7 @@ async def root():
 async def list_models(authorization: Optional[str] = Header(default=None)):
     verify_key(authorization)
     return {
-        name: {
-            "model_id":    cfg["model_id"],
-            "has_thinking": bool(cfg["extra_body"]),
-        }
+        name: {"model_id": cfg["model_id"]}
         for name, cfg in MODELS.items()
     }
 
@@ -190,54 +224,27 @@ async def chat(
             detail=f"Unknown model '{request.model}'. Available: {list(MODELS.keys())}",
         )
 
-    cfg = MODELS[request.model]
-
-    # Build messages list
-    messages = [m.model_dump() for m in request.messages]
-
-    # Inject per-model system prompt at position 0 if not already present
-    if request.inject_system_prompt:
-        system_prompt = get_system_prompt(request.model)
-        if not messages or messages[0].get("role") != "system":
-            messages.insert(0, {"role": "system", "content": system_prompt})
-
-    # Optional params
-    kwargs = {}
-    for field in ("temperature", "max_tokens", "top_p", "presence_penalty", "frequency_penalty"):
-        val = getattr(request, field)
-        if val is not None:
-            kwargs[field] = val
-
-    tools = [t.model_dump() for t in request.tools] if request.tools else None
+    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
 
     return StreamingResponse(
-        stream_nvidia(
-            model_name=request.model,
-            messages=messages,
-            tools=tools,
-            tool_choice=request.tool_choice,
-            kwargs=kwargs,
-            extra_body=cfg["extra_body"],
-        ),
+        stream_nvidia(request.model, messages, tools, request.tool_choice, kwargs),
         media_type="text/event-stream",
-        headers={
-            "Cache-Control":     "no-cache",
-            "X-Accel-Buffering": "no",
-        },
+        headers=SSE_HEADERS,
     )
 
 
-# ── Compatibility: OpenAI-style /v1/chat/completions ──────────────────────────
-
 @app.post("/v1/chat/completions")
 async def openai_compat(
     raw: Request,
     authorization: Optional[str] = Header(default=None),
 ):
-    """
-    Drop-in OpenAI-compatible endpoint.
-    Pass model as one of the gateway model names (e.g. 'Kimi-K2').
-    """
+    """OpenAI-compatible drop-in. Use gateway model names as 'model'."""
     verify_key(authorization)
     body = await raw.json()
 
@@ -248,35 +255,17 @@ async def openai_compat(
             detail=f"Unknown model '{model_name}'. Available: {list(MODELS.keys())}",
         )
 
-    cfg = MODELS[model_name]
-    messages = body.get("messages", [])
-
-    inject = body.get("inject_system_prompt", True)
-    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})
-
-    kwargs = {}
-    for field in ("temperature", "max_tokens", "top_p", "presence_penalty", "frequency_penalty"):
-        if field in body:
-            kwargs[field] = body[field]
-
-    tools      = body.get("tools")
+    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")
 
     return StreamingResponse(
-        stream_nvidia(
-            model_name=model_name,
-            messages=messages,
-            tools=tools,
-            tool_choice=tool_choice,
-            kwargs=kwargs,
-            extra_body=cfg["extra_body"],
-        ),
+        stream_nvidia(model_name, messages, tools, tool_choice, kwargs),
         media_type="text/event-stream",
-        headers={
-            "Cache-Control":     "no-cache",
-            "X-Accel-Buffering": "no",
-        },
-    )
+        headers=SSE_HEADERS,
+    )