from typing import Union
from fastapi import FastAPI
from huggingface_hub import InferenceClient
import os 
import json

app = FastAPI()

@app.get("/")
async def root():
    return {"message": "Hello World"}

@app.post("/chat")
async def chat(message: dict):
    try:
        client = InferenceClient(model="gpt2", token=os.environ.get(api_key))
        input_text = message
        response = client.text_generation(input_text)
        print(response)
        return {"message": response}
    except Exception as e:
        print(f"Error: {e}")
        raise HTTPException(status_code=500, detail="Internal server error")