Update server.py

server.py

@@ -3,6 +3,7 @@ import os
 import sys
 import time
 import uuid
+import asyncio
 from typing import List, Optional, AsyncGenerator, Iterable
 from contextlib import asynccontextmanager
 
@@ -43,6 +44,18 @@ def _parse_cors_origins(value: str) -> List[str]:
 
 CORS_ORIGINS = _parse_cors_origins(os.getenv("CORS_ORIGINS", "*"))
 GZIP_MIN_SIZE = int(os.getenv("GZIP_MIN_SIZE", "500"))
+MAX_REQUEST_BYTES = int(os.getenv("MAX_REQUEST_BYTES", str(1_000_000)))  # 1MB default
+
+# Fast-by-default generation settings (still fully overridable per request)
+DEFAULT_TEMPERATURE = float(os.getenv("DEFAULT_TEMPERATURE", "0.7"))
+DEFAULT_TOP_P = float(os.getenv("DEFAULT_TOP_P", "0.95"))
+DEFAULT_MAX_TOKENS = int(os.getenv("DEFAULT_MAX_TOKENS", "1024"))
+
+# Concurrency guard to keep tail latency low under spikes
+MAX_INFLIGHT_REQUESTS = int(os.getenv("MAX_INFLIGHT_REQUESTS", "200"))
+
+# "Thinking" increases latency; keep opt-in via env
+ENABLE_THINKING = os.getenv("ENABLE_THINKING", "0").strip().lower() in {"1", "true", "yes", "on"}
 
 # Model configurations
 ROX_CORE_MODEL = "minimaxai/minimax-m2.5"
@@ -86,6 +99,7 @@ async def lifespan(app: FastAPI):
         max_retries=max_retries,
         http_client=http_client,
     )
+    app.state.inflight_semaphore = asyncio.Semaphore(MAX_INFLIGHT_REQUESTS)
 
     try:
         yield
@@ -120,6 +134,15 @@ async def add_request_context(request: Request, call_next):
     request_id = request.headers.get("x-request-id") or str(uuid.uuid4())
     start = time.perf_counter()
     try:
+        # Protect server from huge bodies (DoS / latency blowups)
+        cl = request.headers.get("content-length")
+        if cl is not None:
+            try:
+                if int(cl) > MAX_REQUEST_BYTES:
+                    return JSONResponse(status_code=413, content={"error": "Request too large"})
+            except ValueError:
+                return JSONResponse(status_code=400, content={"error": "Invalid Content-Length"})
+
         response: Response = await call_next(request)
     finally:
         elapsed_ms = (time.perf_counter() - start) * 1000.0
@@ -148,6 +171,32 @@ def _client(app_: FastAPI) -> AsyncOpenAI:
         raise RuntimeError("Client not initialized")
     return c
 
+def _semaphore(app_: FastAPI) -> asyncio.Semaphore:
+    s = getattr(app_.state, "inflight_semaphore", None)
+    if s is None:
+        raise RuntimeError("Semaphore not initialized")
+    return s
+
+def _effective_temperature(value: Optional[float]) -> float:
+    return DEFAULT_TEMPERATURE if value is None else value
+
+def _effective_top_p(value: Optional[float]) -> float:
+    return DEFAULT_TOP_P if value is None else value
+
+def _effective_max_tokens(value: Optional[int], cap: int) -> int:
+    v = DEFAULT_MAX_TOKENS if value is None else value
+    if v < 1:
+        v = DEFAULT_MAX_TOKENS
+    return min(v, cap)
+
+def _sse_headers() -> dict:
+    # Helps proxies (nginx) avoid buffering and keeps SSE responsive
+    return {
+        "Cache-Control": "no-cache",
+        "Connection": "keep-alive",
+        "X-Accel-Buffering": "no",
+    }
+
 
 # Helper function for streaming responses
 async def stream_response(
@@ -161,25 +210,26 @@ async def stream_response(
 ) -> AsyncGenerator[str, None]:
     """Stream responses from OpenAI API"""
     try:
-        stream = await _client(app_).chat.completions.create(
-            model=model,
-            messages=messages,
-            temperature=temperature,
-            top_p=top_p,
-            max_tokens=max_tokens,
-            stream=True,
-            extra_body=extra_body
-        )
+        async with _semaphore(app_):
+            stream = await _client(app_).chat.completions.create(
+                model=model,
+                messages=messages,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                stream=True,
+                extra_body=extra_body
+            )
         
-        async for chunk in stream:
-            delta = chunk.choices[0].delta
-            content = getattr(delta, "content", None)
-            if content:
-                yield f"data: {json.dumps({'content': content})}\n\n"
+            async for chunk in stream:
+                delta = chunk.choices[0].delta
+                content = getattr(delta, "content", None)
+                if content:
+                    yield f"data: {json.dumps({'content': content}, separators=(',', ':'))}\n\n"
         
         yield "data: [DONE]\n\n"
     except Exception as e:
-        yield f"data: {json.dumps({'error': str(e)})}\n\n"
+        yield f"data: {json.dumps({'error': str(e)}, separators=(',', ':'))}\n\n"
 
 
 @app.get("/health")
@@ -271,9 +321,9 @@ class ChatMessage(BaseModel):
 
 class ChatRequest(BaseModel):
     messages: List[ChatMessage]
-    temperature: Optional[float] = 0.7
-    top_p: Optional[float] = 0.95
-    max_tokens: Optional[int] = 8192
+    temperature: Optional[float] = None
+    top_p: Optional[float] = None
+    max_tokens: Optional[int] = None
     stream: Optional[bool] = False
 
 
@@ -301,22 +351,27 @@ async def chat(req: ChatRequest):
     """Rox Core - Main conversational model with streaming support"""
     messages = [{"role": "system", "content": ROX_CORE_IDENTITY}]
     messages.extend([m.model_dump() for m in req.messages])
+    temperature = _effective_temperature(req.temperature)
+    top_p = _effective_top_p(req.top_p)
+    max_tokens = _effective_max_tokens(req.max_tokens, 8192)
     
     if req.stream:
         return StreamingResponse(
-            stream_response(app, ROX_CORE_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 8192)),
-            media_type="text/event-stream"
+            stream_response(app, ROX_CORE_MODEL, messages, temperature, top_p, max_tokens),
+            media_type="text/event-stream",
+            headers=_sse_headers(),
         )
     
     try:
