main.py

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.post("/v1/chat/completions")
async def chat_completion(request: ChatRequest):
    model_id = MODEL_ALIASES.get(request.model, request.model)

    headers = {
        'accept': 'text/event-stream',
        'content-type': 'application/json',
        'origin': 'https://www.chatwithmono.xyz',
        'referer': 'https://www.chatwithmono.xyz/',
        'user-agent': 'Mozilla/5.0',
    }

    payload = {
        "messages": [{"role": msg.role, "content": msg.content} for msg in request.messages],
        "model": model_id
    }

    if request.stream:
        async def event_stream():
            chat_id = f"chatcmpl-{unix_id()}"
            created = int(time.time())
            sent_done = False

            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:
                                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"
        return StreamingResponse(event_stream(), media_type="text/event-stream")

    else:
        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),
            }
        })


# ==== 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
    }