import os
from fastapi import FastAPI, Request
from pydantic import BaseModel
import dashscope
from dashscope import Generation

# Set up the API key for Dashscope
YOUR_API_TOKEN = os.getenv('HF_TOKEN')
dashscope.api_key = YOUR_API_TOKEN

app = FastAPI()

math_messages = []

class MathQuery(BaseModel):
    question: str

def get_math_response(user_question):
    global math_messages
    if not math_messages:
        math_messages.append({'role': 'system', 'content': 'You are a helpful math assistant.'})
    math_messages = math_messages[:1]
    
    query = f"User question: {user_question}"
    math_messages.append({'role': 'user', 'content': query})
    
    response = Generation.call(
        model="qwen2-math-72b-instruct",
        messages=math_messages,
        result_format='message',
        stream=False
    )
    
    answer = response.output.choices[0].message.content if response.output else "No response available."
    if answer:
        math_messages.append({'role': 'assistant', 'content': answer})
    
    return answer

@app.post("/math-solver")
async def solve_math(query: MathQuery):
    user_question = query.question
    answer = get_math_response(user_question)
    return {"question": user_question, "answer": answer}

@app.get("/")
def root():
    return {"message": "Welcome to the FastAPI Math Solver!"}