Update server.py

server.py

@@ -1,27 +1,28 @@
 import logging
 import os
 import sys
-from typing import List, Optional, AsyncGenerator
+import time
+import uuid
+from typing import List, Optional, AsyncGenerator, Iterable
 from contextlib import asynccontextmanager
 
 from dotenv import load_dotenv
-from fastapi import FastAPI, HTTPException, Request
+from fastapi import FastAPI, HTTPException, Request, Response
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.middleware.gzip import GZipMiddleware
 from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel, Field
-from openai import OpenAI
+from openai import AsyncOpenAI
+import httpx
 import json
 
 
 # Load environment variables
 load_dotenv()
 
-# Configure minimal logging for production speed
-logging.basicConfig(
-    level=logging.WARNING,
-    format='%(levelname)s - %(message)s'
-)
+# Configure logging (env-controlled)
+LOG_LEVEL = os.getenv("LOG_LEVEL", "WARNING").upper()
+logging.basicConfig(level=LOG_LEVEL, format="%(levelname)s - %(message)s")
 logger = logging.getLogger("rox_ai")
 
 # Check for API key
@@ -30,6 +31,19 @@ NVIDIA_API_KEY = os.getenv("NVIDIA_API_KEY")
 if not NVIDIA_API_KEY:
     raise RuntimeError("NVIDIA_API_KEY not set")
 
+API_BASE_URL = os.getenv("NVIDIA_BASE_URL", "https://integrate.api.nvidia.com/v1")
+
+def _parse_cors_origins(value: str) -> List[str]:
+    v = (value or "").strip()
+    if not v:
+        return []
+    if v == "*":
+        return ["*"]
+    return [o.strip() for o in v.split(",") if o.strip()]
+
+CORS_ORIGINS = _parse_cors_origins(os.getenv("CORS_ORIGINS", "*"))
+GZIP_MIN_SIZE = int(os.getenv("GZIP_MIN_SIZE", "500"))
+
 # Model configurations
 ROX_CORE_MODEL = "minimaxai/minimax-m2.5"
 ROX_TURBO_MODEL = "meta/llama-3.1-8b-instruct"  # Changed to a more reliable model
@@ -50,18 +64,33 @@ ROX_DYNO_IDENTITY = "You are Rox 6 Dyno, an AI model created by Rox AI. Your cre
 ROX_CODER_7_IDENTITY = "You are Rox 7 Coder, an AI model created by Rox AI. Your creator and owner is Mohammad Faiz, the founder of Rox AI company. You are the most advanced coding specialist."
 ROX_VISION_IDENTITY = "You are Rox Vision Max, an AI model created by Rox AI. Your creator and owner is Mohammad Faiz, the founder of Rox AI company. You specialize in visual understanding and multimodal tasks."
 
-# Initialize OpenAI client with timeout optimization
-client = OpenAI(
-    base_url="https://integrate.api.nvidia.com/v1",
-    api_key=NVIDIA_API_KEY,
-    timeout=60.0,
-    max_retries=2
-)
-
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     """Lifespan context manager"""
-    yield
+    # One pooled async HTTP client for all requests (keep-alive, limits, timeouts)
+    timeout_s = float(os.getenv("UPSTREAM_TIMEOUT_SECONDS", "60"))
+    max_retries = int(os.getenv("UPSTREAM_MAX_RETRIES", "2"))
+    max_connections = int(os.getenv("UPSTREAM_MAX_CONNECTIONS", "200"))
+    max_keepalive = int(os.getenv("UPSTREAM_MAX_KEEPALIVE_CONNECTIONS", "50"))
+
+    http_client = httpx.AsyncClient(
+        timeout=httpx.Timeout(timeout_s),
+        limits=httpx.Limits(max_connections=max_connections, max_keepalive_connections=max_keepalive),
+        headers={"User-Agent": "Rox-AI-API/2.0"},
+    )
+    app.state.http_client = http_client
+    app.state.client = AsyncOpenAI(
+        base_url=API_BASE_URL,
+        api_key=NVIDIA_API_KEY,
+        timeout=timeout_s,
+        max_retries=max_retries,
+        http_client=http_client,
+    )
+
+    try:
+        yield
+    finally:
+        await http_client.aclose()
 
 
 # Initialize FastAPI app - optimized for speed
@@ -75,38 +104,64 @@ app = FastAPI(
 )
 
 # GZip compression for faster transfers
-app.add_middleware(GZipMiddleware, minimum_size=500)
+app.add_middleware(GZipMiddleware, minimum_size=GZIP_MIN_SIZE)
 
