frontend vertex key support

index.html

@@ -173,8 +173,8 @@
         <h1>🎨 Gemini Image Generator</h1>
         
         <div class="input-group">
-            <label for="apiKey">API Key:</label>
-            <input type="text" id="apiKey" placeholder="Enter your API key">
+            <label for="vertexKey">Vertex Express Key:</label>
+            <input type="text" id="vertexKey" placeholder="Enter your Vertex Express key">
         </div>
         
         <div class="input-group">
@@ -299,11 +299,11 @@
         }
         
         async function generateImage() {
-            const apiKey = document.getElementById('apiKey').value.trim();
+            const vertexKey = document.getElementById('vertexKey').value.trim();
             const prompt = document.getElementById('prompt').value.trim();
             
-            if (!apiKey) {
-                showError('Please enter your API key');
+            if (!vertexKey) {
+                showError('Please enter your Vertex Express key');
                 return;
             }
             
@@ -346,11 +346,11 @@
                     return;
                 }
                 
-                const response = await fetch('/v1beta/models/gemini-2.5-flash-image-preview:generateContent', {
+                const response = await fetch('/frontend/v1beta/models/gemini-2.5-flash-image-preview:generateContent', {
                     method: 'POST',
                     headers: {
                         'Content-Type': 'application/json',
-                        'x-goog-api-key': apiKey
+                        'x-vertex-express-key': vertexKey
                     },
                     body: JSON.stringify({
                         contents: [{

main.py

@@ -1,9 +1,10 @@
 import os
 import re
 import httpx
-from fastapi import FastAPI, Request, HTTPException, Security
+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
@@ -23,6 +24,7 @@ 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)
@@ -82,33 +84,35 @@ async def gif_worker():
     worker_path = current_dir / "gif.worker.js"
     return FileResponse(worker_path, media_type="application/javascript")
 
-# --- 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)
-
+# --- 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")
+            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"]
-                request_json["generationConfig"]["responseModalities"] = ["TEXT", "IMAGE"]
+            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"] = [
             {
@@ -146,7 +150,7 @@ async def proxy(request: Request, model_path: str, api_key: str = Security(get_a
         ]
         request_body_to_send = json.dumps(request_json).encode('utf-8')
     except json.JSONDecodeError:
-            pass # Not a json body, proxy as is
+        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}"
 
@@ -154,7 +158,7 @@ async def proxy(request: Request, model_path: str, api_key: str = Security(get_a
 
     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']
+        if k.lower() not in ['host', 'authorization', 'x-goog-api-key', 'x-vertex-express-key', 'content-length']
     }
 
     print(request_body_to_send)
@@ -208,12 +212,47 @@ async def proxy(request: Request, model_path: str, api_key: str = Security(get_a
             return Response(
                 content=modified_response_data,
                 status_code=response.status_code,
-                headers={"content-type":response.headers.get("content-type")},
+                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