app.py

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 "<html><body></body></html>"

@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)