from fastapi import FastAPI
from pydantic import BaseModel
from g4f.client import Client
import uvicorn

# Inisialisasi FastAPI
app = FastAPI()

# Model untuk menerima payload untuk chat
class MessageRequest(BaseModel):
    model: str
    messages: list
    web_search: bool

# Model untuk menerima payload untuk generate image
class ImageRequest(BaseModel):
    model: str
    prompt: str
    response_format: str

@app.post("/chat")
async def chat(request: MessageRequest):
    # Inisialisasi Client g4f
    client = Client()
    
    # Menangani response dari g4f untuk chat
    response = client.chat.completions.create(
        model=request.model,
        messages=request.messages,
        web_search=request.web_search
    )
    
    # Mengembalikan response dari API
    return {"response": response}

@app.post("/generate-image")
async def generate_image(request: ImageRequest):
    # Inisialisasi Client g4f
    client = Client()
    
    # Menangani response dari g4f untuk generate image
    response = client.images.generate(
        model=request.model,
        prompt=request.prompt,
        response_format=request.response_format
    )
    
    # Mendapatkan URL gambar yang dihasilkan
    image_url = response
    print(f"Generated image URL: {image_url}")
    
    # Mengembalikan URL gambar yang dihasilkan
    return {"image_url": image_url}

# Menjalankan server dengan Uvicorn jika file ini dieksekusi langsung
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)