main.py

from fastapi import FastAPI, Request
from pydantic import BaseModel
from model_inference import generate_code

app = FastAPI()

class Prompt(BaseModel):
    prompt: str
    max_tokens: int = 128

@app.get("/")
def home():
    return {"message": "🧠 Model is online!"}

@app.post("/generate")
def generate(prompt: Prompt):
    output = generate_code(prompt.prompt, prompt.max_tokens)
    return {"output": output}