from fastapi import FastAPI, Request
from transformers import pipeline
import json
import os

app = FastAPI()

# Set your Hugging Face API token
huggingface_token = os.environ.get("huggingface_token")

# Load the model
generator = pipeline("text-generation", model="EleutherAI/gpt-neo-2.7B")

conversation_history = []

@app.post("/chat")
async def chat(message: dict):
    global conversation_history
    # Extract user message
    user_message = message["message"]
    conversation_history.append({"role": "user", "content": user_message})

    messages = [msg["content"] for msg in conversation_history]

    # Add a system message to instruct the model
    messages.insert(0, "You are a story writer. Whatever the prompt, you always write a short story of 30 words.")

    # Generate response using Hugging Face model
    reply = generator(messages, max_length=30, do_sample=False)[0]["generated_text"]

    conversation_history.append({"role": "assistant", "content": reply})

    return reply