Update main.py

main.py

@@ -1,32 +1,57 @@
-from fastapi import FastAPI
-from fastapi.responses import StreamingResponse, JSONResponse
-from pydantic import BaseModel
-from typing import List, Optional
-import time
+import base64
 import json
 import os
-import base64
+import time
+from typing import List, Optional
+
 import httpx
 from dotenv import load_dotenv
+from fastapi import FastAPI
+from fastapi.responses import JSONResponse, StreamingResponse
+from pydantic import BaseModel
 
-load_dotenv()
+# --- Configuration ---
 
-# Dummy imports for example
-from models import AVAILABLE_MODELS, MODEL_ALIASES  # Ensure these are defined properly
+load_dotenv()
 
-# Env variables
-IMAGE_API_URL = os.environ["IMAGE_API_URL"]
+# Env variables for external services
+IMAGE_API_URL = os.environ.get("IMAGE_API_URL", "https://image.api.example.com") # Add a default for safety
 SNAPZION_UPLOAD_URL = "https://upload.snapzion.com/api/public-upload"
-SNAPZION_API_KEY = os.environ["SNAP"]
+SNAPZION_API_KEY = os.environ.get("SNAP", "") # Add a default for safety
 
-app = FastAPI()
+# --- Dummy Model Definitions ---
+# In a real application, these would be defined properly.
+# For this example, we define them here.
+
+AVAILABLE_MODELS = [
+    {"id": "gpt-4-turbo", "object": "model", "created": int(time.time()), "owned_by": "system"},
+    {"id": "gpt-4o", "object": "model", "created": int(time.time()), "owned_by": "system"},
+    {"id": "gpt-3.5-turbo", "object": "model", "created": int(time.time()), "owned_by": "system"},
+    {"id": "dall-e-3", "object": "model", "created": int(time.time()), "owned_by": "system"},
+    # Add any other models you support
+]
+
+MODEL_ALIASES = {
+    # Example: "gpt-4": "gpt-4-turbo"
+}
+
+# --- FastAPI Application ---
+
+app = FastAPI(
+    title="OpenAI Compatible API",
+    description="An adapter for various services to be compatible with the OpenAI API specification.",
+    version="1.0.0"
+)
 
 def unix_id():
+    """Generates a Unix timestamp-based ID."""
     return str(int(time.time() * 1000))
 
-# === Models ===
+# === API Endpoints ===
+
 @app.get("/v1/models")
 async def list_models():
+    """Lists the available models."""
     return {"object": "list", "data": AVAILABLE_MODELS}
 
 # === Chat Completion ===
@@ -39,9 +64,18 @@ class ChatRequest(BaseModel):
     messages: List[Message]
     model: str
     stream: Optional[bool] = False
+    # Add other common parameters for compatibility if needed
+    # max_tokens: Optional[int] = None
+    # temperature: Optional[float] = None
+    # user: Optional[str] = None
+
 
 @app.post("/v1/chat/completions")
 async def chat_completion(request: ChatRequest):
