Update main.py

main.py

@@ -1,29 +1,46 @@
-from fastapi import FastAPI, Request
+from fastapi import FastAPI
 from fastapi.responses import StreamingResponse, JSONResponse
 from pydantic import BaseModel
 from typing import List, Optional
 import time
 import json
+import os
 import httpx
+from dotenv import load_dotenv
+
+load_dotenv()
+
 from models import AVAILABLE_MODELS, MODEL_ALIASES
 
+# Require IMAGE_API_URL strictly from environment
+IMAGE_API_URL = os.environ["IMAGE_API_URL"]
+
 app = FastAPI()
 
+
 def unix_id():
     return str(int(time.time() * 1000))
 
+
+# ==== Models ====
+
+@app.get("/v1/models")
+async def list_models():
+    return {"object": "list", "data": AVAILABLE_MODELS}
+
+
+# ==== Chat Completion ====
+
 class Message(BaseModel):
     role: str
     content: str
 
+
 class ChatRequest(BaseModel):
     messages: List[Message]
     model: str
     stream: Optional[bool] = False
 
-@app.get("/v1/models")
-async def list_models():
-    return {"object": "list", "data": AVAILABLE_MODELS}
 
 @app.post("/v1/chat/completions")
 async def chat_completion(request: ChatRequest):
@@ -34,7 +51,7 @@ async def chat_completion(request: ChatRequest):
         'content-type': 'application/json',
         'origin': 'https://www.chatwithmono.xyz',
         'referer': 'https://www.chatwithmono.xyz/',
-        'user-agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/137.0.0.0 Safari/537.36',
+        'user-agent': 'Mozilla/5.0',
     }
 
     payload = {
@@ -123,3 +140,45 @@ async def chat_completion(request: ChatRequest):
                 "total_tokens": usage_info.get("promptTokens", 0) + usage_info.get("completionTokens", 0),
             }
         })
+
+
+# ==== Image Generation ====
+
+class ImageGenerationRequest(BaseModel):
+    prompt: str
+    aspect_ratio: Optional[str] = "1:1"
+    n: Optional[int] = 1
+    user: Optional[str] = None
+
+
+@app.post("/v1/images/generations")
+async def generate_images(request: ImageGenerationRequest):
+    results = []
+
+    async with httpx.AsyncClient(timeout=60) as client:
+        for _ in range(request.n):
+            params = {
+                "prompt": request.prompt,
+                "aspect_ratio": request.aspect_ratio,
+                "link": "typegpt.net"
+            }
+
+            resp = await client.get(IMAGE_API_URL, params=params)
+
+            if resp.status_code != 200:
+                return JSONResponse(
+                    status_code=502,
+                    content={"error": "Image generation failed", "details": resp.text}
+                )
+
+            data = resp.json()
+            results.append({
+                "url": data.get("image_link"),
+                "b64_json": data.get("base64_output"),
+                "retries": data.get("attempt")
+            })
+
+    return {
+        "created": int(time.time()),
+        "data": results
+    }