-# CORS - unlimited access
+# CORS - env controlled (default "*")
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
+    allow_origins=CORS_ORIGINS,
+    allow_credentials=(CORS_ORIGINS != ["*"]),
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
+@app.middleware("http")
+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:
+        response: Response = await call_next(request)
+    finally:
+        elapsed_ms = (time.perf_counter() - start) * 1000.0
+        # Keep logs lightweight; only emit at INFO+ if enabled
+        if logger.isEnabledFor(logging.INFO):
+            logger.info("%s %s -> %.2fms id=%s", request.method, request.url.path, elapsed_ms, request_id)
+
+    response.headers["X-Request-Id"] = request_id
+    response.headers["X-Process-Time-Ms"] = f"{elapsed_ms:.2f}"
+    return response
+
 
 # Minimal exception handler
 @app.exception_handler(Exception)
 async def global_exception_handler(request: Request, exc: Exception):
+    logger.exception("Unhandled error on %s %s", request.method, request.url.path)
     return JSONResponse(
         status_code=500,
         content={"error": "Internal server error"}
     )
 
 
-@app.get("/health")
-def health_check():
-    """Health check endpoint"""
-    return {"status": "healthy", "service": "Rox AI", "version": "2.0"}
+def _client(app_: FastAPI) -> AsyncOpenAI:
+    c = getattr(app_.state, "client", None)
+    if c is None:
+        raise RuntimeError("Client not initialized")
+    return c
 
 
 # Helper function for streaming responses
-async def stream_response(model: str, messages: list, temperature: float, top_p: float, max_tokens: int, extra_body: dict = None):
+async def stream_response(
+    app_: FastAPI,
+    model: str,
+    messages: list,
+    temperature: float,
+    top_p: float,
+    max_tokens: int,
+    extra_body: dict | None = None,
+) -> AsyncGenerator[str, None]:
     """Stream responses from OpenAI API"""
     try:
