fix length

main.py

@@ -6,6 +6,7 @@ from fastapi.responses import StreamingResponse
 from fastapi.security import APIKeyHeader, APIKeyQuery
 from itertools import cycle
 import asyncio
+import json
 
 # --- Configuration ---
 PROXY_API_KEY = os.environ.get("PROXY_API_KEY")
@@ -71,13 +72,19 @@ async def proxy(request: Request, model_path: str, api_key: str = Security(get_a
         express_key = next(key_rotator)
     
     project_id = await get_project_id(express_key)
-    
-    request_body = await request.json()
+
+    raw_request_body = await request.body()
+    request_body_to_send = raw_request_body
 
     # Model-specific request body modification
     if "gemini-2.5-flash-image-preview" in model_path:
-        if "generationConfig" in request_body and "thinkingConfig" in request_body.get("generationConfig", {}):
-            del request_body["generationConfig"]["thinkingConfig"]
+        try:
+            request_json = json.loads(raw_request_body)
+            if "generationConfig" in request_json and "thinkingConfig" in request_json.get("generationConfig", {}):
+                del request_json["generationConfig"]["thinkingConfig"]
+                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/us-central1/publishers/google/models/{model_path}?key={express_key}"
 
@@ -86,19 +93,27 @@ async def proxy(request: Request, model_path: str, api_key: str = Security(get_a
             yield chunk
 
     async with httpx.AsyncClient(timeout=None) as client:
+        # Exclude host, auth, and content-length headers from the original request
+        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']
+        }
+
         req = client.build_request(
             method=request.method,
             url=target_url,
-            headers={k: v for k, v in request.headers.items() if k.lower() not in ['host', 'authorization', 'x-goog-api-key']},
-            json=request_body,
+            headers=headers_to_proxy,
+            content=request_body_to_send,
         )
         response = await client.send(req, stream=True)
 
     if "streamGenerateContent" in model_path:
         return StreamingResponse(stream_response(response), media_type=response.headers.get("content-type"))
     else:
-        response_data = await response.aread()
-        return response_data
+        # For non-streaming, read the response and return it.
+        # httpx.Response can be returned directly in FastAPI
+        await response.aread()
+        return response
 
 if __name__ == "__main__":
     import uvicorn