Class SplitTask

ms_agent/tools/split_task.py:11–140

A tool special for task splitting

class SplitTask(ToolBase):
    """A tool special for task splitting"""

    def __init__(self, config: DictConfig, **kwargs):
        super().__init__(config)
        if hasattr(config, 'tools') and hasattr(config.tools, 'split_task'):
            self.tag_prefix = getattr(config.tools.split_task, 'tag_prefix',
                                      'worker-')
        else:
            self.tag_prefix = kwargs.get('tag_prefix', 'worker-')
        self.round = 0

    async def connect(self):
        pass

    async def cleanup(self):
        pass

    async def _get_tools_inner(self):
        return {
            'split_task': [
                Tool(
                    tool_name='split_to_sub_task',
                    server_name='split_task',
                    description=
                    'Split complex task into sub tasks and start them, for example, '
                    'split a website generation task into sub tasks, '
                    'you plan the framework, include code files and classes and functions, and give the detail '
                    'information to the system and query field of the subtask, then '
                    'let each subtask to write a single file',
                    parameters={
                        'type': 'object',
                        'properties': {
                            'tasks': {
                                'type':
                                'array',
                                'description':
                                'MANDATORY: Each element is a dict, which must contains two fields: '
                                '`system`(str) and `query`(str) to start one sub task.'
                            }
                        },
                        'required': ['tasks'],
                        'additionalProperties': False
                    })
            ]
        }

    async def call_tool(self, server_name: str, *, tool_name: str,
                        tool_args: dict):
        """
        1. LLMAgent will be used to start subtask
        2. config will be inherited from the parent task
        3. Supports both parallel and sequential execution modes
        """
        from ms_agent.agent import LLMAgent

        tasks = tool_args.get('tasks')
        execution_mode = tool_args.get(