Spaces:
Runtime error
Runtime error
| import uuid | |
| import pandas as pd | |
| from datetime import datetime | |
| from swarmai.utils.task_queue.TaskQueueBase import TaskQueueBase | |
| from swarmai.utils.task_queue.Task import Task | |
| from swarmai.agents.AgentBase import AgentBase | |
| class PandasQueue(TaskQueueBase): | |
| """Super simple implementatin of the versatile task queue using pandas DataFrame. | |
| Pretty slow, but allows for easy manipulation of tasks, filtering, etc. | |
| Thread-safeness is handeled by the TaskQueueBase class. | |
| In the current swarm architecture the taks should have following attributes: | |
| - task_id: unique identifier of the task | |
| - priority: priority of the task. Task queue will first return high priority tasks. | |
| - task_type: type of the task, so that specific agents can filter tasks | |
| - task_description: description of the task | |
| - status: status of the task, e.g. "pending", "in progress", "completed", "failed", 'cancelled' | |
| """ | |
| def __init__(self, task_types: list, agent_types: list, task_association: dict): | |
| """ | |
| Task association is a dictionary that returns a list of task_types for a given agent_type. | |
| Attributes: | |
| - task_types (list[str]): list of task types that are supported by the task queue | |
| - agent_types (list[str]): list of agent types that are supported by the task queue | |
| - task_association (dict): dictionary that returns a list of task_types for a given agent_type | |
| """ | |
| super().__init__() | |
| self.columns = ["task_id", "priority", "task_type", "task_description", "status", "add_time", "claim_time", "complete_time", "claim_agent_id"] | |
| self.tasks = pd.DataFrame(columns=self.columns) | |
| self.task_types = task_types | |
| self.agent_types = agent_types | |
| self.task_association = task_association | |
| def add_task(self, task: Task) -> bool: | |
| """Adds a task to the queue. | |
| Task attr = (task_id, priority, task_type, task_description, status) | |
| """ | |
| if task.task_type not in self.task_types: | |
| raise ValueError(f"Task type {task.task_type} is not supported.") | |
| if task.task_description is None: | |
| raise ValueError(f"Task description {task.task_description} is not valid.") | |
| if isinstance(task.task_description, str) == False: | |
| raise ValueError(f"Task description {task.task_description} is not valid.") | |
| if task.task_description == "": | |
| raise ValueError(f"Task description {task.task_description} is not valid.") | |
| priority = task.priority | |
| task_type = task.task_type | |
| task_description = task.task_description | |
| status = "pending" | |
| add_time = datetime.now() | |
| task_i = pd.DataFrame([[uuid.uuid4(), priority, task_type, task_description, status, add_time, None, None, None]], columns=self.columns) | |
| self.tasks = pd.concat([self.tasks, task_i], ignore_index=True) | |
| def get_task(self, agent: AgentBase) -> Task: | |
| """Gets the next task from the queue, based on the agent type | |
| """ | |
| supported_tasks = self._get_supported_tasks(agent.agent_type) | |
| df_clone = self.tasks.copy() | |
| # get only pending tasks | |
| df_clone = df_clone[df_clone["status"] == "pending"] | |
| # get only supported tasks | |
| df_clone = df_clone[df_clone["task_type"].isin(supported_tasks)] | |
| if len(df_clone) == 0: | |
| return None | |
| # sort by priority | |
| df_clone = df_clone.sort_values(by="priority", ascending=False) | |
| # get the first task | |
| task = df_clone.iloc[0] | |
| # claim the task | |
| status = "in progress" | |
| claim_time = datetime.now() | |
| claim_agent_id = agent.agent_id | |
| task_obj = Task(task_id=task["task_id"], priority=task["priority"], task_type=task["task_type"], task_description=task["task_description"], status=status) | |
| # update the task in the queue | |
| df_i = pd.DataFrame([[task["task_id"], task["priority"], task["task_type"], task["task_description"], status, task["add_time"], claim_time, None, claim_agent_id]], columns=self.columns) | |
| self.tasks = self.tasks[self.tasks["task_id"] != task["task_id"]] | |
| self.tasks = pd.concat([self.tasks, df_i], ignore_index=True) | |
| return task_obj | |
| def complete_task(self, task_id): | |
| """Completes the task with the given task_id. | |
| """ | |
| task = self.tasks[self.tasks["task_id"] == task_id] | |
| if len(task) == 0: | |
| """In case task was deleted from the queue""" | |
| return False | |
| task = task.iloc[0] | |
| if task["status"] != "in progress": | |
| return False | |
| status = "completed" | |
| complete_time = datetime.now() | |
| df_i = pd.DataFrame([[task["task_id"], task["priority"], task["task_type"], task["task_description"], status, task["add_time"], task["claim_time"], complete_time, task["claim_agent_id"]]], columns=self.columns) | |
| self.tasks = self.tasks[self.tasks["task_id"] != task["task_id"]] | |
| self.tasks = pd.concat([self.tasks, df_i], ignore_index=True) | |
| return True | |
| def reset_task(self, task_id: str): | |
| task = self.tasks[self.tasks["task_id"] == task_id] | |
| if len(task) == 0: | |
| """In case task was deleted from the queue""" | |
| return False | |
| task = task.iloc[0] | |
| status = "pending" | |
| df_i = pd.DataFrame([[task["task_id"], task["priority"], task["task_type"], task["task_description"], status, task["add_time"], None, None, None]], columns=self.columns) | |
| self.tasks = self.tasks[self.tasks["task_id"] != task["task_id"]] | |
| self.tasks = pd.concat([self.tasks, df_i], ignore_index=True) | |
| return True | |
| def _get_supported_tasks(self, agent_type): | |
| """Returns a list of supported tasks for a given agent type. | |
| """ | |
| if agent_type not in self.agent_types: | |
| raise ValueError(f"Agent type {agent_type} is not supported.") | |
| if self.task_association is None: | |
| # get all present task types | |
| return self.task_types | |
| return self.task_association[agent_type] | |
| def get_all_tasks(self): | |
| """Returns all tasks in the queue. | |
| Allows the manager model to bush up the tasks list to delete duplicates or unnecessary tasks. | |
| """ | |
| raise NotImplementedError |