import os import re import httpx from fastapi import FastAPI, Request, HTTPException, Security from fastapi.responses import StreamingResponse, Response from fastapi.security import APIKeyHeader, APIKeyQuery from itertools import cycle import asyncio import json # --- 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() # --- 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.") # --- Proxy Endpoint --- @app.post("/v1beta/models/{model_path:path}") async def proxy(request: Request, model_path: str, api_key: str = Security(get_api_key)): async with key_lock: express_key = next(key_rotator) project_id = await get_project_id(express_key) 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"] print(request_json["generationConfig"]) if "generationConfig" in request_json and "responseMimeType" in request_json.get("generationConfig", {}): del request_json["generationConfig"]["responseMimeType"] request_json["generationConfig"] request_json["generationConfig"]["responseModalities"] = ["TEXT", "IMAGE"] # 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', '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() if __name__ == "__main__": import uvicorn # Hugging Face Spaces run on port 7860 uvicorn.run(app, host="0.0.0.0", port=7860)