import os import base64 import random import httpx from urllib.parse import quote from fastapi import APIRouter, Request, HTTPException, Header from fastapi.responses import Response, JSONResponse, StreamingResponse import re from typing import Optional, Any import json from helper.assets import ( save_base64_image, cleanup_image, is_base64_image, ) import asyncio from helper.ratelimit import ( enforce_rate_limit, resolve_rate_limit_identity, check_audio_rate_limit, check_video_rate_limit, check_image_rate_limit, MAX_CHAT_PROMPT_BYTES, MAX_CHAT_PROMPT_CHARS, MAX_GROQ_PROMPT_BYTES, MAX_GROQ_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, MAX_MEDIA_PROMPT_CHARS, extract_user_text, calculate_messages_size, normalize_prompt_value, enforce_prompt_size, resolve_bound_subject, get_usage_snapshot_for_subject, ) from helper.keywords import * from uuid import uuid4 from time import time from typing import Dict, List, Optional, Tuple router = APIRouter(prefix="/gen") PKEY = os.getenv("POLLINATIONS_KEY", "") PKEY2 = os.getenv("POLLINATIONS2_KEY", "") PKEY3 = os.getenv("POLLINATIONS3_KEY", "") AIRFORCE_KEY = os.getenv("AIRFORCE") AIRFORCE_VIDEO_MODEL = "grok-imagine-video" AIRFORCE_API_URL = "https://api.airforce/v1/images/generations" valid_ratios = {"3:2", "2:3", "1:1", "", None} ratios = {"3:2", "2:3", "1:1"} valid_modes = {"normal", "fun", "", None} modes = {"normal", "fun"} MODEL_MAP = { "llama-3.1-8b-instant": "Meta Llama 3.1 8B Instant", "gpt-4o-mini": "OpenAI GPT 4o Mini", "nemotron-3-super": "NVIDIA Nemotron 3 Super", "openai/gpt-oss-120b": "OpenAI GPT-OSS 120B", "openai/gpt-oss-20b": "OpenAI GPT-OSS 20B", "qwen-3-235b-a22b-instruct-2507": "Qwen3 Instruct", "llama-3.3-70b-versatile": "Meta Llama 3.3 70B Versatile", "meta-llama/llama-4-scout-17b-16e-instruct": "Meta Llama 4 Scout", } FALLBACK_MODEL = "meta-llama/llama-4-scout-17b-16e-instruct" FALLBACK_PROVIDER = "groq" # Header that API-key authenticated clients send so we know to stream # thinking tokens back to them. API_KEY_HEADER = "x-api-key" # ────────────────────────────────────────────── # CENTRAL ROUTING LOGIC # ────────────────────────────────────────────── def route_chat( messages: List[Dict[str, Any]], uses_tools: bool = False, ) -> Tuple[str, str]: """ Inspect messages and return (chosen_model, provider). This is the single source of truth for model selection. No API calls, no side-effects — pure routing logic. """ total_chars, total_bytes = calculate_messages_size(messages) prompt_text = extract_user_text(messages) long_context = is_long_context(messages) code_present = contains_code(prompt_text) math_heavy = is_math_heavy(prompt_text) structured_task = is_structured_task(prompt_text) multi_q = multiple_questions(prompt_text) code_heavy = is_code_heavy(prompt_text, code_present, long_context) has_images = contains_images(messages) score = 0 if long_context: score += 3 if math_heavy: score += 3 if structured_task: score += 2 if code_present: score += 2 if multi_q: score += 1 for kw in REASONING_KEYWORDS: if kw in prompt_text: score += 1 score = min(score, 10) # ── multimodal fast-path ────────────────── if has_images: return "gpt-4o-mini", "navy vision" # ── tool-use branch ────────────────────── if uses_tools: if long_context: return "nemotron-3-super", "navy" if score >= 6: return "nemotron-3-super", "navy" if score >= 4: return "openai/gpt-oss-120b", "groq" return "openai/gpt-oss-20b", "groq" # ── code branch ────────────────────────── if code_present: if code_heavy and score >= 6: return "o3-mini", "navy" if score >= 4: return "llama-3.3-70b-versatile", "groq" # ── general reasoning branch ───────────── if score >= 6: return "sonar", "navy" if score >= 4: return "meta-llama/llama-4-scout-17b-16e-instruct", "groq" # ── default ────────────────────────────── chosen_model, provider = "llama-3.1-8b-instant", "groq" # Groq context-size guard — promote to navy if too large if provider == "groq" and ( total_chars > MAX_GROQ_PROMPT_CHARS or total_bytes > MAX_GROQ_PROMPT_BYTES ): return "gpt-4o-mini", "navy" return chosen_model, provider def _log_routing( chosen_model: str, provider: str, messages: List[Dict[str, Any]], uses_tools: bool, ) -> None: prompt_text = extract_user_text(messages) long_context = is_long_context(messages) code_present = contains_code(prompt_text) math_heavy = is_math_heavy(prompt_text) structured_task = is_structured_task(prompt_text) multi_q = multiple_questions(prompt_text) has_images = contains_images(messages) print( f"\n[ADVANCED ROUTER]\n" f" Uses tools: {uses_tools}\n" f" Long context: {long_context}\n" f" Code present: {code_present}\n" f" Math heavy: {math_heavy}\n" f" Structured: {structured_task}\n" f" Multi-question:{multi_q}\n" f" Has images: {has_images}\n" f" → Selected: {chosen_model} ({provider})\n" ) # ────────────────────────────────────────────── # CENTRAL HTTP CALL # ────────────────────────────────────────────── def _get_provider_url_and_key(provider: str) -> Tuple[str, str]: """Return (url, api_key) for the given provider, raising on misconfiguration.""" if provider == "groq": keys = [k.strip() for k in os.getenv("GROQ_KEY", "").split(",") if k.strip()] if not keys: raise HTTPException(500, "Missing GROQ_KEY(s)") return "https://api.groq.com/openai/v1/chat/completions", random.choice(keys) if provider == "cerebras": keys = [k.strip() for k in os.getenv("CER_KEY", "").split(",") if k.strip()] if not keys: raise HTTPException(500, "Missing CER_KEY(s)") return "https://api.cerebras.ai/v1/chat/completions", random.choice(keys) if provider == "navy vision": keys = [k.strip() for k in os.getenv("NAVY_KEY", "").split(",") if k.strip()] if not keys: raise HTTPException(500, "Missing NAVY_KEY(s)") return "https://api.navy/v1/chat/completions", random.choice(keys) if provider == "navy": keys = [k.strip() for k in os.getenv("NAVY_TEXT_ONLY", "").split(",") if k.strip()] if not keys: raise HTTPException(500, "Missing NAVY_TEXT_ONLY key(s)") return "https://api.navy/v1/chat/completions", random.choice(keys) raise HTTPException(500, f"Unknown provider: {provider!r}") async def call_chat_completions( messages: List[Dict[str, Any]], model: str, provider: str, extra_body: Optional[Dict[str, Any]] = None, ) -> Dict[str, Any]: """ Resilient chat-completions call designed to survive Cloudflare 524 timeouts. Strategy: 1. Ask the upstream for a *streaming* response so bytes arrive before Cloudflare's ~100 s idle timeout fires. 2. Accumulate the stream into a single synthetic non-streaming payload so callers don't need to change. 3. Retry up to 2 times (with a short back-off) on 502/503/524. 4. On exhausted retries fall through to the Groq fallback. """ url, api_key = _get_provider_url_and_key(provider) headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"} # Always request streaming upstream — we reassemble below. body: Dict[str, Any] = {"model": model, "messages": messages, "stream": True} if extra_body: body.update(extra_body) body["stream"] = True # force streaming even if caller passed stream=False TRANSIENT = {502, 503, 524, 429} MAX_ATTEMPTS = 3 last_exc: Optional[Exception] = None for attempt in range(MAX_ATTEMPTS): if attempt: await asyncio.sleep(2 ** attempt) # 2 s, 4 s try: async with httpx.AsyncClient(timeout=httpx.Timeout(300.0, read=300.0)) as client: async with client.stream("POST", url, json=body, headers=headers) as r: # Transient upstream error — retry. if r.status_code in TRANSIENT: body_bytes = await r.aread() last_exc = HTTPException( status_code=r.status_code, detail=body_bytes.decode("utf-8", errors="replace")[:500], ) print(f"[call_chat_completions] attempt {attempt+1} got {r.status_code}, retrying…") continue if r.status_code != 200: body_bytes = await r.aread() raise HTTPException( status_code=r.status_code, detail=body_bytes.decode("utf-8", errors="replace")[:1000], ) # ── Reassemble streaming SSE into a single response object ── accumulated_content = "" accumulated_reasoning = "" tool_calls_map: Dict[int, Dict[str, Any]] = {} usage: Dict[str, Any] = {} finish_reason: Optional[str] = None resp_id = "" resp_model = model async for line in r.aiter_lines(): if not line or not line.startswith("data:"): continue raw = line[5:].strip() if raw == "[DONE]": break try: obj = json.loads(raw) except Exception: continue if not isinstance(obj, dict): continue resp_id = resp_id or obj.get("id", "") resp_model = obj.get("model", resp_model) if "usage" in obj and obj["usage"]: usage = obj["usage"] choices = obj.get("choices") or [] if not choices: continue choice = choices[0] finish_reason = choice.get("finish_reason") or finish_reason delta = choice.get("delta") or {} # Accumulate text content. dc = delta.get("content") if dc: accumulated_content += dc # Accumulate reasoning / thinking tokens. dr = delta.get("reasoning_content") or delta.get("reasoning") if dr: accumulated_reasoning += dr # Accumulate tool-call argument chunks (streamed as fragments). for tc_delta in (delta.get("tool_calls") or []): idx = tc_delta.get("index", 0) if idx not in tool_calls_map: tool_calls_map[idx] = { "id": tc_delta.get("id", ""), "type": tc_delta.get("type", "function"), "function": {"name": "", "arguments": ""}, } existing = tool_calls_map[idx] if tc_delta.get("id"): existing["id"] = tc_delta["id"] fn_delta = tc_delta.get("function") or {} if fn_delta.get("name"): existing["function"]["name"] += fn_delta["name"] if fn_delta.get("arguments"): existing["function"]["arguments"] += fn_delta["arguments"] # Reassemble into a standard non-streaming response shape. tool_calls_list = [tool_calls_map[i] for i in sorted(tool_calls_map)] message: Dict[str, Any] = {"role": "assistant", "content": accumulated_content} if accumulated_reasoning: message["reasoning_content"] = accumulated_reasoning if tool_calls_list: message["tool_calls"] = tool_calls_list return { "id": resp_id, "object": "chat.completion", "model": resp_model, "choices": [ { "index": 0, "message": message, "finish_reason": finish_reason or "stop", } ], "usage": usage, } except HTTPException: raise except (httpx.RemoteProtocolError, httpx.ReadError, httpx.ConnectError) as exc: last_exc = exc print(f"[call_chat_completions] attempt {attempt+1} network error: {exc}, retrying…") continue # All attempts exhausted — fall back to Groq. print(f"[call_chat_completions] all attempts failed ({last_exc}), falling back to Groq") fb_url, fb_key = _get_provider_url_and_key(FALLBACK_PROVIDER) fb_headers = {"Authorization": f"Bearer {fb_key}", "Content-Type": "application/json"} fallback_body = { "model": FALLBACK_MODEL, "messages": messages, "stream": False, } if extra_body: # Forward tools/tool_choice but not stream override. for k in ("tools", "tool_choice"): if k in extra_body: fallback_body[k] = extra_body[k] async with httpx.AsyncClient(timeout=httpx.Timeout(120.0)) as client: fb_r = await client.post(fb_url, json=fallback_body, headers=fb_headers) if fb_r.status_code != 200: raise HTTPException( status_code=fb_r.status_code, detail=f"Primary and fallback both failed. Fallback: {fb_r.text[:500]}", ) return fb_r.json() def _extract_text_from_response(data: Dict[str, Any]) -> str: try: return data["choices"][0]["message"]["content"] or "" except Exception: return "" def _extract_usage(data: Dict[str, Any]) -> Tuple[int, int]: usage = data.get("usage", {}) input_tok = usage.get("prompt_tokens") or usage.get("input_tokens", 0) output_tok = usage.get("completion_tokens") or usage.get("output_tokens", 0) return input_tok, output_tok # ────────────────────────────────────────────── # HELPER: image generation # ────────────────────────────────────────────── def is_cinematic_image_prompt(prompt: str) -> bool: for kw in CREATIVE_KEYWORDS: if kw in prompt.lower(): return True return False def _is_api_key_request(request: Request) -> bool: """ Return True when the caller authenticated with an API key rather than a session cookie / browser auth. We use this to decide whether to forward think-tag / reasoning_content tokens to the client. """ return bool( request.headers.get(API_KEY_HEADER) or request.headers.get("authorization", "").lower().startswith("bearer ") ) def _inject_reasoning_into_chunk(obj: Dict[str, Any]) -> Dict[str, Any]: """ Some navy models return thinking tokens in a non-standard ``reasoning_content`` field inside each delta. When that field is present we wrap it in … and prepend it to the regular ``content`` delta so that every SSE-speaking client sees a single, unified text stream. The original ``reasoning_content`` field is preserved so clients that know about it can still use it directly. """ try: delta = obj["choices"][0]["delta"] except (KeyError, IndexError, TypeError): return obj reasoning = delta.get("reasoning_content") or delta.get("reasoning") or "" content = delta.get("content") or "" if reasoning and isinstance(reasoning, str): # Wrap in tags and prepend to the visible content delta. wrapped = f"{reasoning}" delta["content"] = wrapped + content # Keep the raw field so native clients can parse it too. delta["reasoning_content"] = reasoning obj["choices"][0]["delta"] = delta return obj def _normalize_usage_block(obj: Dict[str, Any]) -> Dict[str, Any]: """Rewrite the usage block to a canonical shape (in-place, returns obj).""" if "usage" not in obj or not isinstance(obj.get("usage"), dict): return obj u = obj["usage"] input_tok = u.get("prompt_tokens") or u.get("input_tokens", 0) output_tok = u.get("completion_tokens") or u.get("output_tokens", 0) obj["usage"] = { "prompt_tokens": input_tok, "completion_tokens": output_tok, "total_tokens": input_tok + output_tok, "input_tokens": input_tok, "output_tokens": output_tok, } return obj # ────────────────────────────────────────────── # IMAGE GENERATION # ────────────────────────────────────────────── @router.post("/image") @router.get("/image/{prompt}") async def generate_image( request: Request, prompt: str = None, authorization: str = Header(None), x_client_id: str = Header(None), ): timeout = httpx.Timeout(300.0, read=300.0) if prompt is None: payload = await request.json() prompt = payload.get("prompt") mode = payload.get("mode") image_urls = payload.get("image_urls") else: mode = request.query_params.get("mode") image_urls = request.query_params.getlist("image_urls") prompt = normalize_prompt_value(prompt, "prompt") enforce_prompt_size(prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Image prompt") await check_image_rate_limit(request, authorization, x_client_id) chosen_model = "zimage" if is_cinematic_image_prompt(prompt): chosen_model = "flux" if isinstance(mode, str): m = mode.strip().lower() if m == "fantasy": chosen_model = "flux" elif m == "realistic": chosen_model = "zimage" has_input_image = bool(image_urls) temp_assets = [] if has_input_image: chosen_model = "klein" params = {"model": chosen_model, "key": PKEY2} if has_input_image: processed = [] for img in image_urls[:2]: if is_base64_image(img): image_id = save_base64_image(img) temp_assets.append(image_id) served = f"{request.base_url}asset-cdn/assets/{image_id}" processed.append(served) else: processed.append(img) params["image"] = "|".join(processed) encoded_prompt = quote(prompt, safe="") query = "&".join(f"{k}={quote(str(v), safe='')}" for k, v in params.items()) url = f"https://gen.pollinations.ai/image/{encoded_prompt}?{query}" try: async with httpx.AsyncClient(timeout=timeout) as client: resp = await client.get(url) finally: for aid in temp_assets: cleanup_image(aid) if resp.status_code != 200: raise HTTPException(500, f"Pollinations error: {resp.status_code}") return Response(content=resp.content, media_type="image/jpeg") # ────────────────────────────────────────────── # SFX GENERATION # ────────────────────────────────────────────── @router.get("/sfx/{prompt}") @router.post("/sfx") async def gensfx( request: Request, prompt: str = None, authorization: str = Header(None), x_client_id: str = Header(None), ): if prompt is None: payload = await request.json() prompt = payload.get("prompt") prompt = normalize_prompt_value(prompt, "prompt") enforce_prompt_size(prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Audio prompt") await check_audio_rate_limit(request, authorization, x_client_id) url = f"https://gen.pollinations.ai/audio/{prompt}?model=acestep&key={PKEY}" async with httpx.AsyncClient(timeout=None) as client: resp = await client.get(url) if resp.status_code != 200: return JSONResponse( status_code=resp.status_code, content={"success": False, "error": "Upstream music/sfx generation failed"}, ) return Response(resp.content, media_type="audio/mpeg") # ────────────────────────────────────────────── # TTS GENERATION # ────────────────────────────────────────────── @router.get("/tts/{prompt}") @router.post("/tts") async def gentts( request: Request, prompt: str = None, authorization: str = Header(None), x_client_id: str = Header(None), ): if prompt is None: payload = await request.json() prompt = payload.get("prompt") prompt = normalize_prompt_value(prompt, "prompt") enforce_prompt_size(prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Audio prompt") await check_audio_rate_limit(request, authorization, x_client_id) url = f"https://gen.pollinations.ai/audio/{prompt}?key={PKEY3}" async with httpx.AsyncClient(timeout=None) as client: resp = await client.get(url) if resp.status_code != 200: return JSONResponse( status_code=resp.status_code, content={"success": False, "error": "Upstream audio generation failed"}, ) return Response(resp.content, media_type="audio/mpeg") # ────────────────────────────────────────────── # VIDEO GENERATION (Pollinations) # ────────────────────────────────────────────── @router.get("/video/{prompt}") @router.post("/video") @router.head("/video") async def genvideo( request: Request, prompt: str = None, authorization: str = Header(None), x_client_id: str = Header(None), ): if request.method == "HEAD": return Response( status_code=200, headers={ "Y-prompt": "string — required. The text prompt used to generate the video.", "Y-ratio": "string — optional. Aspect ratio of the output video.", "Y-ratio-values": "3:2,2:3,1:1", "Y-ratio-default": "3:2", "Y-mode": "string — optional. Controls generation style.", "Y-mode-values": "normal,fun", "Y-mode-default": "normal", "Y-duration": "integer — optional. Duration in seconds (1–10).", "Y-duration-default": "5", "Y-image_urls": "array — optional. Up to 2 image URLs for conditioning.", "Y-image_urls-max": "2", "Y-response_format": "video/mp4", "Y-model": "grok-video", }, ) aspectRatio = "3:2" inputMode = "normal" duration = 5 image_urls = None if prompt is None: user_body = await request.json() prompt = user_body.get("prompt") ratio = user_body.get("ratio") mode = user_body.get("mode") image_urls = user_body.get("image_urls") duration = user_body.get("duration", 5) if ratio not in valid_ratios: raise HTTPException(400, f"Invalid aspect ratio '{ratio}'. Must be one of 3:2, 2:3, or 1:1.") if ratio in ratios: aspectRatio = ratio if mode not in valid_modes: raise HTTPException(400, f"Invalid mode '{mode}'. Must be 'normal' or 'fun'.") if mode in modes: inputMode = mode if image_urls: if not isinstance(image_urls, list): raise HTTPException(400, "image_urls must be a list") if len(image_urls) > 2: raise HTTPException(400, "You may provide at most two image URLs") try: duration = max(1, min(10, int(duration))) except (TypeError, ValueError): duration = 5 prompt = normalize_prompt_value(prompt, "prompt") enforce_prompt_size(prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Video prompt") await check_video_rate_limit(request, authorization, x_client_id) RATIO_MAP = {"3:2": "16:9", "2:3": "9:16", "1:1": "9:16"} pollinations_ratio = RATIO_MAP.get(aspectRatio, "16:9") encoded_prompt = quote(prompt, safe="") params = { "model": "ltx-2", "duration": duration, "aspectRatio": pollinations_ratio, "seed": -1, } temp_assets = [] if image_urls: processed_urls = [] for img in image_urls[:2]: if is_base64_image(img): image_id = save_base64_image(img) temp_assets.append(image_id) served_url = f"{request.base_url}asset-cdn/assets/{image_id}" processed_urls.append(served_url) else: processed_urls.append(img) params["image"] = "|".join(processed_urls) if inputMode == "fun": params["enhance"] = "true" query_string = "&".join(f"{k}={quote(str(v), safe='')}" for k, v in params.items()) url = f"https://gen.pollinations.ai/image/{encoded_prompt}?{query_string}&key={PKEY}" print(f"[VIDEO GEN] Pollinations URL: {url}") resp = None try: async with httpx.AsyncClient(timeout=600) as client: resp = await client.get(url) finally: for aid in temp_assets: cleanup_image(aid) if resp is None: raise HTTPException(502, "Video generation request failed") if resp.status_code != 200: body_text = "" try: body_text = resp.text except Exception: pass return JSONResponse( status_code=resp.status_code, content={ "success": False, "error": "Upstream video generation failed", "status_code": resp.status_code, "message": body_text[:1000], }, ) if not resp.content: raise HTTPException(502, "Pollinations returned empty response") return Response( content=resp.content, media_type="video/mp4", headers={ "Content-Length": str(len(resp.content)), "Accept-Ranges": "bytes", }, ) # ────────────────────────────────────────────── # VIDEO GENERATION (Airforce) # ────────────────────────────────────────────── @router.get("/video/airforce/{prompt}") @router.post("/video/airforce") async def genvideo_airforce( request: Request, prompt: str = None, authorization: str = Header(None), x_client_id: str = Header(None), ): if request.method == "HEAD": return Response( status_code=200, headers={ "Y-prompt": "string — required. The text prompt used to generate the video.", "Y-ratio": "string — optional. Aspect ratio of the output video.", "Y-ratio-values": "3:2,2:3,1:1", "Y-ratio-default": "3:2", "Y-mode": "string — optional. Controls generation style.", "Y-mode-values": "normal,fun", "Y-mode-default": "normal", "Y-duration": "integer — optional. Duration in seconds.", "Y-duration-default": "5", "Y-image_urls": "array — optional. Up to 2 image URLs for conditioning.", "Y-image_urls-max": "2", "Y-response_format": "video/mp4", "Y-model": "grok-imagine-video", }, ) aspectRatio = "3:2" inputMode = "normal" image_urls = None if prompt is None: user_body = await request.json() prompt = user_body.get("prompt") ratio = user_body.get("ratio") mode = user_body.get("mode") image_urls = user_body.get("image_urls") if ratio not in valid_ratios: raise HTTPException(400, f"Invalid aspect ratio {ratio}. Must be one of 3:2, 2:3, or 1:1. Default is 3:2") if ratio in ratios: aspectRatio = ratio if mode not in valid_modes: raise HTTPException(400, f"Invalid mode {mode}. Must be 'normal' or 'fun'. Default is normal") if mode in modes: inputMode = mode if image_urls: if not isinstance(image_urls, list): raise HTTPException(400, "image_urls must be a list") if len(image_urls) > 2: raise HTTPException(400, "You may provide at most two image URLs") prompt = normalize_prompt_value(prompt, "prompt") enforce_prompt_size(prompt, MAX_MEDIA_PROMPT_CHARS, MAX_MEDIA_PROMPT_BYTES, "Video prompt") await check_video_rate_limit(request, authorization, x_client_id) payload = { "model": AIRFORCE_VIDEO_MODEL, "prompt": prompt, "n": 1, "size": "1024x1024", "response_format": "b64_json", "sse": False, "mode": inputMode, "aspectRatio": aspectRatio, } if image_urls: payload["image_urls"] = image_urls async with httpx.AsyncClient(timeout=600) as client: resp = await client.post( AIRFORCE_API_URL, headers={"Authorization": f"Bearer {AIRFORCE_KEY}", "Content-Type": "application/json"}, json=payload, ) if resp.status_code != 200: return JSONResponse(status_code=resp.status_code, content=resp.json()) if not resp.content: raise HTTPException(502, "api.airforce returned empty response") try: result = resp.json() b64_video = result["data"][0]["b64_json"] except Exception: raise HTTPException(502, f"Invalid api.airforce response: {resp.text[:500]}") if not b64_video: raise HTTPException(502, "Airforce returned empty b64_json") video_bytes = base64.b64decode(b64_video) return Response( content=video_bytes, media_type="video/mp4", headers={ "Content-Length": str(len(video_bytes)), "Accept-Ranges": "bytes", }, ) # ────────────────────────────────────────────── # CHAT COMPLETIONS (/gen/chat/completions) # ────────────────────────────────────────────── async def _check_chat_rate_limit( request: Request, authorization: Optional[str], client_id: Optional[str] = None, ): return await enforce_rate_limit(request, authorization, "cloudChatDaily", client_id) @router.post("/chat/completions") async def generate_text( request: Request, authorization: Optional[str] = Header(None), x_client_id: Optional[str] = Header(None), ): body = await request.json() messages = body.get("messages", []) if not isinstance(messages, list) or len(messages) == 0: raise HTTPException(400, "messages[] is required") uses_tools = ( "tools" in body and isinstance(body["tools"], list) and len(body["tools"]) > 0 ) or ("tool_choice" in body and body["tool_choice"] not in [None, "none"]) chosen_model, provider = route_chat(messages, uses_tools=uses_tools) _log_routing(chosen_model, provider, messages, uses_tools) await _check_chat_rate_limit(request, authorization, x_client_id) # Determine whether the caller is an API-key client that should receive # raw thinking tokens. forward_thinking = _is_api_key_request(request) body["model"] = chosen_model stream = body.get("stream", False) url, api_key = _get_provider_url_and_key(provider) headers = {"Authorization": f"Bearer {api_key}"} if stream: body["stream"] = True async def stream_fallback(client: httpx.AsyncClient): fallback_body = { "model": FALLBACK_MODEL, "messages": body["messages"], "stream": True, } fb_url, fb_key = _get_provider_url_and_key(FALLBACK_PROVIDER) fb_headers = {"Authorization": f"Bearer {fb_key}"} print("[FALLBACK] Starting Groq fallback stream") async with client.stream("POST", fb_url, json=fallback_body, headers=fb_headers) as r: if r.status_code >= 400: err = (await r.aread()).decode("utf-8", errors="replace") yield f'data: {{"error": "Fallback provider failed: {err[:500]}"}}\n\n' return async for line in r.aiter_lines(): if not line: yield "\n" continue yield (line if line.startswith("data:") else f"data: {line}\n\n") + "\n" async def stream_primary(client: httpx.AsyncClient): try: async with client.stream("POST", url, json=body, headers=headers) as r: if r.status_code >= 400: print("[STREAM FALLBACK] Primary provider failed → switching to fallback") async for chunk in stream_fallback(client): yield chunk return async for line in r.aiter_lines(): if not line: yield "\n" continue if line.startswith("data:"): try: obj = json.loads(line[5:].strip()) if isinstance(obj, dict) and isinstance(obj.get("error"), dict): async for chunk in stream_fallback(client): yield chunk return except Exception: pass yield line + "\n" except Exception as e: print(f"[STREAM ERROR] {e}") async for chunk in stream_fallback(client): yield chunk async def event_generator(): sent_metadata = False async with httpx.AsyncClient(timeout=None) as client: async for chunk in stream_primary(client): # ── emit router metadata once as the very first SSE frame ── if not sent_metadata: meta = { "router_metadata": { "model_name": MODEL_MAP.get(chosen_model, chosen_model) } } yield f"data: {json.dumps(meta)}\n\n" sent_metadata = True # ── pass [DONE] straight through ────────────────────────── if "data: [DONE]" in chunk: yield chunk continue # ── process data: … lines ───────────────────────────────── if chunk.startswith("data:"): raw = chunk[5:].strip() try: obj = json.loads(raw) except Exception: # Not valid JSON — forward verbatim (keeps partial # chunks from blocking the stream). yield chunk continue if not isinstance(obj, dict): yield chunk continue # Normalize usage block whenever it appears. _normalize_usage_block(obj) # ── thinking / reasoning tokens ─────────────────────── # Navy models may embed thinking in two ways: # # 1. As delta.reasoning_content (separate field) # 2. Inline inside delta.content wrapped in … # # For API-key callers we always surface both forms. # For browser/session callers we strip reasoning_content # so it doesn't confuse UI clients that don't expect it, # but tags already present in content are left # alone (they arrived that way from upstream). if forward_thinking: # Merge reasoning_content into content as # … and keep the raw field. obj = _inject_reasoning_into_chunk(obj) else: # Strip the non-standard field so browser clients # don't see unexpected keys. try: delta = obj["choices"][0]["delta"] delta.pop("reasoning_content", None) delta.pop("reasoning", None) obj["choices"][0]["delta"] = delta except (KeyError, IndexError, TypeError): pass yield f"data: {json.dumps(obj)}\n\n" continue # ── any other line (comments, keep-alives, …) ───────────── yield chunk return StreamingResponse( event_generator(), media_type="text/event-stream", headers={ "Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no", }, ) # ── non-streaming ───────────────────────── async with httpx.AsyncClient(timeout=None) as client: r = await client.post(url, json=body, headers=headers) # navy-vision fallback if provider == "navy vision" and r.status_code >= 400: print("[FALLBACK] Navy vision failed — switching to fallback") fb_url, fb_key = _get_provider_url_and_key(FALLBACK_PROVIDER) fallback_body = dict(body) fallback_body["model"] = FALLBACK_MODEL r = await client.post( fb_url, json=fallback_body, headers={"Authorization": f"Bearer {fb_key}"}, ) content_type = (r.headers.get("content-type") or "").lower() if "application/json" in content_type: try: payload = r.json() except Exception: payload = {"error": "Upstream returned invalid JSON"} else: # Normalize usage fields. _normalize_usage_block(payload) # ── thinking tokens in non-streaming responses ──────────────────── # Some navy models put thinking content in # message.reasoning_content. For API-key callers we prepend it to # message.content wrapped in …; for others we drop # the non-standard field. try: message = payload["choices"][0]["message"] reasoning = ( message.pop("reasoning_content", None) or message.pop("reasoning", None) or "" ) if reasoning and isinstance(reasoning, str): if forward_thinking: existing = message.get("content") or "" message["content"] = f"{reasoning}{existing}" # Restore the raw field for clients that want it. message["reasoning_content"] = reasoning # else: already popped — nothing to do. payload["choices"][0]["message"] = message except (KeyError, IndexError, TypeError): pass payload.setdefault("router_metadata", {})["model_name"] = MODEL_MAP.get( chosen_model, chosen_model ) else: payload = { "error": "Upstream returned non-JSON response", "status_code": r.status_code, "message": r.text[:1000], } return JSONResponse(status_code=r.status_code, content=payload) # ────────────────────────────────────────────── # PROMPT ANALYZE (/gen/prompt_analyze) # ────────────────────────────────────────────── @router.post("/prompt_analyze") async def analyze_prompt(request: Request): body = await request.json() messages = body.get("prompt", []) if not isinstance(messages, list) or len(messages) == 0: raise HTTPException(400, "messages[] is required") uses_tools = ( "tools" in body and isinstance(body["tools"], list) and len(body["tools"]) > 0 ) or ("tool_choice" in body and body["tool_choice"] not in [None, "none"]) chosen_model, _ = route_chat(messages, uses_tools=uses_tools) return {MODEL_MAP.get(chosen_model, chosen_model)} # ────────────────────────────────────────────── # MODELS LIST # ────────────────────────────────────────────── @router.get("/models") def return_models_openai(): return { "object": "list", "data": [ { "id": "lightning", "object": "model", "created": 1767225600, "owned_by": "inferenceport-ai", } ], } # ────────────────────────────────────────────── # RESPONSES API (/gen/responses) # ────────────────────────────────────────────── def _resp_id(prefix: str) -> str: return f"{prefix}_{uuid4().hex}" def _resp_ts() -> int: return int(time()) def _content_to_text(content: Any) -> str: if isinstance(content, str): return content if isinstance(content, list): parts = [] for item in content: if isinstance(item, dict) and item.get("type") in ("input_text", "output_text", "text"): txt = item.get("text") if isinstance(txt, str): parts.append(txt) return "".join(parts) return "" def _responses_input_to_messages( input_data: Any, instructions: Optional[str] = None, ) -> List[Dict[str, Any]]: messages: List[Dict[str, Any]] = [] if instructions: messages.append({"role": "developer", "content": instructions}) if isinstance(input_data, str): messages.append({"role": "user", "content": input_data}) return messages if isinstance(input_data, list): for item in input_data: if isinstance(item, str): messages.append({"role": "user", "content": item}) continue if not isinstance(item, dict): continue role = item.get("role", "user") text = _content_to_text(item.get("content", "")) if text: messages.append({"role": role, "content": text}) return messages def _build_responses_payload( model: str, text: str, response_id: str, input_tokens: int = 0, output_tokens: int = 0, tool_calls: Optional[List[Dict[str, Any]]] = None, ) -> Dict[str, Any]: # Build content: text part first, then one function_call part per tool call content: List[Dict[str, Any]] = [] if text: content.append({"type": "output_text", "text": text, "annotations": []}) for tc in (tool_calls or []): fn = tc.get("function", {}) content.append({ "type": "tool_use", "id": tc.get("id", _resp_id("tool")), "name": fn.get("name", ""), "input": json.loads(fn["arguments"]) if fn.get("arguments") else {}, }) # Top-level output items: one message item (text) + one per tool call output_items: List[Dict[str, Any]] = [] if text or not tool_calls: output_items.append({ "id": _resp_id("msg"), "type": "message", "role": "assistant", "status": "completed", "content": [c for c in content if c["type"] == "output_text"], }) for tc in (tool_calls or []): fn = tc.get("function", {}) output_items.append({ "id": tc.get("id", _resp_id("tool")), "type": "function_call", "call_id": tc.get("id", ""), "name": fn.get("name", ""), "arguments": fn.get("arguments", "{}"), "status": "completed", }) return { "id": response_id, "object": "response", "created_at": _resp_ts(), "status": "completed", "completed_at": _resp_ts(), "error": None, "incomplete_details": None, "instructions": None, "max_output_tokens": None, "model": model, "output": output_items, "output_text": text, "usage": { "input_tokens": input_tokens, "output_tokens": output_tokens, "total_tokens": input_tokens + output_tokens, }, } @router.post("/responses") async def create_responses( request: Request, authorization: Optional[str] = Header(None), x_client_id: Optional[str] = Header(None), ): body = await request.json() model = body.get("model") input_data = body.get("input") instructions = body.get("instructions") stream = body.get("stream", False) tools = body.get("tools") tool_choice = body.get("tool_choice") if not model: raise HTTPException(400, "model is required") if input_data is None: raise HTTPException(400, "input is required") messages = _responses_input_to_messages(input_data, instructions=instructions) if not messages: raise HTTPException(400, "input could not be parsed") uses_tools = bool(tools) or (tool_choice not in [None, "none"]) # Build extra fields to forward upstream extra_body: Dict[str, Any] = {} if tools: extra_body["tools"] = tools if tool_choice is not None: extra_body["tool_choice"] = tool_choice chosen_model, provider = route_chat(messages, uses_tools=uses_tools) _log_routing(chosen_model, provider, messages, uses_tools=uses_tools) await _check_chat_rate_limit(request, authorization, x_client_id) async def _generate() -> Tuple[str, List[Dict[str, Any]], int, int]: data = await call_chat_completions( messages, chosen_model, provider, extra_body=extra_body or None ) input_tokens, output_tokens = _extract_usage(data) message = data.get("choices", [{}])[0].get("message", {}) text = message.get("content") or "" tool_calls = message.get("tool_calls") or [] return text, tool_calls, input_tokens, output_tokens # ── non-streaming ───────────────────────── if stream is False: text, tool_calls, input_tokens, output_tokens = await _generate() response_id = _resp_id("resp") return JSONResponse( content=_build_responses_payload( chosen_model, text, response_id, input_tokens, output_tokens, tool_calls ) ) # ── streaming ───────────────────────────── async def event_stream(): response_id = _resp_id("resp") item_id = _resp_id("item") ts = _resp_ts() def sse(event_type: str, data: dict) -> str: """Emit a properly-formed SSE frame with both event: and data: lines. The OpenAI SDK dispatches on the `event:` field — without it most events are silently dropped.""" return f"event: {event_type}\ndata: {json.dumps(data)}\n\n" # 1. response.created yield sse("response.created", { "type": "response.created", "response": { "id": response_id, "object": "response", "created_at": ts, "status": "in_progress", "model": model, "output": [], "usage": None, }, }) # 2. response.in_progress yield sse("response.in_progress", { "type": "response.in_progress", "response": { "id": response_id, "object": "response", "created_at": ts, "status": "in_progress", "model": model, }, }) # ── Run _generate() in the background, pinging every 15 s ────────────── # Without keepalive bytes, Cloudflare (524) and Codex both drop the # connection while the model is thinking or accumulating tool arguments. # SSE comment lines (": ping") are invisible to application code but # reset every proxy's idle-timeout counter. PING_INTERVAL = 15 # seconds gen_task: asyncio.Task = asyncio.ensure_future(_generate()) while not gen_task.done(): try: await asyncio.wait_for(asyncio.shield(gen_task), timeout=PING_INTERVAL) except asyncio.TimeoutError: yield ": ping\n\n" except Exception: break # real error — handled below try: text, tool_calls, input_tokens, output_tokens = gen_task.result() except HTTPException as exc: yield sse("response.failed", { "type": "response.failed", "response": { "id": response_id, "object": "response", "created_at": ts, "status": "failed", "model": chosen_model, "error": {"code": "upstream_error", "message": exc.detail}, }, }) yield "data: [DONE]\n\n" return except Exception as exc: yield sse("response.failed", { "type": "response.failed", "response": { "id": response_id, "object": "response", "created_at": ts, "status": "failed", "model": chosen_model, "error": {"code": "upstream_error", "message": str(exc)}, }, }) yield "data: [DONE]\n\n" return output_index = 0 # ── text output (only emitted if there is text content) ────────────── if text: yield sse("response.output_item.added", { "type": "response.output_item.added", "response_id": response_id, "output_index": output_index, "item": {"id": item_id, "type": "message", "role": "assistant", "status": "in_progress", "content": []}, }) yield sse("response.content_part.added", { "type": "response.content_part.added", "response_id": response_id, "item_id": item_id, "output_index": output_index, "content_index": 0, "part": {"type": "output_text", "text": "", "annotations": []}, }) chunk_size = 64 for i in range(0, len(text), chunk_size): yield sse("response.output_text.delta", { "type": "response.output_text.delta", "response_id": response_id, "item_id": item_id, "output_index": output_index, "content_index": 0, "delta": text[i : i + chunk_size], }) yield sse("response.output_text.done", { "type": "response.output_text.done", "response_id": response_id, "item_id": item_id, "output_index": output_index, "content_index": 0, "text": text, }) yield sse("response.content_part.done", { "type": "response.content_part.done", "response_id": response_id, "item_id": item_id, "output_index": output_index, "content_index": 0, "part": {"type": "output_text", "text": text, "annotations": []}, }) yield sse("response.output_item.done", { "type": "response.output_item.done", "response_id": response_id, "output_index": output_index, "item": {"id": item_id, "type": "message", "role": "assistant", "status": "completed", "content": [{"type": "output_text", "text": text, "annotations": []}]}, }) output_index += 1 # ── tool call outputs (one item per call) ───────────────────────────── for tc in (tool_calls or []): fn = tc.get("function", {}) tc_id = tc.get("id", _resp_id("tool")) tc_item = { "id": tc_id, "type": "function_call", "call_id": tc_id, "name": fn.get("name", ""), "arguments": fn.get("arguments", "{}"), "status": "completed", } yield sse("response.output_item.added", { "type": "response.output_item.added", "response_id": response_id, "output_index": output_index, "item": {**tc_item, "status": "in_progress"}, }) yield sse("response.function_call_arguments.delta", { "type": "response.function_call_arguments.delta", "response_id": response_id, "item_id": tc_id, "output_index": output_index, "call_id": tc_id, "delta": fn.get("arguments", "{}"), }) yield sse("response.function_call_arguments.done", { "type": "response.function_call_arguments.done", "response_id": response_id, "item_id": tc_id, "output_index": output_index, "call_id": tc_id, "arguments": fn.get("arguments", "{}"), }) yield sse("response.output_item.done", { "type": "response.output_item.done", "response_id": response_id, "output_index": output_index, "item": tc_item, }) output_index += 1 # ── response.completed ──────────────────────────────────────────────── yield sse("response.completed", { "type": "response.completed", "response": _build_responses_payload( chosen_model, text, response_id, input_tokens, output_tokens, tool_calls ), }) yield "data: [DONE]\n\n" return StreamingResponse( event_stream(), media_type="text/event-stream", headers={"Cache-Control": "no-cache", "Connection": "keep-alive", "X-Accel-Buffering": "no"}, )