from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from model import ModelManager  # Import directly from model.py

app = FastAPI(title="BanglaLLM API", description="An API for BanglaLLM GPT model", version="1.0.0")

# Load the model
model_manager = ModelManager()

# Define request schema
class Query(BaseModel):
    prompt: str
    max_length: int = 100
    temperature: float = 0.7
    top_k: int = 50
    top_p: float = 0.9

@app.get("/")
def read_root():
    return {"message": "Welcome to the BanglaLLM API"}

@app.post("/generate")
def generate(query: Query):
    try:
        response = model_manager.generate_response(
            prompt=query.prompt,
            max_length=query.max_length,
            temperature=query.temperature,
            top_k=query.top_k,
            top_p=query.top_p
        )
        return {"response": response}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))