from fastapi import FastAPI, HTTPException
from fastapi.responses import PlainTextResponse
from pydantic import BaseModel
import os
from google import genai
from google.genai import types

app = FastAPI()

class GenerateRequest(BaseModel):
    system_prompt: str
    user_input: str

@app.post("/generate", response_class=PlainTextResponse)
def generate(req: GenerateRequest):
    client = genai.Client(api_key=os.environ.get("GEMINI_API_KEY"))
    model = "gemini-flash-latest"
    
    contents = [
        types.Content(
            role="user",
            parts=[types.Part.from_text(text=req.user_input)]
        )
    ]
    
    config = types.GenerateContentConfig(
        thinking_config=types.ThinkingConfig(thinking_budget=0),
        response_mime_type="application/json",
        system_instruction=[types.Part.from_text(text=req.system_prompt)]
    )

    output = ""
    try:
        for chunk in client.models.generate_content_stream(
            model=model,
            contents=contents,
            config=config
        ):
            output += chunk.text
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    
    return output