from fastapi import FastAPI
from pydantic import BaseModel
from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
from huggingface_hub import login
import os

# Authenticate with Hugging Face
login(token=os.getenv("HF_TOKEN"))

app = FastAPI()

pipe = pipeline("text-generation", model="mistralai/Mistral-Small-24B-Instruct-2501")
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-Small-24B-Instruct-2501")
model = AutoModelForCausalLM.from_pretrained("mistralai/Mistral-Small-24B-Instruct-2501")

class ChatRequest(BaseModel):
    messages: list

@app.get("/")
def read_root():
    return {"message": "Waredocs LLM API is running!"}

@app.get("/ask")
def ask_question(prompt: str):
    messages = [{"role": "user", "content": prompt}]
    result = pipe(prompt, max_length=200)
    return result[0]