-        completion = await _client(app).chat.completions.create(
-            model=ROX_CORE_MODEL,
-            messages=messages,
-            temperature=req.temperature,
-            top_p=req.top_p,
-            max_tokens=min(req.max_tokens, 8192),
-            stream=False
-        )
+        async with _semaphore(app):
+            completion = await _client(app).chat.completions.create(
+                model=ROX_CORE_MODEL,
+                messages=messages,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                stream=False
+            )
         return {"content": completion.choices[0].message.content or ""}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -327,22 +382,27 @@ async def turbo(req: ChatRequest):
     """Rox 2.1 Turbo - Fast and efficient with streaming"""
     messages = [{"role": "system", "content": ROX_TURBO_IDENTITY}]
     messages.extend([m.model_dump() for m in req.messages])
+    temperature = _effective_temperature(req.temperature)
+    top_p = _effective_top_p(req.top_p)
+    max_tokens = _effective_max_tokens(req.max_tokens, 8192)
     
     if req.stream:
         return StreamingResponse(
-            stream_response(app, ROX_TURBO_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 8192)),
-            media_type="text/event-stream"
+            stream_response(app, ROX_TURBO_MODEL, messages, temperature, top_p, max_tokens),
+            media_type="text/event-stream",
+            headers=_sse_headers(),
         )
     
     try:
-        completion = await _client(app).chat.completions.create(
-            model=ROX_TURBO_MODEL,
-            messages=messages,
-            temperature=req.temperature,
-            top_p=req.top_p,
-            max_tokens=min(req.max_tokens, 8192),
-            stream=False
-        )
+        async with _semaphore(app):
+            completion = await _client(app).chat.completions.create(
+                model=ROX_TURBO_MODEL,
+                messages=messages,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                stream=False
+            )
         return {"content": completion.choices[0].message.content or ""}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -353,30 +413,35 @@ async def coder(req: ChatRequest):
     """Rox 3.5 Coder - Specialized coding with streaming"""
     messages = [{"role": "system", "content": ROX_CODER_IDENTITY}]
     messages.extend([m.model_dump() for m in req.messages])
+    temperature = _effective_temperature(req.temperature)
+    top_p = _effective_top_p(req.top_p)
+    max_tokens = _effective_max_tokens(req.max_tokens, 16384)
     
     extra_body = {
         "top_k": 20,
         "presence_penalty": 0,
         "repetition_penalty": 1,
-        "chat_template_kwargs": {"enable_thinking": True}
+        "chat_template_kwargs": {"enable_thinking": ENABLE_THINKING}
     }
     
     if req.stream:
         return StreamingResponse(
-            stream_response(app, ROX_CODER_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 16384), extra_body),
-            media_type="text/event-stream"
+            stream_response(app, ROX_CODER_MODEL, messages, temperature, top_p, max_tokens, extra_body),
+            media_type="text/event-stream",
+            headers=_sse_headers(),
         )
     
     try:
