import os import time import logging import asyncio import httpx from fastapi import FastAPI, HTTPException, Header, Depends from fastapi.responses import HTMLResponse from pydantic import BaseModel from typing import List, Optional # --- CUSTOM LOG FILTER TO BLOCK BOT TRAFFIC FROM THE CONSOLE --- class BotLogFilter(logging.Filter): def filter(self, record: logging.LogRecord) -> bool: msg = record.getMessage() if "GET /v1/" in msg or "POST /v1/" in msg or "Started server" in msg or "Application" in msg or "Uvicorn running" in msg: return True if "HTTP" in msg: return False return True logging.getLogger("uvicorn.access").addFilter(BotLogFilter()) # --------------------------------------------------------------- app = FastAPI(title="GenAI Multi-Proxy Server (Rate-Limit Protected)") AGNES_API_KEY = os.getenv("AGNES_API_KEY") REQUIRED_STATIC_TOKEN = "sk-apikeyyoudummy" # Global timeout config TIMEOUT_CONFIG = httpx.Timeout(320.0, connect=30.0, read=300.0, write=30.0) # Video Rate Limiter Lock & Timestamp video_lock = asyncio.Lock() last_video_time = 0.0 class ChatMessage(BaseModel): role: str content: str class ChatCompletionRequest(BaseModel): model: str = "agnes-image-2.1-flash-1k" messages: List[ChatMessage] temperature: Optional[float] = 1.0 stream: Optional[bool] = False async def verify_static_token(authorization: Optional[str] = Header(None)) -> str: if not authorization or not authorization.startswith("Bearer "): raise HTTPException(status_code=401, detail="Missing or malformed Authorization header.") token = authorization.split(" ")[1] if token != REQUIRED_STATIC_TOKEN: raise HTTPException(status_code=403, detail="Invalid API Key provided.") if not AGNES_API_KEY: raise HTTPException(status_code=500, detail="Server misconfiguration: AGNES_API_KEY env variable not set.") return AGNES_API_KEY # --- IMAGE HANDLER WITH 2K -> 1K FALLBACK --- async def handle_image_generation(prompt: str, target_size: str, upstream_key: str) -> str: headers = {"Authorization": f"Bearer {upstream_key}", "Content-Type": "application/json"} current_size = target_size max_retries = 3 async with httpx.AsyncClient(timeout=TIMEOUT_CONFIG) as client: for attempt in range(max_retries): payload = { "model": "agnes-image-2.1-flash", "prompt": prompt, "size": current_size, "ratio": "1:1", "extra_body": {"response_format": "url"} } try: print(f"[PROXY LOG] -> Sending image request ({current_size}) [Attempt {attempt + 1}/{max_retries}]...") res = await client.post("https://apihub.agnes-ai.com/v1/images/generations", json=payload, headers=headers) # Check for Rate Limit 429 if res.status_code == 429: if current_size == "2K": print(f"[PROXY LOG] !!! 2K tier rate limit hit! Automatically falling back to 1K tier...") current_size = "1K" await asyncio.sleep(0.5) continue # Retry immediately using 1K payload else: print(f"[PROXY LOG] !!! Agnes Rate Limit Hit on 1K tier: {res.text}") raise HTTPException(status_code=429, detail="Upstream rate limit reached. Please wait a moment before trying again.") # Upstream busy or temporary error retries if res.status_code in [502, 503, 504] or "upstream_error" in res.text or "do_request_failed" in res.text: print(f"[PROXY LOG] !!! Upstream busy or network hiccup. Retrying in 1.5 seconds...") await asyncio.sleep(1.5) continue if res.status_code != 200: print(f"[PROXY LOG] !!! Agnes API Error: {res.text}") raise HTTPException(status_code=res.status_code, detail=f"Agnes Image Error: {res.text}") img_url = res.json()["data"][0]["url"] print(f"[PROXY LOG] -> Agnes image ready ({current_size}). Returning direct URL: {img_url}") return img_url except httpx.RequestError as exc: print(f"[PROXY LOG] !!! Connection exception on attempt {attempt + 1}: {exc}") if attempt == max_retries - 1: raise HTTPException(status_code=503, detail="Agnes AI gateway completely unreachable.") await asyncio.sleep(1.5) raise HTTPException(status_code=503, detail="Agnes AI backend is currently overloaded. Please try again.") # --- RATE-LIMITED VIDEO HANDLER --- async def handle_video_generation(prompt: str, upstream_key: str) -> str: global last_video_time headers = {"Authorization": f"Bearer {upstream_key}", "Content-Type": "application/json"} payload = { "model": "agnes-video-v2.0", "prompt": prompt, "width": 832, "height": 448, "num_frames": 81, "frame_rate": 24 } # Enforce 1 request per minute queue locally async with video_lock: elapsed = time.time() - last_video_time if elapsed < 60.0: wait_time = 60.0 - elapsed print(f"[PROXY LOG] -> Video rate limit throttle: queueing request for {wait_time:.1f}s...") await asyncio.sleep(wait_time) max_retries = 3 async with httpx.AsyncClient(timeout=TIMEOUT_CONFIG) as client: for attempt in range(max_retries): try: print(f"[PROXY LOG] -> Submitting video task to Agnes [Attempt {attempt + 1}/{max_retries}]...") res = await client.post("https://apihub.agnes-ai.com/v1/videos", json=payload, headers=headers) if res.status_code == 429: print(f"[PROXY LOG] !!! Agnes Video Rate Limit Hit: {res.text}") raise HTTPException(status_code=429, detail="Video engine rate limit reached (1 request/min). Please try again shortly.") if res.status_code in [502, 503, 504] or "upstream_error" in res.text or "do_request_failed" in res.text: print(f"[PROXY LOG] !!! Agnes busy or error. Retrying task submission in 1.5s...") await asyncio.sleep(1.5) continue if res.status_code != 200: print(f"[PROXY LOG] !!! Agnes Task Creation Error: {res.text}") raise HTTPException(status_code=res.status_code, detail=f"Agnes API Error: {res.text}") task_data = res.json() video_id = task_data.get("video_id") # Update successful generation timestamp last_video_time = time.time() tracking_url = f"https://apihub.agnes-ai.com/agnesapi?video_id={video_id}" print(f"[PROXY LOG] -> Task successfully queued. Returning tracking URL immediately: {tracking_url}") return tracking_url except httpx.RequestError as exc: print(f"[PROXY LOG] !!! Network exception during task creation: {exc}") if attempt == max_retries - 1: raise HTTPException(status_code=503, detail="Agnes video creation gateway completely unreachable.") await asyncio.sleep(1.5) raise HTTPException(status_code=503, detail="Agnes backend is currently overloaded. Please try again.") # --- ROUTES --- @app.get("/", response_class=HTMLResponse) async def read_root(): return "" @app.get("/v1/models") async def list_models(): now = int(time.time()) return { "object": "list", "data": [ {"id": "agnes-image-2.1-flash-1k", "object": "model", "created": now, "owned_by": "custom-proxy"}, {"id": "agnes-image-2.1-flash-2k", "object": "model", "created": now, "owned_by": "custom-proxy"}, {"id": "agnes-video-v2.0", "object": "model", "created": now, "owned_by": "custom-proxy"} ] } @app.post("/v1/chat/completions") async def chat_completions_proxy(request: ChatCompletionRequest, upstream_key: str = Depends(verify_static_token)): print(f"[PROXY LOG] === Incoming Request Received ===") if not request.messages: raise HTTPException(status_code=400, detail="No messages provided.") user_prompt = request.messages[-1].content model_name = request.model.lower() if request.model else "agnes-image-2.1-flash-1k" # Routing matching target model selection if "video" in model_name: final_url = await handle_video_generation(user_prompt, upstream_key) else: target_size = "2K" if "-2k" in model_name else "1K" final_url = await handle_image_generation(user_prompt, target_size, upstream_key) print(f"[PROXY LOG] === Request Successfully Completed ===") prompt_tokens = 20000 completion_tokens = 40000 total_tokens = prompt_tokens + completion_tokens return { "id": f"chatcmpl-{int(time.time())}", "object": "chat.completion", "created": int(time.time()), "model": request.model, "choices": [ { "index": 0, "message": { "role": "assistant", "content": final_url }, "finish_reason": "stop" } ], "usage": { "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, "total_tokens": total_tokens } } if __name__ == "__main__": import uvicorn uvicorn.run("app:app", host="0.0.0.0", port=7860, reload=False)