import os
from fastapi import FastAPI, Header, HTTPException
from pydantic import BaseModel
from transformers import pipeline

API_KEY = os.getenv("API_KEY")

app = FastAPI()

generator = pipeline(
    "text-generation",
    model="mistralai/Mistral-7B-Instruct-v0.3",
    max_new_tokens=256
)

class PromptRequest(BaseModel):
    prompt: str

@app.post("/generate")
def generate(data: PromptRequest, x_api_key: str = Header(None)):
    if API_KEY is None:
        raise HTTPException(status_code=500, detail="API key not configured")

    if x_api_key != API_KEY:
        raise HTTPException(status_code=401, detail="Invalid API Key")

    result = generator(data.prompt)
    return {"response": result[0]["generated_text"]}