-        completion = await _client(app).chat.completions.create(
-            model=ROX_CODER_MODEL,
-            messages=messages,
-            temperature=req.temperature,
-            top_p=req.top_p,
-            max_tokens=min(req.max_tokens, 16384),
-            stream=False,
-            extra_body=extra_body
-        )
+        async with _semaphore(app):
+            completion = await _client(app).chat.completions.create(
+                model=ROX_CODER_MODEL,
+                messages=messages,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                stream=False,
+                extra_body=extra_body
+            )
         return {"content": completion.choices[0].message.content or ""}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -387,25 +452,30 @@ async def turbo45(req: ChatRequest):
     """Rox 4.5 Turbo - Advanced reasoning with streaming"""
     messages = [{"role": "system", "content": ROX_TURBO_45_IDENTITY}]
     messages.extend([m.model_dump() for m in req.messages])
+    temperature = _effective_temperature(req.temperature)
+    top_p = _effective_top_p(req.top_p)
+    max_tokens = _effective_max_tokens(req.max_tokens, 8192)
     
-    extra_body = {"chat_template_kwargs": {"thinking": True}}
+    extra_body = {"chat_template_kwargs": {"thinking": ENABLE_THINKING}} if ENABLE_THINKING else None
     
     if req.stream:
         return StreamingResponse(
-            stream_response(app, ROX_TURBO_45_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 8192), extra_body),
-            media_type="text/event-stream"
+            stream_response(app, ROX_TURBO_45_MODEL, messages, temperature, top_p, max_tokens, extra_body),
+            media_type="text/event-stream",
+            headers=_sse_headers(),
         )
     
     try:
-        completion = await _client(app).chat.completions.create(
-            model=ROX_TURBO_45_MODEL,
-            messages=messages,
-            temperature=req.temperature,
-            top_p=req.top_p,
-            max_tokens=min(req.max_tokens, 8192),
-            stream=False,
-            extra_body=extra_body
-        )
+        async with _semaphore(app):
+            completion = await _client(app).chat.completions.create(
+                model=ROX_TURBO_45_MODEL,
+                messages=messages,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                stream=False,
+                extra_body=extra_body
+            )
         return {"content": completion.choices[0].message.content or ""}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -416,25 +486,30 @@ async def ultra(req: ChatRequest):
     """Rox 5 Ultra - Most advanced with streaming"""
     messages = [{"role": "system", "content": ROX_ULTRA_IDENTITY}]
     messages.extend([m.model_dump() for m in req.messages])
+    temperature = _effective_temperature(req.temperature)
+    top_p = _effective_top_p(req.top_p)
+    max_tokens = _effective_max_tokens(req.max_tokens, 8192)
     
-    extra_body = {"chat_template_kwargs": {"thinking": True}}
+    extra_body = {"chat_template_kwargs": {"thinking": ENABLE_THINKING}} if ENABLE_THINKING else None
     
     if req.stream:
         return StreamingResponse(
-            stream_response(app, ROX_ULTRA_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 8192), extra_body),
-            media_type="text/event-stream"
+            stream_response(app, ROX_ULTRA_MODEL, messages, temperature, top_p, max_tokens, extra_body),
+            media_type="text/event-stream",
+            headers=_sse_headers(),
         )
     
     try:
-        completion = await _client(app).chat.completions.create(
-            model=ROX_ULTRA_MODEL,
-            messages=messages,
-            temperature=req.temperature,
-            top_p=req.top_p,
-            max_tokens=min(req.max_tokens, 8192),
-            stream=False,
-            extra_body=extra_body
-        )
+        async with _semaphore(app):
+            completion = await _client(app).chat.completions.create(
+                model=ROX_ULTRA_MODEL,
+                messages=messages,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                stream=False,
+                extra_body=extra_body
+            )
         return {"content": completion.choices[0].message.content or ""}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -445,25 +520,30 @@ async def dyno(req: ChatRequest):
     """Rox 6 Dyno - Extended context with streaming"""
     messages = [{"role": "system", "content": ROX_DYNO_IDENTITY}]
     messages.extend([m.model_dump() for m in req.messages])
+    temperature = _effective_temperature(req.temperature)
+    top_p = _effective_top_p(req.top_p)
+    max_tokens = _effective_max_tokens(req.max_tokens, 16384)
     
-    extra_body = {"chat_template_kwargs": {"thinking": True}}
+    extra_body = {"chat_template_kwargs": {"thinking": ENABLE_THINKING}} if ENABLE_THINKING else None
     
     if req.stream:
         return StreamingResponse(
-            stream_response(app, ROX_DYNO_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 16384), extra_body),
-            media_type="text/event-stream"
+            stream_response(app, ROX_DYNO_MODEL, messages, temperature, top_p, max_tokens, extra_body),
+            media_type="text/event-stream",
+            headers=_sse_headers(),
         )
     
     try:
