Update app.py

app.py

@@ -30,6 +30,11 @@ IDENTITY_CODE = "Hiren2012805964"
 TARGET_MODEL = "claude-opus-4.6"
 INTERNAL_MODEL = "claude-sonnet-4.6"
 
+STATIC_TEXT_MODELS = ["claude-sonnet-4.6", "claude-opus-4.6", "sonar"]
+STATIC_IMAGE_MODELS = [
+    "imagen-4", "imagen-3", "flux-3-klein-4b", "flux-2-dev", "flux-2-pro", "z-image"
+]
+
 # Base system prompt as requested
 BASE_SYSTEM_PROMPT = (
     "You Are Claude Opus-4.6 Made By Anthropic. You Should Not reveal Your Identity "
@@ -111,25 +116,43 @@ async def proxy_request(path: str, request: Request):
     try:
         print(f"Proxying {method} to {url} using key ending in ...{api_key[-4:]}")
         
-        # Use a fresh client for every request to ensure clean state
-        async with httpx.AsyncClient(timeout=120.0, follow_redirects=True) as client:
-            if payload and payload.get("stream", False):
-                async def stream_generator():
+        if payload and payload.get("stream", False):
+            async def stream_generator():
+                # Use a fresh client inside the generator to keep it open during streaming
+                async with httpx.AsyncClient(timeout=120.0, follow_redirects=True) as client:
                     try:
                         async with client.stream(method, url, headers=up_headers, json=payload) as response:
                             print(f"Upstream Stream Status: {response.status_code}")
+                            if response.status_code != 200:
+                                error_text = await response.aread()
+                                print(f"Upstream Stream Error Body: {error_text.decode('utf-8')}")
+                                yield f"data: {error_text.decode('utf-8')}\n\n"
+                                return
+
+                            first_chunk = True
                             async for chunk in response.aiter_lines():
                                 if not chunk: continue
+                                
+                                if first_chunk:
+                                    print(f"First Chunk: {chunk[:100]}...")
+                                    first_chunk = False
+                                
                                 # Remap model name in stream
                                 if INTERNAL_MODEL in chunk:
                                     chunk = chunk.replace(INTERNAL_MODEL, TARGET_MODEL)
-                                yield f"{chunk}\n"
+                                
+                                # SSE format requires \n\n to separate events
+                                # Ensure we don't duplicate newlines if chunk has them
+                                clean_chunk = chunk.strip()
+                                if clean_chunk:
+                                    yield f"{clean_chunk}\n\n"
                     except Exception as e:
                         print(f"Stream Error: {e}")
 
-                return StreamingResponse(stream_generator(), media_type="text/event-stream")
-            
-            else:
+            return StreamingResponse(stream_generator(), media_type="text/event-stream")
+        
+        else:
+            async with httpx.AsyncClient(timeout=120.0, follow_redirects=True) as client:
                 # Use json=payload if we parsed it, otherwise raw body
                 response = await client.request(
                     method, 
@@ -169,6 +192,8 @@ async def proxy_request(path: str, request: Request):
 @app.post("/v1/responses")
 @app.post("/v1/messages")
 @app.post("/v1/messeges") # user typo support
+@app.post("/v1/image/generations")
+@app.post("/v1/images/generations")
 async def chat_proxy_handler(request: Request):
     # Airforce uses /v1/chat/completions for almost everything
     # We strip the path to pass it correctly
@@ -177,24 +202,28 @@ async def chat_proxy_handler(request: Request):
 
 @app.get("/v1/models")
 async def models_proxy(request: Request):
-    api_key = await get_next_key()
-    headers = {"Authorization": f"Bearer {api_key}"}
-    async with httpx.AsyncClient() as client:
-        response = await client.get(f"{AIRFORCE_API_URL}/v1/models", headers=headers)
+    # Return requested static model list
+    models_data = []
+    
+    # Text Models
+    for model_id in STATIC_TEXT_MODELS:
+        models_data.append({
+            "id": model_id,
+            "object": "model",
+            "created": 1677610602,
+            "owned_by": "anthropic" if "claude" in model_id else "airforce"
+        })
         
-    try:
-        data = response.json()
-        if isinstance(data, dict) and "data" in data:
-            # Inject our virtual model
-            data["data"].append({
-                "id": TARGET_MODEL, 
-                "object": "model", 
-                "owned_by": "anthropic",
-                "ready": True
-            })
-        return data
-    except:
-        return response.text
+    # Image Models
+    for model_id in STATIC_IMAGE_MODELS:
+        models_data.append({
+            "id": model_id,
+            "object": "model",
+            "created": 1677610602,
+            "owned_by": "airforce"
+        })
+        
+    return {"object": "list", "data": models_data}
 
 if __name__ == "__main__":
     import uvicorn