Spaces:
Runtime error
Runtime error
File size: 6,392 Bytes
b3509ba |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 |
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 |