+    """
+    Handles chat completion requests, supporting both streaming and non-streaming responses.
+    This endpoint is updated to match the OpenAI streaming chunk format precisely.
+    """
     model_id = MODEL_ALIASES.get(request.model, request.model)
 
     headers = {
@@ -61,82 +95,140 @@ async def chat_completion(request: ChatRequest):
         async def event_stream():
             chat_id = f"chatcmpl-{unix_id()}"
             created = int(time.time())
-            sent_done = False
+            
+            is_first_chunk = True
+            usage_info = None
 
-            async with httpx.AsyncClient(timeout=120) as client:
-                async with client.stream("POST", "https://www.chatwithmono.xyz/api/chat", headers=headers, json=payload) as response:
-                    async for line in response.aiter_lines():
-                        if line.startswith("0:"):
-                            try:
-                                content_piece = json.loads(line[2:])
-                                chunk_data = {
-                                    "id": chat_id,
-                                    "object": "chat.completion.chunk",
-                                    "created": created,
-                                    "model": model_id,
-                                    "choices": [{
-                                        "delta": {"content": content_piece},
-                                        "index": 0,
-                                        "finish_reason": None
-                                    }]
-                                }
-                                yield f"data: {json.dumps(chunk_data)}\n\n"
-                            except:
+            try:
+                async with httpx.AsyncClient(timeout=120) as client:
+                    async with client.stream("POST", "https://www.chatwithmono.xyz/api/chat", headers=headers, json=payload) as response:
+                        response.raise_for_status()
+                        
+                        async for line in response.aiter_lines():
+                            if not line:
                                 continue
-                        elif line.startswith(("e:", "d:")) and not sent_done:
-                            sent_done = True
-                            done_chunk = {
-                                "id": chat_id,
-                                "object": "chat.completion.chunk",
-                                "created": created,
-                                "model": model_id,
-                                "choices": [{
-                                    "delta": {},
-                                    "index": 0,
-                                    "finish_reason": "stop"
-                                }]
-                            }
-                            yield f"data: {json.dumps(done_chunk)}\n\ndata: [DONE]\n\n"
+
+                            if line.startswith("0:"):
+                                try:
+                                    content_piece = json.loads(line[2:])
+                                    delta = {
+                                        "content": content_piece,
+                                        "function_call": None,
+                                        "tool_calls": None
+                                    }
+                                    if is_first_chunk:
+                                        delta["role"] = "assistant"
+                                        is_first_chunk = False
+
+                                    chunk_data = {
+                                        "id": chat_id,
+                                        "object": "chat.completion.chunk",
+                                        "created": created,
+                                        "model": model_id,
+                                        "choices": [{"index": 0, "delta": delta, "finish_reason": None}],
+                                        "usage": None
+                                    }
+                                    yield f"data: {json.dumps(chunk_data)}\n\n"
+                                except json.JSONDecodeError:
+                                    continue
+                            
+                            elif line.startswith(("e:", "d:")):
+                                try:
+                                    end_data = json.loads(line[2:])
+                                    usage_info = end_data.get("usage")
+                                except (json.JSONDecodeError, AttributeError):
+                                    pass
+                                break
+                
+                # After the loop, send the final chunk with finish_reason and usage
+                final_usage = None
+                if usage_info:
+                    prompt_tokens = usage_info.get("promptTokens", 0)
+                    completion_tokens = usage_info.get("completionTokens", 0)
+                    final_usage = {
+                        "prompt_tokens": prompt_tokens,
+                        "completion_tokens": completion_tokens,
+                        "total_tokens": prompt_tokens + completion_tokens,
+                    }
+
+                done_chunk = {
+                    "id": chat_id,
+                    "object": "chat.completion.chunk",
+                    "created": created,
+                    "model": model_id,
+                    "choices": [{
+                        "index": 0,
+                        "delta": {"role": "assistant", "content": None, "function_call": None, "tool_calls": None},
+                        "finish_reason": "stop"
+                    }],
+                    "usage": final_usage
+                }
+                yield f"data: {json.dumps(done_chunk)}\n\n"
+
+            except httpx.HTTPStatusError as e:
+                error_content = {
+                    "error": {
+                        "message": f"Upstream API error: {e.response.status_code}. Details: {e.response.text}",
+                        "type": "upstream_error",
+                        "code": str(e.response.status_code)
+                    }
+                }
+                yield f"data: {json.dumps(error_content)}\n\n"
+            
+            finally:
+                yield "data: [DONE]\n\n"
+
         return StreamingResponse(event_stream(), media_type="text/event-stream")
     else:
+        # Non-streaming logic remains the same
         assistant_response = ""
         usage_info = {}
 
-        async with httpx.AsyncClient(timeout=120) as client:
-            async with client.stream("POST", "https://www.chatwithmono.xyz/api/chat", headers=headers, json=payload) as response:
-                async for chunk in response.aiter_lines():
-                    if chunk.startswith("0:"):
-                        try:
-                            piece = json.loads(chunk[2:])
-                            assistant_response += piece
-                        except:
-                            continue
-                    elif chunk.startswith(("e:", "d:")):
-                        try:
-                            data = json.loads(chunk[2:])
-                            usage_info = data.get("usage", {})
-                        except:
-                            continue
-
-        return JSONResponse(content={
-            "id": f"chatcmpl-{unix_id()}",
-            "object": "chat.completion",
-            "created": int(time.time()),
-            "model": model_id,
-            "choices": [{
-                "index": 0,
-                "message": {
-                    "role": "assistant",
-                    "content": assistant_response
-                },
-                "finish_reason": "stop"
-            }],
-            "usage": {
-                "prompt_tokens": usage_info.get("promptTokens", 0),
-                "completion_tokens": usage_info.get("completionTokens", 0),
-                "total_tokens": usage_info.get("promptTokens", 0) + usage_info.get("completionTokens", 0),
-            }
-        })
+        try:
+            async with httpx.AsyncClient(timeout=120) as client:
+                async with client.stream("POST", "https://www.chatwithmono.xyz/api/chat", headers=headers, json=payload) as response:
+                    response.raise_for_status()
+                    async for chunk in response.aiter_lines():
+                        if chunk.startswith("0:"):
+                            try:
+                                piece = json.loads(chunk[2:])
+                                assistant_response += piece
+                            except:
+                                continue
+                        elif chunk.startswith(("e:", "d:")):
+                            try:
+                                data = json.loads(chunk[2:])
+                                usage_info = data.get("usage", {})
+                            except:
+                                continue
+            
+            return JSONResponse(content={
+                "id": f"chatcmpl-{unix_id()}",
+                "object": "chat.completion",
+                "created": int(time.time()),
+                "model": model_id,
+                "choices": [{
+                    "index": 0,
+                    "message": {"role": "assistant", "content": assistant_response},
+                    "finish_reason": "stop"
+                }],
+                "usage": {
+                    "prompt_tokens": usage_info.get("promptTokens", 0),
+                    "completion_tokens": usage_info.get("completionTokens", 0),
+                    "total_tokens": usage_info.get("promptTokens", 0) + usage_info.get("completionTokens", 0),
+                }
+            })
+        except httpx.HTTPStatusError as e:
+            return JSONResponse(
+                status_code=e.response.status_code,
+                content={
+                    "error": {
+                        "message": f"Upstream API error. Details: {e.response.text}",
+                        "type": "upstream_error"
+                    }
+                }
+            )
+
 
 # === Image Generation ===
 
@@ -149,83 +241,58 @@ class ImageGenerationRequest(BaseModel):
 
 @app.post("/v1/images/generations")
 async def generate_images(request: ImageGenerationRequest):
+    """Handles image generation requests."""
     results = []
 
-    async with httpx.AsyncClient(timeout=60) as client:
-        for _ in range(request.n):
-            model = request.model or "default"
-
-            if model in ["gpt-image-1", "dall-e-3", "dall-e-2", "nextlm-image-1"]:
-                headers = {
-                    'Content-Type': 'application/json',
-                    'User-Agent': 'Mozilla/5.0',
-                    'Referer': 'https://www.chatwithmono.xyz/',
-                    'sec-ch-ua-platform': '"Windows"',
-                    'sec-ch-ua': '"Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"',
-                    'sec-ch-ua-mobile': '?0',
-                }
-
-                payload = {
-                    "prompt": request.prompt,
-                    "model": model
-                }
+    try:
+        async with httpx.AsyncClient(timeout=120) as client:
+            for _ in range(request.n):
+                model = request.model or "default"
 
-                resp = await client.post("https://www.chatwithmono.xyz/api/image", headers=headers, json=payload)
-                if resp.status_code != 200:
-                    return JSONResponse(
-                        status_code=502,
-                        content={"error": f"{model} generation failed", "details": resp.text}
-                    )
+                if model in ["gpt-image-1", "dall-e-3", "dall-e-2", "nextlm-image-1"]:
+                    headers = {
+                        'Content-Type': 'application/json',
+                        'User-Agent': 'Mozilla/5.0',
+                        'Referer': 'https://www.chatwithmono.xyz/',
+                    }
+                    payload = {"prompt": request.prompt, "model": model}
+                    resp = await client.post("https://www.chatwithmono.xyz/api/image", headers=headers, json=payload)
+                    resp.raise_for_status()
+                    data = resp.json()
+                    b64_image = data.get("image")
+                    if not b64_image:
+                        return JSONResponse(status_code=502, content={"error": "Missing base64 image in response"})
 
-                data = resp.json()
-                b64_image = data.get("image")
-                if not b64_image:
-                    return JSONResponse(status_code=502, content={"error": "Missing base64 image in response"})
+                    if SNAPZION_API_KEY:
+                        upload_headers = {"Authorization": SNAPZION_API_KEY}
+                        upload_files = {'file': ('image.png', base64.b64decode(b64_image), 'image/png')}
+                        upload_resp = await client.post(SNAPZION_UPLOAD_URL, headers=upload_headers, files=upload_files)
+                        upload_resp.raise_for_status()
+                        upload_data = upload_resp.json()
+                        image_url = upload_data.get("url")
+                    else:
+                        image_url = f"data:image/png;base64,{b64_image}"
 
-                data_uri = f"data:image/png;base64,{b64_image}"
+                    results.append({"url": image_url, "b64_json": b64_image, "revised_prompt": data.get("revised_prompt")})
+                else:
+                    params = {"prompt": request.prompt, "aspect_ratio": request.aspect_ratio, "link": "typegpt.net"}
+                    resp = await client.get(IMAGE_API_URL, params=params)
+                    resp.raise_for_status()
+                    data = resp.json()
+                    results.append({"url": data.get("image_link"), "b64_json": data.get("base64_output")})
+    except httpx.HTTPStatusError as e:
+        return JSONResponse(
+            status_code=502,
+            content={"error": f"Image generation failed. Upstream error: {e.response.status_code}", "details": e.response.text}
+        )
+    except Exception as e:
+        return JSONResponse(status_code=500, content={"error": "An internal error occurred.", "details": str(e)})
 
-                upload_headers = {"Authorization": SNAPZION_API_KEY}
-                upload_files = {
-                    'file': ('image.png', base64.b64decode(b64_image), 'image/png')
-                }
 
-                upload_resp = await client.post(SNAPZION_UPLOAD_URL, headers=upload_headers, files=upload_files)
-                if upload_resp.status_code != 200:
-                    return JSONResponse(
-                        status_code=502,
-                        content={"error": "Upload failed", "details": upload_resp.text}
-                    )
-
-                upload_data = upload_resp.json()
-                results.append({
-                    "url": upload_data.get("url"),
-                    "b64_json": b64_image,
-                    "data_uri": data_uri,
-                    "model": model
-                })
-            else:
-                params = {
-                    "prompt": request.prompt,
-                    "aspect_ratio": request.aspect_ratio,
-                    "link": "typegpt.net"
-                }
+    return {"created": int(time.time()), "data": results}
 
-                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"),
-                    "model": "default"
-                })
-
-    return {
-        "created": int(time.time()),
-        "data": results
-    }
+if __name__ == "__main__":
+    import uvicorn
+    # Make sure you have a .env file with SNAP and IMAGE_API_URL
+    # Example: uvicorn your_script_name:app --reload --port 8000
+    uvicorn.run(app, host="0.0.0.0", port=8000)