-        stream = client.chat.completions.create(
+        stream = await _client(app_).chat.completions.create(
             model=model,
             messages=messages,
             temperature=temperature,
@@ -116,9 +171,11 @@ async def stream_response(model: str, messages: list, temperature: float, top_p:
             extra_body=extra_body
         )
         
-        for chunk in stream:
-            if chunk.choices[0].delta.content:
-                yield f"data: {json.dumps({'content': chunk.choices[0].delta.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})}\n\n"
         
         yield "data: [DONE]\n\n"
     except Exception as e:
@@ -243,21 +300,21 @@ class HFResponseItem(BaseModel):
 async def chat(req: ChatRequest):
     """Rox Core - Main conversational model with streaming support"""
     messages = [{"role": "system", "content": ROX_CORE_IDENTITY}]
-    messages.extend([m.dict() for m in req.messages])
+    messages.extend([m.model_dump() for m in req.messages])
     
     if req.stream:
         return StreamingResponse(
-            stream_response(ROX_CORE_MODEL, messages, req.temperature, req.top_p, req.max_tokens),
+            stream_response(app, ROX_CORE_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 8192)),
             media_type="text/event-stream"
         )
     
     try:
-        completion = client.chat.completions.create(
+        completion = await _client(app).chat.completions.create(
             model=ROX_CORE_MODEL,
             messages=messages,
             temperature=req.temperature,
             top_p=req.top_p,
-            max_tokens=req.max_tokens,
+            max_tokens=min(req.max_tokens, 8192),
             stream=False
         )
         return {"content": completion.choices[0].message.content or ""}
@@ -269,21 +326,21 @@ async def chat(req: ChatRequest):
 async def turbo(req: ChatRequest):
     """Rox 2.1 Turbo - Fast and efficient with streaming"""
     messages = [{"role": "system", "content": ROX_TURBO_IDENTITY}]
-    messages.extend([m.dict() for m in req.messages])
+    messages.extend([m.model_dump() for m in req.messages])
     
     if req.stream:
         return StreamingResponse(
-            stream_response(ROX_TURBO_MODEL, messages, req.temperature, req.top_p, req.max_tokens),
+            stream_response(app, ROX_TURBO_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 8192)),
             media_type="text/event-stream"
         )
     
     try:
-        completion = client.chat.completions.create(
+        completion = await _client(app).chat.completions.create(
             model=ROX_TURBO_MODEL,
             messages=messages,
             temperature=req.temperature,
             top_p=req.top_p,
-            max_tokens=req.max_tokens,
+            max_tokens=min(req.max_tokens, 8192),
             stream=False
         )
         return {"content": completion.choices[0].message.content or ""}
@@ -295,7 +352,7 @@ async def turbo(req: ChatRequest):
 async def coder(req: ChatRequest):
     """Rox 3.5 Coder - Specialized coding with streaming"""
     messages = [{"role": "system", "content": ROX_CODER_IDENTITY}]
-    messages.extend([m.dict() for m in req.messages])
+    messages.extend([m.model_dump() for m in req.messages])
     
     extra_body = {
         "top_k": 20,
@@ -306,12 +363,12 @@ async def coder(req: ChatRequest):
     
     if req.stream:
         return StreamingResponse(
-            stream_response(ROX_CODER_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 16384), extra_body),
+            stream_response(app, ROX_CODER_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 16384), extra_body),
             media_type="text/event-stream"
         )
     
     try:
-        completion = client.chat.completions.create(
+        completion = await _client(app).chat.completions.create(
             model=ROX_CODER_MODEL,
             messages=messages,
             temperature=req.temperature,
@@ -329,18 +386,18 @@ async def coder(req: ChatRequest):
 async def turbo45(req: ChatRequest):
     """Rox 4.5 Turbo - Advanced reasoning with streaming"""
     messages = [{"role": "system", "content": ROX_TURBO_45_IDENTITY}]
-    messages.extend([m.dict() for m in req.messages])
+    messages.extend([m.model_dump() for m in req.messages])
     
     extra_body = {"chat_template_kwargs": {"thinking": True}}
     
     if req.stream:
         return StreamingResponse(
-            stream_response(ROX_TURBO_45_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 8192), extra_body),
+            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"
         )
     
     try:
-        completion = client.chat.completions.create(
+        completion = await _client(app).chat.completions.create(
             model=ROX_TURBO_45_MODEL,
             messages=messages,
             temperature=req.temperature,
@@ -358,18 +415,18 @@ async def turbo45(req: ChatRequest):
 async def ultra(req: ChatRequest):
     """Rox 5 Ultra - Most advanced with streaming"""
     messages = [{"role": "system", "content": ROX_ULTRA_IDENTITY}]
-    messages.extend([m.dict() for m in req.messages])
+    messages.extend([m.model_dump() for m in req.messages])
     
     extra_body = {"chat_template_kwargs": {"thinking": True}}
     
     if req.stream:
         return StreamingResponse(
-            stream_response(ROX_ULTRA_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 8192), extra_body),
+            stream_response(app, ROX_ULTRA_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 8192), extra_body),
             media_type="text/event-stream"
         )
     
     try:
-        completion = client.chat.completions.create(
+        completion = await _client(app).chat.completions.create(
             model=ROX_ULTRA_MODEL,
             messages=messages,
             temperature=req.temperature,
@@ -387,18 +444,18 @@ async def ultra(req: ChatRequest):
 async def dyno(req: ChatRequest):
     """Rox 6 Dyno - Extended context with streaming"""
     messages = [{"role": "system", "content": ROX_DYNO_IDENTITY}]
-    messages.extend([m.dict() for m in req.messages])
+    messages.extend([m.model_dump() for m in req.messages])
     
     extra_body = {"chat_template_kwargs": {"thinking": True}}
     
     if req.stream:
         return StreamingResponse(
-            stream_response(ROX_DYNO_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 16384), extra_body),
+            stream_response(app, ROX_DYNO_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 16384), extra_body),
             media_type="text/event-stream"
         )
     
     try:
-        completion = client.chat.completions.create(
+        completion = await _client(app).chat.completions.create(
             model=ROX_DYNO_MODEL,
             messages=messages,
             temperature=req.temperature,
@@ -416,7 +473,7 @@ async def dyno(req: ChatRequest):
 async def coder7(req: ChatRequest):
     """Rox 7 Coder - Most advanced coding with streaming"""
     messages = [{"role": "system", "content": ROX_CODER_7_IDENTITY}]
-    messages.extend([m.dict() for m in req.messages])
+    messages.extend([m.model_dump() for m in req.messages])
     
     extra_body = {
         "chat_template_kwargs": {
@@ -427,12 +484,12 @@ async def coder7(req: ChatRequest):
     
     if req.stream:
         return StreamingResponse(
-            stream_response(ROX_CODER_7_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 16384), extra_body),
+            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"
         )
     
     try:
-        completion = client.chat.completions.create(
+        completion = await _client(app).chat.completions.create(
             model=ROX_CODER_7_MODEL,
             messages=messages,
             temperature=req.temperature,
@@ -450,16 +507,16 @@ async def coder7(req: ChatRequest):
 async def vision(req: ChatRequest):
     """Rox Vision Max - Visual understanding with streaming"""
     messages = [{"role": "system", "content": ROX_VISION_IDENTITY}]
-    messages.extend([m.dict() for m in req.messages])
+    messages.extend([m.model_dump() for m in req.messages])
     
     if req.stream:
         return StreamingResponse(
-            stream_response(ROX_VISION_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 8192)),
+            stream_response(app, ROX_VISION_MODEL, messages, req.temperature, req.top_p, min(req.max_tokens, 8192)),
             media_type="text/event-stream"
         )
     
     try:
-        completion = client.chat.completions.create(
+        completion = await _client(app).chat.completions.create(
             model=ROX_VISION_MODEL,
             messages=messages,
             temperature=req.temperature,
@@ -482,7 +539,7 @@ async def hf_generate(req: HFRequest):
     ]
     
     try:
-        completion = client.chat.completions.create(
+        completion = await _client(app).chat.completions.create(
             model=ROX_CORE_MODEL,
             messages=messages,
             temperature=params.temperature or 0.7,