LongArticleTaskServer/applications/tasks/task_scheduler.py

from applications.utils import task_schedule_response
from applications.tasks.monitor_tasks import check_kimi_balance


class TaskScheduler:
    def __init__(self, data, log_service):
        self.data = data
        self.log_client = log_service

    async def deal(self):
        task_name = self.data.get("task_name")
        if not task_name:
            await self.log_client.log(
                contents={
                    "task": task_name,
                    "function": "task_scheduler_deal",
                    "message": "not task name in params",
                    "status": "fail",
                    "data": self.data,
                }
            )
            return await task_schedule_response.fail_response(
                error_code="4002", error_message="task_name must be input"
            )

        match task_name:
            case "check_kimi_balance":
                response = await check_kimi_balance()
                await self.log_client.log(
                    contents={
                        "task": task_name,
                        "function": "task_scheduler_deal",
                        "message": "check_kimi_balance task execute successfully",
                        "status": "success",
                        "data": response,
                    }
                )
                return await task_schedule_response.success_response(
                    task_name=task_name, data=response
                )
            case _:
                await self.log_client.log(
                    contents={
                        "task": task_name,
                        "function": "task_scheduler_deal",
                        "message": "wrong task input",
                        "status": "success",
                        "data": self.data,
                    }
                )
                return await task_schedule_response.fail_response(
                    error_code="4001", error_message="wrong task name input"
                )