-        completion = await _client(app).chat.completions.create(
-            model=ROX_DYNO_MODEL,
-            messages=messages,
-            temperature=req.temperature,
-            top_p=req.top_p,
-            max_tokens=min(req.max_tokens, 16384),
-            stream=False,
-            extra_body=extra_body
-        )
+        async with _semaphore(app):
+            completion = await _client(app).chat.completions.create(
+                model=ROX_DYNO_MODEL,
+                messages=messages,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                stream=False,
+                extra_body=extra_body
+            )
         return {"content": completion.choices[0].message.content or ""}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -474,30 +554,35 @@ async def coder7(req: ChatRequest):
     """Rox 7 Coder - Most advanced coding with streaming"""
     messages = [{"role": "system", "content": ROX_CODER_7_IDENTITY}]
     messages.extend([m.model_dump() for m in req.messages])
+    temperature = _effective_temperature(req.temperature)
+    top_p = _effective_top_p(req.top_p)
+    max_tokens = _effective_max_tokens(req.max_tokens, 16384)
     
     extra_body = {
         "chat_template_kwargs": {
-            "enable_thinking": True,
+            "enable_thinking": ENABLE_THINKING,
             "clear_thinking": False
         }
     }
     
     if req.stream:
         return StreamingResponse(
-            stream_response(app, ROX_CODER_7_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 16384), extra_body),
-            media_type="text/event-stream"
+            stream_response(app, ROX_CODER_7_MODEL, messages, temperature, top_p, max_tokens, extra_body),
+            media_type="text/event-stream",
+            headers=_sse_headers(),
         )
     
     try:
-        completion = await _client(app).chat.completions.create(
-            model=ROX_CODER_7_MODEL,
-            messages=messages,
-            temperature=req.temperature,
-            top_p=req.top_p,
-            max_tokens=min(req.max_tokens, 16384),
-            stream=False,
-            extra_body=extra_body
-        )
+        async with _semaphore(app):
+            completion = await _client(app).chat.completions.create(
+                model=ROX_CODER_7_MODEL,
+                messages=messages,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                stream=False,
+                extra_body=extra_body
+            )
         return {"content": completion.choices[0].message.content or ""}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -508,22 +593,27 @@ async def vision(req: ChatRequest):
     """Rox Vision Max - Visual understanding with streaming"""
     messages = [{"role": "system", "content": ROX_VISION_IDENTITY}]
     messages.extend([m.model_dump() for m in req.messages])
+    temperature = _effective_temperature(req.temperature)
+    top_p = _effective_top_p(req.top_p)
+    max_tokens = _effective_max_tokens(req.max_tokens, 8192)
     
     if req.stream:
         return StreamingResponse(
-            stream_response(app, ROX_VISION_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 8192)),
-            media_type="text/event-stream"
+            stream_response(app, ROX_VISION_MODEL, messages, temperature, top_p, max_tokens),
+            media_type="text/event-stream",
+            headers=_sse_headers(),
         )
     
     try:
-        completion = await _client(app).chat.completions.create(
-            model=ROX_VISION_MODEL,
-            messages=messages,
-            temperature=req.temperature,
-            top_p=req.top_p,
-            max_tokens=min(req.max_tokens, 8192),
-            stream=False
-        )
+        async with _semaphore(app):
+            completion = await _client(app).chat.completions.create(
+                model=ROX_VISION_MODEL,
+                messages=messages,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                stream=False
+            )
         return {"content": completion.choices[0].message.content or ""}
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
@@ -539,14 +629,18 @@ async def hf_generate(req: HFRequest):
     ]
     
     try:
-        completion = await _client(app).chat.completions.create(
-            model=ROX_CORE_MODEL,
-            messages=messages,
-            temperature=params.temperature or 0.7,
-            top_p=params.top_p or 0.95,
-            max_tokens=params.max_new_tokens or 8192,
-            stream=False
-        )
+        temperature = _effective_temperature(params.temperature)
+        top_p = _effective_top_p(params.top_p)
+        max_tokens = _effective_max_tokens(params.max_new_tokens, 8192)
+        async with _semaphore(app):
+            completion = await _client(app).chat.completions.create(
+                model=ROX_CORE_MODEL,
+                messages=messages,
+                temperature=temperature,
+                top_p=top_p,
+                max_tokens=max_tokens,
+                stream=False
+            )
         return [{"generated_text": completion.choices[0].message.content or ""}]
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))