from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from huggingface_hub import InferenceClient

# Инициализация FastAPI
app = FastAPI()

# Инициализация клиента для модели
client = InferenceClient("Qwen/Qwen2.5-Coder-32B-Instruct")

# Модель данных для запроса
class ChatRequest(BaseModel):
    message: str
    history: list[tuple[str, str]] = []
    system_message: str = "You are a friendly Chatbot."
    max_tokens: int = 512
    temperature: float = 0.7
    top_p: float = 0.95

# Маршрут для обработки запросов
@app.post("/chat")
async def chat(request: ChatRequest):
    try:
        # Формируем сообщения для модели
        messages = [{"role": "system", "content": request.system_message}]

        # Добавляем историю диалога
        for user_msg, assistant_msg in request.history:
            if user_msg:
                messages.append({"role": "user", "content": user_msg})
            if assistant_msg:
                messages.append({"role": "assistant", "content": assistant_msg})

        # Добавляем текущее сообщение пользователя
        messages.append({"role": "user", "content": request.message})

        # Получаем ответ от модели
        response = ""
        for message in client.chat_completion(
            messages,
            max_tokens=request.max_tokens,
            stream=True,
            temperature=request.temperature,
            top_p=request.top_p,
        ):
            token = message.choices[0].delta.content
            response += token

        # Возвращаем ответ
        return {"response": response}

    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

# Запуск приложения (для локального тестирования)
if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=7860)