my_modelscope_agent/agent.py

import importlib
from typing import Dict, List, Optional, Union

from .agent_types import AgentType
from .llm import LLM
from .output_parser import OutputParser, get_output_parser
from .output_wrapper import display
from .prompt import PromptGenerator, get_prompt_generator
from .retrieve import KnowledgeRetrieval, ToolRetrieval
from .tools import TOOL_INFO_LIST


class AgentExecutor:
    def custom_run_init(self,
                        task: str,
                        remote: bool = False,
                        print_info: bool = False,
                        append_files: list = []) -> List[Dict]:

        tool_list = self.retrieve_tools(task)
        knowledge_list = self.get_knowledge(task)

        self.prompt_generator.init_prompt(
            task, tool_list, knowledge_list, append_files=append_files)
        function_list = self.prompt_generator.get_function_list(tool_list)

        llm_result, exec_result = '', ''

        idx = 0
        final_res = []

        return tool_list, knowledge_list, function_list, llm_result, exec_result, idx, final_res, remote, print_info

    def custom_gene_prompt(self, llm_result, exec_result, idx):
        idx += 1

        # generate prompt and call llm
        llm_artifacts = self.prompt_generator.generate(
            llm_result, exec_result)

        return llm_artifacts, idx

    def custom_parse_llm(self, llm_artifacts, llm_result, idx, final_res, remote, print_info):
        if print_info:
            print(f'|LLM inputs in round {idx}: {llm_artifacts}')

        # parse and get tool name and arguments
        try:
            action, action_args = self.output_parser.parse_response(
                llm_result)
        except ValueError as e:
            return [{'exec_result': f'{e}'}]

        if action is None:
            # in chat mode, the final result of last instructions should be updated to prompt history
            _ = self.prompt_generator.generate(llm_result, '')

            # for summarize
            # display(llm_result, {}, idx, self.agent_type)
            return [{'end_res': final_res}]

        if action in self.available_tool_list:
            action_args = self.parse_action_args(action_args)
            tool = self.tool_list[action]

            # TODO @wenmeng.zwm remove this hack logic for image generation
            if action == 'image_gen' and self.seed:
                action_args['seed'] = self.seed
            try:
                exec_result = tool(**action_args, remote=remote)
                if print_info:
                    print(f'|exec_result: {exec_result}')

                # parse exec result and store result to agent state
                final_res.append(exec_result)
                self.parse_exec_result(exec_result)
            except Exception as e:
                exec_result = f'Action call error: {action}: {action_args}. \n Error message: {e}'
                return [{'exec_result': exec_result}]
        else:
            exec_result = f"Unknown action: '{action}'. "
            return [{'exec_result': exec_result}]

        # display result
        # display(llm_result, exec_result, idx, self.agent_type)

        return [{'no_stop': {'llm_result': llm_result, 'exec_result': exec_result, 'idx': idx, 'final_res': final_res}}]

    def __init__(self,
                 llm: LLM,
                 tool_cfg: Optional[Dict] = {},
                 agent_type: AgentType = AgentType.DEFAULT,
                 additional_tool_list: Optional[Dict] = {},
                 prompt_generator: Optional[PromptGenerator] = None,
                 output_parser: Optional[OutputParser] = None,
                 tool_retrieval: Optional[Union[bool, ToolRetrieval]] = True,
                 knowledge_retrieval: Optional[KnowledgeRetrieval] = None):
        """
        the core class of ms agent. It is responsible for the interaction between user, llm and tools,
        and return the execution result to user.

        Args:
            llm (LLM): llm model, can be load from local or a remote server.
            tool_cfg (Optional[Dict]): cfg of default tools
            agent_type (AgentType, optional): agent type. Defaults to AgentType.DEFAULT, decide which type of agent
            reasoning type to use
            additional_tool_list (Optional[Dict], optional): user-defined additional tool list. Defaults to {}.
            prompt_generator (Optional[PromptGenerator], optional): this module is responsible for generating prompt
            according to interaction result. Defaults to use MSPromptGenerator.
            output_parser (Optional[OutputParser], optional): this module is responsible for parsing output of llm
            to executable actions. Defaults to use MsOutputParser.
            tool_retrieval (Optional[Union[bool, ToolRetrieval]], optional): Retrieve related tools by input task,
            since most of the tools may be useless for LLM in specific task.
            If it is bool type and is True, will use default tool_retrieval. Defaults to True.
            knowledge_retrieval (Optional[KnowledgeRetrieval], optional): If user want to use extra knowledge,
            this component can be used to retrieve related knowledge. Defaults to None.
        """

        self.llm = llm

        self.agent_type = agent_type
        self.llm.set_agent_type(agent_type)
        self.prompt_generator = prompt_generator or get_prompt_generator(
            agent_type)
        self.output_parser = output_parser or get_output_parser(agent_type)

        self._init_tools(tool_cfg, additional_tool_list)

        if isinstance(tool_retrieval, bool) and tool_retrieval:
            tool_retrieval = ToolRetrieval()
        self.tool_retrieval = tool_retrieval
        if self.tool_retrieval:
            self.tool_retrieval.construct(
                [str(t) for t in self.tool_list.values()])
        self.knowledge_retrieval = knowledge_retrieval
        self.reset()
        self.seed = None

    def _init_tools(self,
                    tool_cfg: Dict = {},
                    additional_tool_list: Dict = {}):
        """init tool list of agent. We provide a default tool list, which is initialized by a cfg file.
        user can also provide user-defined tools by additional_tool_list.
        The key of additional_tool_list is tool name, and the value is corresponding object.

        Args:
            tool_cfg (Dict): default tool cfg.
            additional_tool_list (Dict, optional): user-defined tools. Defaults to {}.
        """
        self.tool_list = {}
        tool_info_list = {**TOOL_INFO_LIST, **additional_tool_list}
        # tools_module = importlib.import_module('modelscope_agent.tools')
        from . import tools as tools_module

        for tool_name in tool_cfg.keys():
            if tool_cfg[tool_name].get('use', False):
                assert tool_name in tool_info_list, f'Invalid tool name: {tool_name}, ' \
                                                    f'available ones are: {tool_info_list.keys()}'
                tool_class_name = tool_info_list[tool_name]
                tool_class = getattr(tools_module, tool_class_name)
                tool_name = tool_class.name
                self.tool_list[tool_name] = tool_class(tool_cfg)

        self.tool_list = {**self.tool_list, **additional_tool_list}
        # self.available_tool_list = deepcopy(self.tool_list)
        self.set_available_tools(self.tool_list.keys())

    def set_available_tools(self, available_tool_list):
        # TODO @wenmeng.zwm refine tool init
        for t in available_tool_list:
            if t not in self.tool_list:
                raise ValueError(
                    f'Unsupported tools found:{t}, please check, valid ones: {self.tool_list.keys()}'
                )

        self.available_tool_list = {
            k: self.tool_list[k]
            for k in available_tool_list
        }

    def retrieve_tools(self, query: str) -> List[str]:
        """retrieve tools given query

        Args:
            query (str): query

        """
        if self.tool_retrieval:
            retrieve_tools = self.tool_retrieval.retrieve(query)
            self.set_available_tools(available_tool_list=retrieve_tools.keys())
        return self.available_tool_list.values()

    def get_knowledge(self, query: str) -> List[str]:
        """retrieve knowledge given query

        Args:
            query (str): query

        """
        return self.knowledge_retrieval.retrieve(
            query) if self.knowledge_retrieval else []

    def run(self,
            task: str,
            remote: bool = False,
            print_info: bool = False,
            append_files: list = []) -> List[Dict]:
        """ use llm and tools to execute task given by user

        Args:
            task (str): concrete task
            remote (bool, optional): whether to execute tool in remote mode. Defaults to False.
            print_info (bool, optional): whether to print prompt info. Defaults to False.

        Returns:
            List[Dict]: execute result. One task may need to interact with llm multiple times,
            so a list of dict is returned. Each dict contains the result of one interaction.
        """

        # retrieve tools
        tool_list = self.retrieve_tools(task)
        knowledge_list = self.get_knowledge(task)

        self.prompt_generator.init_prompt(
            task, tool_list, knowledge_list, append_files=append_files)
        function_list = self.prompt_generator.get_function_list(tool_list)

        llm_result, exec_result = '', ''

        idx = 0
        final_res = []

        while True:
            idx += 1

            # generate prompt and call llm
            llm_artifacts = self.prompt_generator.generate(
                llm_result, exec_result)
            try:
                llm_result = self.llm.generate(llm_artifacts, function_list)
            except RuntimeError as e:
                return [{'exec_result': str(e)}]

            if print_info:
                print(f'|LLM inputs in round {idx}: {llm_artifacts}')

            # parse and get tool name and arguments
            try:
                action, action_args = self.output_parser.parse_response(
                    llm_result)
            except ValueError as e:
                return [{'exec_result': f'{e}'}]

            if action is None:
                # in chat mode, the final result of last instructions should be updated to prompt history
                _ = self.prompt_generator.generate(llm_result, '')

                # for summarize
                display(llm_result, {}, idx, self.agent_type)
                return final_res

            if action in self.available_tool_list:
                action_args = self.parse_action_args(action_args)
                tool = self.tool_list[action]

                # TODO @wenmeng.zwm remove this hack logic for image generation
                if action == 'image_gen' and self.seed:
                    action_args['seed'] = self.seed
                try:
                    exec_result = tool(**action_args, remote=remote)
                    if print_info:
                        print(f'|exec_result: {exec_result}')

                    # parse exec result and store result to agent state
                    final_res.append(exec_result)
                    self.parse_exec_result(exec_result)
                except Exception as e:
                    exec_result = f'Action call error: {action}: {action_args}. \n Error message: {e}'
                    return [{'exec_result': exec_result}]
            else:
                exec_result = f"Unknown action: '{action}'. "
                return [{'exec_result': exec_result}]

            # display result
            display(llm_result, exec_result, idx, self.agent_type)

    def stream_run(self,
                   task: str,
                   remote: bool = True,
                   print_info: bool = False,
                   append_files: list = []) -> Dict:
        """this is a stream version of run, which can be used in scenario like gradio.
        It will yield the result of each interaction, so that the caller can display the result

        Args:
            task (str): concrete task
            remote (bool, optional): whether to execute tool in remote mode. Defaults to True.
            print_info (bool, optional): whether to print prompt info. Defaults to False.
            files that individually used in each run, no need to record to global state

        Yields:
            Iterator[Dict]: iterator of llm response and tool execution result
        """

        # retrieve tools
        tool_list = self.retrieve_tools(task)
        knowledge_list = self.get_knowledge(task)

        self.prompt_generator.init_prompt(
            task,
            tool_list,
            knowledge_list,
            append_files=append_files,
        )
        function_list = self.prompt_generator.get_function_list(tool_list)

        llm_result, exec_result = '', ''

        idx = 0

        while True:
            idx += 1
            llm_artifacts = self.prompt_generator.generate(
                llm_result, exec_result)
            if print_info:
                print(f'|LLM inputs in round {idx}:\n{llm_artifacts}')

            llm_result = ''
            try:
                for s in self.llm.stream_generate(llm_artifacts,
                                                  function_list):
                    llm_result += s
                    yield {'llm_text': s}
            except RuntimeError:
                s = self.llm.generate(llm_artifacts)
                llm_result += s
                yield {'llm_text': s}
            except Exception as e:
                yield {'llm_text': str(e)}

            # parse and get tool name and arguments
            try:
                action, action_args = self.output_parser.parse_response(
                    llm_result)
            except ValueError as e:
                yield {'exec_result': f'{e}'}
                return

            if action is None:
                # in chat mode, the final result of last instructions should be updated to prompt history
                _ = self.prompt_generator.generate(llm_result, '')
                yield {'is_final': True}
                return

            if action in self.available_tool_list:
                # yield observation to as end of action input symbol asap
                yield {'llm_text': 'Observation: '}
                action_args = self.parse_action_args(action_args)
                tool = self.tool_list[action]

                # TODO @wenmeng.zwm remove this hack logic for image generation
                if action == 'image_gen' and self.seed:
                    action_args['seed'] = self.seed
                try:
                    exec_result = tool(**action_args, remote=remote)
                    yield {'exec_result': exec_result}

                    # parse exec result and update state
                    self.parse_exec_result(exec_result)
                except Exception as e:
                    exec_result = f'Action call error: {action}: {action_args}. \n Error message: {e}'
                    yield {'exec_result': exec_result}
                    self.prompt_generator.reset()
                    return
            else:
                exec_result = f"Unknown action: '{action}'. "
                yield {'exec_result': exec_result}
                self.prompt_generator.reset()
                return

    def reset(self):
        """
        clear history and agent state
        """
        self.prompt_generator.reset()
        self.agent_state = {}

    def parse_action_args(self, action_args):
        """
        replace action_args in str to Image/Video/Audio Wrapper, so that tool can handle them
        """
        parsed_action_args = {}
        for name, arg in action_args.items():
            try:
                true_arg = self.agent_state.get(arg, arg)
            except Exception as e:
                print(f'Error when parsing action args: {e}, using fall back')
                true_arg = arg
            parsed_action_args[name] = true_arg
        return parsed_action_args

    def parse_exec_result(self, exec_result, *args, **kwargs):
        """
        update exec result to agent state.
        key is the str representation of the result.
        """
        for k, v in exec_result.items():
            self.agent_state[str(v)] = v