import os import re import httpx from fastapi import FastAPI, Request, HTTPException, Security, Header from fastapi.responses import StreamingResponse, Response, FileResponse from fastapi.security import APIKeyHeader, APIKeyQuery import logging from fastapi.staticfiles import StaticFiles from itertools import cycle import asyncio import json from pathlib import Path # --- Configuration --- PROXY_API_KEY = os.environ.get("PROXY_API_KEY") VERTEX_EXPRESS_KEYS_STR = os.environ.get("VERTEX_EXPRESS_KEYS") VERTEX_EXPRESS_KEYS = [key.strip() for key in VERTEX_EXPRESS_KEYS_STR.split(',')] if VERTEX_EXPRESS_KEYS_STR else [] if not VERTEX_EXPRESS_KEYS: raise ValueError("VERTEX_EXPRESS_KEYS environment variable not set or empty.") # --- Globals --- app = FastAPI() project_id_cache = {} key_rotator = cycle(VERTEX_EXPRESS_KEYS) key_lock = asyncio.Lock() logger = logging.getLogger(__name__) # --- API Key Security --- api_key_query = APIKeyQuery(name="key", auto_error=False) api_key_header = APIKeyHeader(name="x-goog-api-key", auto_error=False) async def get_api_key( key_query: str = Security(api_key_query), key_header: str = Security(api_key_header), ): if PROXY_API_KEY: if key_query == PROXY_API_KEY: return key_query if key_header == PROXY_API_KEY: return key_header raise HTTPException(status_code=401, detail="Invalid or missing API Key") else: # If no PROXY_API_KEY is set, authentication is skipped return None # --- Project ID Extraction --- async def get_project_id(key: str): if key in project_id_cache: return project_id_cache[key] url = f"https://aiplatform.googleapis.com/v1/publishers/google/models/gemini-2.6-pro:generateContent?key={key}" headers = {'Content-Type': 'application/json'} data = '{}' async with httpx.AsyncClient() as client: try: response = await client.post(url, headers=headers, data=data) response.raise_for_status() except httpx.HTTPStatusError as e: if e.response.status_code == 404: error_message = e.response.json().get("error", {}).get("message", "") match = re.search(r"projects/([^/]+)/locations/", error_message) if match: project_id = match.group(1) project_id_cache[key] = project_id return project_id raise HTTPException(status_code=500, detail=f"Failed to extract project ID: {e.response.text}") raise HTTPException(status_code=500, detail="Could not extract project ID from any key.") # --- Frontend Routes --- @app.get("/") async def frontend(): # Get the directory where this script is located current_dir = Path(__file__).parent html_path = current_dir / "index.html" return FileResponse(html_path) @app.get("/gif.worker.js") async def gif_worker(): # Serve the GIF worker script current_dir = Path(__file__).parent worker_path = current_dir / "gif.worker.js" return FileResponse(worker_path, media_type="application/javascript") # --- Shared Model Calling Logic --- async def call_model(request: Request, model_path: str, express_key: str, project_id: str): """ Shared function to handle model calling logic for both proxy endpoints. """ raw_request_body = await request.body() request_body_to_send = raw_request_body try: request_json = json.loads(raw_request_body) if "gemini-2.0-flash-exp-image-generation" in model_path: model_path = model_path.replace("gemini-2.0-flash-exp-image-generation", "gemini-2.5-flash-image-preview") if "generationConfig" not in request_json: request_json["generationConfig"] = {} # Model-specific request body modification if "gemini-2.5-flash-image-preview" in model_path: if "generationConfig" in request_json and "thinkingConfig" in request_json.get("generationConfig", {}): del request_json["generationConfig"]["thinkingConfig"] if "generationConfig" in request_json and "responseMimeType" in request_json.get("generationConfig", {}): del request_json["generationConfig"]["responseMimeType"] request_json["generationConfig"]["responseModalities"] = ["TEXT", "IMAGE"] # Ensure contents have role field if "contents" in request_json: for content in request_json["contents"]: if "role" not in content: content["role"] = "user" request_json["safetySettings"] = [ { "category": "HARM_CATEGORY_HATE_SPEECH", "threshold": "BLOCK_NONE" }, { "category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE" }, { "category": "HARM_CATEGORY_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE" }, { "category": "HARM_CATEGORY_HARASSMENT", "threshold": "BLOCK_NONE" }, { "category": "HARM_CATEGORY_IMAGE_HATE", "threshold": "BLOCK_NONE" }, { "category": "HARM_CATEGORY_IMAGE_DANGEROUS_CONTENT", "threshold": "BLOCK_NONE" }, { "category": "HARM_CATEGORY_IMAGE_HARASSMENT", "threshold": "BLOCK_NONE" }, { "category": "HARM_CATEGORY_IMAGE_SEXUALLY_EXPLICIT", "threshold": "BLOCK_NONE" } ] request_body_to_send = json.dumps(request_json).encode('utf-8') except json.JSONDecodeError: pass # Not a json body, proxy as is target_url = f"https://aiplatform.googleapis.com/v1/projects/{project_id}/locations/global/publishers/google/models/{model_path}?key={express_key}" client = httpx.AsyncClient(timeout=None) headers_to_proxy = { k: v for k, v in request.headers.items() if k.lower() not in ['host', 'authorization', 'x-goog-api-key', 'x-vertex-express-key', 'content-length'] } print(request_body_to_send) if "streamGenerateContent" in model_path: target_url = target_url + "&alt=sse" req = client.build_request( method=request.method, url=target_url, headers=headers_to_proxy, content=request_body_to_send, ) response = await client.send(req, stream=True) if response.status_code != 200: try: response_data = await response.aread() return Response( content=response_data, status_code=response.status_code, headers=dict(response.headers), ) finally: await response.aclose() await client.aclose() if "streamGenerateContent" in model_path: async def stream_generator(): try: async for line in response.aiter_lines(): print(line) yield f"{line}\n" finally: await response.aclose() await client.aclose() return StreamingResponse(stream_generator(), media_type=response.headers.get("content-type")) else: try: response_data = await response.aread() response_json = json.loads(response_data) if 'candidates' in response_json: for candidate in response_json.get('candidates', []): if 'content' in candidate and 'parts' in candidate.get('content', {}): candidate['content']['parts'] = [part for part in candidate['content']['parts'] if part] modified_response_data = json.dumps(response_json).encode('utf-8') return Response( content=modified_response_data, status_code=response.status_code, headers={"content-type": response.headers.get("content-type")}, ) finally: await response.aclose() await client.aclose() # --- Frontend-specific endpoint (no authentication required) --- @app.post("/frontend/v1beta/models/{model_name}:{function_name}") async def frontend_proxy( model_name: str, function_name: str, request: Request, vertex_express_key: str = Header(..., alias="x-vertex-express-key") ): """ Frontend-specific proxy endpoint that only requires a Vertex Express key. No proxy authentication needed. """ try: # Get or extract project ID for this key project_id = await get_project_id(vertex_express_key) # Use shared model calling logic model_path = f"{model_name}:{function_name}" return await call_model(request, model_path, vertex_express_key, project_id) except Exception as e: logger.error(f"Frontend proxy error: {str(e)}") raise HTTPException(status_code=500, detail=f"Proxy error: {str(e)}") # --- Proxy Endpoint --- @app.post("/v1beta/models/{model_path:path}") async def proxy(request: Request, model_path: str, _: str = Security(get_api_key)): async with key_lock: express_key = next(key_rotator) project_id = await get_project_id(express_key) # Use shared model calling logic return await call_model(request, model_path, express_key, project_id) if __name__ == "__main__": import uvicorn # Hugging Face Spaces run on port 7860 uvicorn.run(app, host="0.0.0.0", port=7860)