import os
from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

# ✅ Use /tmp instead of /data to avoid permission errors
os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"

app = FastAPI()

# Load model and tokenizer
model_id = "microsoft/phi-3-mini-4k-instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir="/tmp/huggingface")
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32, cache_dir="/tmp/huggingface")

chat_history = []

class ChatRequest(BaseModel):
    message: str

@app.post("/chat")
async def chat(request: ChatRequest):
    global chat_history
    messages = [{"role": "system", "content": "You are a helpful assistant."}]
    for user, bot in chat_history:
        messages.append({"role": "user", "content": user})
        messages.append({"role": "assistant", "content": bot})
    messages.append({"role": "user", "content": request.message})

    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    inputs = tokenizer(prompt, return_tensors="pt")
    outputs = model.generate(**inputs, max_new_tokens=200, do_sample=True, temperature=0.7)
    response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)

    chat_history.append((request.message, response.strip()))
    return {"response": response.strip()}