from fastapi import FastAPI, Request, HTTPException
import requests
import os
from pydantic import BaseModel

app = FastAPI()


API_URL = "https://api-inference.huggingface.co/models/meta-llama/Meta-Llama-3-8B-Instruct"
HF_TOKEN = os.getenv("HF_API_KEY")  # Load from environment variable
headers = {"Authorization": f"Bearer {HF_TOKEN}"}


class PromptRequest(BaseModel):
    prompt: str


@app.post("/generate")
async def generate_text(data: PromptRequest):
    try:
        response = requests.post(API_URL, headers=headers, json={"inputs": data.prompt})
        response.raise_for_status()
        return response.json()
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))