Agent Module

erniebot_agent.agents

Agent

The base class for agents.

Typically, this class should be the base class for custom agent classes. A class derived from this class must implement how the agent orchestates the components to complete tasks.

Attributes:

Name	Type	Description
llm	BaseERNIEBot	The LLM that the agent uses.
memory	Memory	The message storage that keeps the chat history.

Source code in erniebot-agent/src/erniebot_agent/agents/agent.py

class Agent(GradioMixin, BaseAgent[BaseERNIEBot]):
    """The base class for agents.

    Typically, this class should be the base class for custom agent classes. A
    class derived from this class must implement how the agent orchestates the
    components to complete tasks.

    Attributes:
        llm: The LLM that the agent uses.
        memory: The message storage that keeps the chat history.
    """

    llm: BaseERNIEBot
    memory: Memory

    def __init__(
        self,
        llm: BaseERNIEBot,
        tools: Union[ToolManager, Iterable[BaseTool]],
        *,
        memory: Optional[Memory] = None,
        system: Optional[str] = None,
        callbacks: Optional[Union[CallbackManager, Iterable[CallbackHandler]]] = None,
        file_manager: Optional[FileManager] = None,
        plugins: Optional[List[str]] = None,
    ) -> None:
        """Initialize an agent.

        Args:
            llm: An LLM for the agent to use.
            tools: Tools for the agent to use.
            memory: A memory object that equips the agent to remember chat
                history. If not specified, a new WholeMemory object will be instantiated.
            system: A message that tells the LLM how to interpret the
                conversations.
            callbacks: A list of callback handlers for the agent to use. If
                `None`, a default list of callbacks will be used.
            file_manager: A file manager for the agent to interact with files.
                If `None`, a global file manager that can be shared among
                different components will be implicitly created and used.
            plugins: A list of names of the plugins for the agent to use. If
                `None`, the agent will use a default list of plugins. Set
                `plugins` to `[]` to disable the use of plugins.
        """
        super().__init__()
        self.llm = llm
        if isinstance(tools, ToolManager):
            self._tool_manager = tools
        else:
            self._tool_manager = ToolManager(tools)
        if memory is None:
            memory = self._create_default_memory()
        self.memory = memory

        self.system = SystemMessage(system) if system is not None else system
        if self.system is not None:
            self.memory.set_system_message(self.system)

        if callbacks is None:
            callbacks = get_default_callbacks()
        if isinstance(callbacks, CallbackManager):
            self._callback_manager = callbacks
        else:
            self._callback_manager = CallbackManager(callbacks)
        self._file_manager = file_manager or get_default_file_manager()
        self._plugins = plugins
        self._init_file_needs_url()

    @final
    async def run(self, prompt: str, files: Optional[Sequence[File]] = None) -> AgentResponse:
        """Run the agent asynchronously.

        Args:
            prompt: A natural language text describing the task that the agent
                should perform.
            files: A list of files that the agent can use to perform the task.

        Returns:
            Response from the agent.
        """
        if files:
            await self._ensure_managed_files(files)
        await self._callback_manager.on_run_start(agent=self, prompt=prompt)
        try:
            agent_resp = await self._run(prompt, files)
        except BaseException as e:
            await self._callback_manager.on_run_error(agent=self, error=e)
            raise e
        else:
            await self._callback_manager.on_run_end(agent=self, response=agent_resp)
        return agent_resp

    @final
    async def run_stream(
        self, prompt: str, files: Optional[Sequence[File]] = None
    ) -> AsyncIterator[Tuple[AgentStep, List[Message]]]:
        """Run the agent asynchronously, returning an async iterator of responses.

        Args:
            prompt: A natural language text describing the task that the agent
                should perform.
            files: A list of files that the agent can use to perform the task.
        Returns:
            Iterator of responses from the agent.
        """
        if files:
            await self._ensure_managed_files(files)
        await self._callback_manager.on_run_start(agent=self, prompt=prompt)
        try:
            async for step, msg in self._run_stream(prompt, files):
                yield (step, msg)
        except BaseException as e:
            await self._callback_manager.on_run_error(agent=self, error=e)
            raise e
        else:
            await self._callback_manager.on_run_end(
                agent=self,
                response=AgentResponse(
                    text="Agent run stopped.",
                    chat_history=self.memory.get_messages(),
                    steps=[step],
                    status="STOPPED",
                ),
            )

    @final
    async def run_llm(
        self,
        messages: List[Message],
        **llm_opts: Any,
    ) -> LLMResponse:
        """Run the LLM asynchronously, returning final response.

        Args:
            messages: The input messages.
            llm_opts: Options to pass to the LLM.

        Returns:
            Response from the LLM.
        """
        await self._callback_manager.on_llm_start(agent=self, llm=self.llm, messages=messages)
        try:
            llm_resp = await self._run_llm(messages, **(llm_opts or {}))
        except (Exception, KeyboardInterrupt) as e:
            await self._callback_manager.on_llm_error(agent=self, llm=self.llm, error=e)
            raise e
        else:
            await self._callback_manager.on_llm_end(agent=self, llm=self.llm, response=llm_resp)
        return llm_resp

    @final
    async def run_llm_stream(
        self,
        messages: List[Message],
        **llm_opts: Any,
    ) -> AsyncIterator[LLMResponse]:
        """Run the LLM asynchronously, returning an async iterator of responses

        Args:
            messages: The input messages.
            llm_opts: Options to pass to the LLM.

        Returns:
            Iterator of responses from the LLM.
        """
        llm_resp = None
        await self._callback_manager.on_llm_start(agent=self, llm=self.llm, messages=messages)
        try:
            # The LLM will return an async iterator.
            async for llm_resp in self._run_llm_stream(messages, **(llm_opts or {})):
                yield llm_resp
        except (Exception, KeyboardInterrupt) as e:
            await self._callback_manager.on_llm_error(agent=self, llm=self.llm, error=e)
            raise e
        else:
            await self._callback_manager.on_llm_end(agent=self, llm=self.llm, response=llm_resp)
        return

    @final
    async def run_tool(self, tool_name: str, tool_args: str) -> ToolResponse:
        """Run the specified tool asynchronously.

        Args:
            tool_name: The name of the tool to run.
            tool_args: The tool arguments in JSON format.

        Returns:
            Response from the tool.
        """
        tool = self._tool_manager.get_tool(tool_name)
        await self._callback_manager.on_tool_start(agent=self, tool=tool, input_args=tool_args)
        try:
            tool_resp = await self._run_tool(tool, tool_args)
        except (Exception, KeyboardInterrupt) as e:
            await self._callback_manager.on_tool_error(agent=self, tool=tool, error=e)
            raise e
        else:
            await self._callback_manager.on_tool_end(agent=self, tool=tool, response=tool_resp)
        return tool_resp

    def load_tool(self, tool: BaseTool) -> None:
        """Load a tool into the agent.

        Args:
            tool: The tool to load.
        """
        self._tool_manager.add_tool(tool)

    def unload_tool(self, tool: Union[BaseTool, str]) -> None:
        """Unload a tool from the agent.

        Args:
            tool: The tool to unload.
        """
        if isinstance(tool, str):
            tool = self.get_tool(tool)
        self._tool_manager.remove_tool(tool)

    def get_tool(self, tool_name: str) -> BaseTool:
        """Get a tool by name."""
        return self._tool_manager.get_tool(tool_name)

    def get_tools(self) -> List[BaseTool]:
        """Get the tools that the agent can choose from."""
        return self._tool_manager.get_tools()

    def reset_memory(self) -> None:
        """Clear the chat history."""
        self.memory.clear_chat_history()

    def get_file_manager(self) -> FileManager:
        # Can we create a lazy proxy for the global file manager and simply set
        # and use `self._file_manager`?
        if self._file_manager is None:
            file_manager = GlobalFileManagerHandler().get()
        else:
            file_manager = self._file_manager
        return file_manager

    @abc.abstractmethod
    async def _run(self, prompt: str, files: Optional[Sequence[File]] = None) -> AgentResponse:
        raise NotImplementedError

    @abc.abstractmethod
    async def _run_stream(
        self, prompt: str, files: Optional[Sequence[File]] = None
    ) -> AsyncIterator[Tuple[AgentStep, List[Message]]]:
        """
        Abstract asynchronous generator method that should be implemented by subclasses.
        This method should yield a sequence of (AgentStep, List[Message]) tuples based on the given
        prompt and optionally accompanying files.
        """
        if TYPE_CHECKING:
            # HACK
            # This conditional block is strictly for static type-checking purposes (e.g., mypy)
            # and will not be executed.
            only_for_mypy_type_check: Tuple[AgentStep, List[Message]] = (DEFAULT_FINISH_STEP, [])
            yield only_for_mypy_type_check

    async def _run_llm(self, messages: List[Message], **opts: Any) -> LLMResponse:
        """Run the LLM with the given messages and options.

        Args:
            messages: The input messages.
            opts: Options to pass to the LLM.

        Returns:
            Response from the LLM.
        """
        for reserved_opt in ("stream", "system", "plugins"):
            if reserved_opt in opts:
                raise TypeError(f"`{reserved_opt}` should not be set.")

        if "functions" not in opts:
            functions = self._tool_manager.get_tool_schemas()
        else:
            functions = opts.pop("functions")

        if hasattr(self.llm, "system"):
            _logger.warning(
                "The `system` message has already been set in the agent;"
                "the `system` message configured in ERNIEBot will become ineffective."
            )
        opts["system"] = self.system.content if self.system is not None else None
        opts["plugins"] = self._plugins
        llm_ret = await self.llm.chat(messages, stream=False, functions=functions, **opts)
        return LLMResponse(message=llm_ret)

    async def _run_llm_stream(self, messages: List[Message], **opts: Any) -> AsyncIterator[LLMResponse]:
        """Run the LLM, yielding an async iterator of responses.

        Args:
            messages: The input messages.
            opts: Options to pass to the LLM.

        Returns:
            Async iterator of responses from the LLM.
        """
        for reserved_opt in ("stream", "system", "plugins"):
            if reserved_opt in opts:
                raise TypeError(f"`{reserved_opt}` should not be set.")

        if "functions" not in opts:
            functions = self._tool_manager.get_tool_schemas()
        else:
            functions = opts.pop("functions")

        if hasattr(self.llm, "system"):
            _logger.warning(
                "The `system` message has already been set in the agent;"
                "the `system` message configured in ERNIEBot will become ineffective."
            )
        opts["system"] = self.system.content if self.system is not None else None
        opts["plugins"] = self._plugins
        llm_ret = await self.llm.chat(messages, stream=True, functions=functions, **opts)
        async for msg in llm_ret:
            yield LLMResponse(message=msg)

    async def _run_tool(self, tool: BaseTool, tool_args: str) -> ToolResponse:
        parsed_tool_args = self._parse_tool_args(tool_args)
        file_manager = self.get_file_manager()
        # XXX: Sniffing is less efficient and probably unnecessary.
        # Can we make a protocol to statically recognize file inputs and outputs
        # or can we have the tools introspect about this?
        input_files = file_manager.sniff_and_extract_files_from_dict(parsed_tool_args)
        tool_ret = await tool(**parsed_tool_args)
        if isinstance(tool_ret, dict):
            output_files = file_manager.sniff_and_extract_files_from_dict(tool_ret)
        else:
            output_files = []
        tool_ret_json = json.dumps(tool_ret, ensure_ascii=False)
        return ToolResponse(json=tool_ret_json, input_files=input_files, output_files=output_files)

    def _create_default_memory(self) -> Memory:
        return WholeMemory()

    def _init_file_needs_url(self):
        self.file_needs_url = False
        if self._plugins:
            for plugin in self._plugins:
                if plugin not in _PLUGINS_WO_FILE_IO:
                    self.file_needs_url = True

    def _parse_tool_args(self, tool_args: str) -> Dict[str, Any]:
        try:
            args_dict = json.loads(tool_args)
        except json.JSONDecodeError:
            raise ValueError(f"`tool_args` cannot be parsed as JSON. `tool_args`: {tool_args}")

        if not isinstance(args_dict, dict):
            raise ValueError(f"`tool_args` cannot be interpreted as a dict. `tool_args`: {tool_args}")
        return args_dict

    async def _ensure_managed_files(self, files: Sequence[File]) -> None:
        def _raise_exception(file: File) -> NoReturn:
            raise FileError(f"{repr(file)} is not managed by the file manager of the agent.")

        file_manager = self.get_file_manager()
        for file in files:
            try:
                managed_file = file_manager.look_up_file_by_id(file.id)
            except FileError:
                _raise_exception(file)
            if file is not managed_file:
                _raise_exception(file)

__init__

__init__(llm: BaseERNIEBot, tools: Union[ToolManager, Iterable[BaseTool]], *, memory: Optional[Memory] = None, system: Optional[str] = None, callbacks: Optional[Union[CallbackManager, Iterable[CallbackHandler]]] = None, file_manager: Optional[FileManager] = None, plugins: Optional[List[str]] = None) -> None

Initialize an agent.

Parameters:

Name	Type	Description	Default
llm	BaseERNIEBot	An LLM for the agent to use.	required
tools	Union[ToolManager, Iterable[BaseTool]]	Tools for the agent to use.	required
memory	Optional[Memory]	A memory object that equips the agent to remember chat history. If not specified, a new WholeMemory object will be instantiated.	None
system	Optional[str]	A message that tells the LLM how to interpret the conversations.	None
callbacks	Optional[Union[CallbackManager, Iterable[CallbackHandler]]]	A list of callback handlers for the agent to use. If None, a default list of callbacks will be used.	None
file_manager	Optional[FileManager]	A file manager for the agent to interact with files. If None, a global file manager that can be shared among different components will be implicitly created and used.	None
plugins	Optional[List[str]]	A list of names of the plugins for the agent to use. If None, the agent will use a default list of plugins. Set plugins to [] to disable the use of plugins.	None

Source code in erniebot-agent/src/erniebot_agent/agents/agent.py

def __init__(
    self,
    llm: BaseERNIEBot,
    tools: Union[ToolManager, Iterable[BaseTool]],
    *,
    memory: Optional[Memory] = None,
    system: Optional[str] = None,
    callbacks: Optional[Union[CallbackManager, Iterable[CallbackHandler]]] = None,
    file_manager: Optional[FileManager] = None,
    plugins: Optional[List[str]] = None,
) -> None:
    """Initialize an agent.

    Args:
        llm: An LLM for the agent to use.
        tools: Tools for the agent to use.
        memory: A memory object that equips the agent to remember chat
            history. If not specified, a new WholeMemory object will be instantiated.
        system: A message that tells the LLM how to interpret the
            conversations.
        callbacks: A list of callback handlers for the agent to use. If
            `None`, a default list of callbacks will be used.
        file_manager: A file manager for the agent to interact with files.
            If `None`, a global file manager that can be shared among
            different components will be implicitly created and used.
        plugins: A list of names of the plugins for the agent to use. If
            `None`, the agent will use a default list of plugins. Set
            `plugins` to `[]` to disable the use of plugins.
    """
    super().__init__()
    self.llm = llm
    if isinstance(tools, ToolManager):
        self._tool_manager = tools
    else:
        self._tool_manager = ToolManager(tools)
    if memory is None:
        memory = self._create_default_memory()
    self.memory = memory

    self.system = SystemMessage(system) if system is not None else system
    if self.system is not None:
        self.memory.set_system_message(self.system)

    if callbacks is None:
        callbacks = get_default_callbacks()
    if isinstance(callbacks, CallbackManager):
        self._callback_manager = callbacks
    else:
        self._callback_manager = CallbackManager(callbacks)
    self._file_manager = file_manager or get_default_file_manager()
    self._plugins = plugins
    self._init_file_needs_url()

get_tool

get_tool(tool_name: str) -> BaseTool

Get a tool by name.

Source code in erniebot-agent/src/erniebot_agent/agents/agent.py

def get_tool(self, tool_name: str) -> BaseTool:
    """Get a tool by name."""
    return self._tool_manager.get_tool(tool_name)

get_tools

get_tools() -> List[BaseTool]

Get the tools that the agent can choose from.

Source code in erniebot-agent/src/erniebot_agent/agents/agent.py

def get_tools(self) -> List[BaseTool]:
    """Get the tools that the agent can choose from."""
    return self._tool_manager.get_tools()

load_tool

load_tool(tool: BaseTool) -> None

Load a tool into the agent.

Parameters:

Name	Type	Description	Default
tool	BaseTool	The tool to load.	required

Source code in erniebot-agent/src/erniebot_agent/agents/agent.py

def load_tool(self, tool: BaseTool) -> None:
    """Load a tool into the agent.

    Args:
        tool: The tool to load.
    """
    self._tool_manager.add_tool(tool)

reset_memory

reset_memory() -> None

Clear the chat history.

Source code in erniebot-agent/src/erniebot_agent/agents/agent.py

def reset_memory(self) -> None:
    """Clear the chat history."""
    self.memory.clear_chat_history()

run async

run(prompt: str, files: Optional[Sequence[File]] = None) -> AgentResponse

Run the agent asynchronously.

Parameters:

Name	Type	Description	Default
prompt	str	A natural language text describing the task that the agent should perform.	required
files	Optional[Sequence[File]]	A list of files that the agent can use to perform the task.	None

Returns:

Type	Description
AgentResponse	Response from the agent.

Source code in erniebot-agent/src/erniebot_agent/agents/agent.py

@final
async def run(self, prompt: str, files: Optional[Sequence[File]] = None) -> AgentResponse:
    """Run the agent asynchronously.

    Args:
        prompt: A natural language text describing the task that the agent
            should perform.
        files: A list of files that the agent can use to perform the task.

    Returns:
        Response from the agent.
    """
    if files:
        await self._ensure_managed_files(files)
    await self._callback_manager.on_run_start(agent=self, prompt=prompt)
    try:
        agent_resp = await self._run(prompt, files)
    except BaseException as e:
        await self._callback_manager.on_run_error(agent=self, error=e)
        raise e
    else:
        await self._callback_manager.on_run_end(agent=self, response=agent_resp)
    return agent_resp

run_llm async

run_llm(messages: List[Message], **llm_opts: Any) -> LLMResponse

Run the LLM asynchronously, returning final response.

Parameters:

Name	Type	Description	Default
messages	List[Message]	The input messages.	required
llm_opts	Any	Options to pass to the LLM.	{}

Returns:

Type	Description
LLMResponse	Response from the LLM.

Source code in erniebot-agent/src/erniebot_agent/agents/agent.py

@final
async def run_llm(
    self,
    messages: List[Message],
    **llm_opts: Any,
) -> LLMResponse:
    """Run the LLM asynchronously, returning final response.

    Args:
        messages: The input messages.
        llm_opts: Options to pass to the LLM.

    Returns:
        Response from the LLM.
    """
    await self._callback_manager.on_llm_start(agent=self, llm=self.llm, messages=messages)
    try:
        llm_resp = await self._run_llm(messages, **(llm_opts or {}))
    except (Exception, KeyboardInterrupt) as e:
        await self._callback_manager.on_llm_error(agent=self, llm=self.llm, error=e)
        raise e
    else:
        await self._callback_manager.on_llm_end(agent=self, llm=self.llm, response=llm_resp)
    return llm_resp

run_llm_stream async

run_llm_stream(messages: List[Message], **llm_opts: Any) -> AsyncIterator[LLMResponse]

Run the LLM asynchronously, returning an async iterator of responses

Parameters:

Name	Type	Description	Default
messages	List[Message]	The input messages.	required
llm_opts	Any	Options to pass to the LLM.	{}

Returns:

Type	Description
AsyncIterator[LLMResponse]	Iterator of responses from the LLM.

Source code in erniebot-agent/src/erniebot_agent/agents/agent.py

@final
async def run_llm_stream(
    self,
    messages: List[Message],
    **llm_opts: Any,
) -> AsyncIterator[LLMResponse]:
    """Run the LLM asynchronously, returning an async iterator of responses

    Args:
        messages: The input messages.
        llm_opts: Options to pass to the LLM.

    Returns:
        Iterator of responses from the LLM.
    """
    llm_resp = None
    await self._callback_manager.on_llm_start(agent=self, llm=self.llm, messages=messages)
    try:
        # The LLM will return an async iterator.
        async for llm_resp in self._run_llm_stream(messages, **(llm_opts or {})):
            yield llm_resp
    except (Exception, KeyboardInterrupt) as e:
        await self._callback_manager.on_llm_error(agent=self, llm=self.llm, error=e)
        raise e
    else:
        await self._callback_manager.on_llm_end(agent=self, llm=self.llm, response=llm_resp)
    return

run_stream async

run_stream(prompt: str, files: Optional[Sequence[File]] = None) -> AsyncIterator[Tuple[AgentStep, List[Message]]]

Run the agent asynchronously, returning an async iterator of responses.

Parameters:

Name	Type	Description	Default
prompt	str	A natural language text describing the task that the agent should perform.	required
files	Optional[Sequence[File]]	A list of files that the agent can use to perform the task.	None

Returns: Iterator of responses from the agent.

Source code in erniebot-agent/src/erniebot_agent/agents/agent.py

@final
async def run_stream(
    self, prompt: str, files: Optional[Sequence[File]] = None
) -> AsyncIterator[Tuple[AgentStep, List[Message]]]:
    """Run the agent asynchronously, returning an async iterator of responses.

    Args:
        prompt: A natural language text describing the task that the agent
            should perform.
        files: A list of files that the agent can use to perform the task.
    Returns:
        Iterator of responses from the agent.
    """
    if files:
        await self._ensure_managed_files(files)
    await self._callback_manager.on_run_start(agent=self, prompt=prompt)
    try:
        async for step, msg in self._run_stream(prompt, files):
            yield (step, msg)
    except BaseException as e:
        await self._callback_manager.on_run_error(agent=self, error=e)
        raise e
    else:
        await self._callback_manager.on_run_end(
            agent=self,
            response=AgentResponse(
                text="Agent run stopped.",
                chat_history=self.memory.get_messages(),
                steps=[step],
                status="STOPPED",
            ),
        )

run_tool async

run_tool(tool_name: str, tool_args: str) -> ToolResponse

Run the specified tool asynchronously.

Parameters:

Name	Type	Description	Default
tool_name	str	The name of the tool to run.	required
tool_args	str	The tool arguments in JSON format.	required

Returns:

Type	Description
ToolResponse	Response from the tool.

Source code in erniebot-agent/src/erniebot_agent/agents/agent.py

@final
async def run_tool(self, tool_name: str, tool_args: str) -> ToolResponse:
    """Run the specified tool asynchronously.

    Args:
        tool_name: The name of the tool to run.
        tool_args: The tool arguments in JSON format.

    Returns:
        Response from the tool.
    """
    tool = self._tool_manager.get_tool(tool_name)
    await self._callback_manager.on_tool_start(agent=self, tool=tool, input_args=tool_args)
    try:
        tool_resp = await self._run_tool(tool, tool_args)
    except (Exception, KeyboardInterrupt) as e:
        await self._callback_manager.on_tool_error(agent=self, tool=tool, error=e)
        raise e
    else:
        await self._callback_manager.on_tool_end(agent=self, tool=tool, response=tool_resp)
    return tool_resp

unload_tool

unload_tool(tool: Union[BaseTool, str]) -> None

Unload a tool from the agent.

Parameters:

Name	Type	Description	Default
tool	Union[BaseTool, str]	The tool to unload.	required

Source code in erniebot-agent/src/erniebot_agent/agents/agent.py

def unload_tool(self, tool: Union[BaseTool, str]) -> None:
    """Unload a tool from the agent.

    Args:
        tool: The tool to unload.
    """
    if isinstance(tool, str):
        tool = self.get_tool(tool)
    self._tool_manager.remove_tool(tool)

FunctionAgent

An agent driven by function calling.

The orchestration capabilities of a function agent are powered by the function calling ability of LLMs. Typically, a function agent asks the LLM to generate a response that can be parsed into an action (e.g., calling a tool with given arguments), and then the agent takes that action, which forms an agent step. The agent repeats this process until the maximum number of steps is reached or the LLM considers the task finished.

Attributes:

Name	Type	Description
llm	BaseERNIEBot	The LLM that the agent uses.
memory	Memory	The message storage that keeps the chat history.
max_steps	int	The maximum number of steps in each agent run.

Source code in erniebot-agent/src/erniebot_agent/agents/function_agent.py

class FunctionAgent(Agent):
    """An agent driven by function calling.

    The orchestration capabilities of a function agent are powered by the
    function calling ability of LLMs. Typically, a function agent asks the LLM
    to generate a response that can be parsed into an action (e.g., calling a
    tool with given arguments), and then the agent takes that action, which
    forms an agent step. The agent repeats this process until the maximum number
    of steps is reached or the LLM considers the task finished.

    Attributes:
        llm: The LLM that the agent uses.
        memory: The message storage that keeps the chat history.
        max_steps: The maximum number of steps in each agent run.
    """

    llm: BaseERNIEBot
    memory: Memory
    max_steps: int

    def __init__(
        self,
        llm: BaseERNIEBot,
        tools: Union[ToolManager, Iterable[BaseTool]],
        *,
        memory: Optional[Memory] = None,
        system: Optional[str] = None,
        callbacks: Optional[Union[CallbackManager, Iterable[CallbackHandler]]] = None,
        file_manager: Optional[FileManager] = None,
        plugins: Optional[List[str]] = None,
        max_steps: Optional[int] = None,
        first_tools: Optional[Sequence[BaseTool]] = [],
    ) -> None:
        """Initialize a function agent.

        Args:
            llm: An LLM for the agent to use.
            tools: A list of tools for the agent to use.
            memory: A memory object that equips the agent to remember chat
                history. If `None`, a `WholeMemory` object will be used.
            system: A message that tells the LLM how to interpret the
                conversations. If `None`, the system message contained in
                `memory` will be used.
            callbacks: A list of callback handlers for the agent to use. If
                `None`, a default list of callbacks will be used.
            file_manager: A file manager for the agent to interact with files.
                If `None`, a global file manager that can be shared among
                different components will be implicitly created and used.
            plugins: A list of names of the plugins for the agent to use. If
                `None`, the agent will use a default list of plugins. Set
                `plugins` to `[]` to disable the use of plugins.
            max_steps: The maximum number of steps in each agent run. If `None`,
                use a default value.
            first_tools: Tools scheduled to be called sequentially at the
                beginning of each agent run.

        Raises:
            ValueError: if `max_steps` is non-positive.
            RuntimeError: if tools in first_tools but not in tools list.

        """
        super().__init__(
            llm=llm,
            tools=tools,
            memory=memory,
            system=system,
            callbacks=callbacks,
            file_manager=file_manager,
            plugins=plugins,
        )
        if max_steps is not None:
            if max_steps <= 0:
                raise ValueError("Invalid `max_steps` value")
            self.max_steps = max_steps
        else:
            self.max_steps = _MAX_STEPS

        if first_tools:
            self._first_tools = first_tools
            for tool in self._first_tools:
                if tool not in self.get_tools():
                    raise RuntimeError("The tool in `first_tools` must be in the tools list.")
        else:
            self._first_tools = []

    async def _run(self, prompt: str, files: Optional[Sequence[File]] = None) -> AgentResponse:
        chat_history: List[Message] = []
        steps_taken: List[AgentStep] = []

        run_input = await HumanMessage.create_with_files(
            prompt, files or [], include_file_urls=self.file_needs_url
        )

        num_steps_taken = 0
        chat_history.append(run_input)

        for tool in self._first_tools:
            curr_step, new_messages = await self._call_first_tools(chat_history, selected_tool=tool)
            if not isinstance(curr_step, EndStep):
                chat_history.extend(new_messages)
                num_steps_taken += 1
                steps_taken.append(curr_step)
            else:
                # If tool choice not work, skip this round
                _logger.warning(f"Selected tool [{tool.tool_name}] not work")

        while num_steps_taken < self.max_steps:
            curr_step, new_messages = await self._step(chat_history)
            chat_history.extend(new_messages)
            if isinstance(curr_step, ToolStep):
                steps_taken.append(curr_step)

            elif isinstance(curr_step, PluginStep):
                steps_taken.append(curr_step)
                # 预留 调用了Plugin之后不结束的接口

                # 此处为调用了Plugin之后直接结束的Plugin
                curr_step = DEFAULT_FINISH_STEP

            if isinstance(curr_step, EndStep):
                response = self._create_finished_response(chat_history, steps_taken, curr_step)
                self.memory.add_message(chat_history[0])
                self.memory.add_message(chat_history[-1])
                return response
            num_steps_taken += 1
        response = self._create_stopped_response(chat_history, steps_taken)
        return response

    async def _call_first_tools(
        self, chat_history: List[Message], selected_tool: Optional[BaseTool] = None
    ) -> Tuple[AgentStep, List[Message]]:
        input_messages = self.memory.get_messages() + chat_history
        if selected_tool is None:
            llm_resp = await self.run_llm(messages=input_messages)
            return await self._process_step(llm_resp, chat_history)

        tool_choice = {"type": "function", "function": {"name": selected_tool.tool_name}}
        llm_resp = await self.run_llm(
            messages=input_messages,
            functions=[selected_tool.function_call_schema()],  # only regist one tool
            tool_choice=tool_choice,
        )
        return await self._process_step(llm_resp, chat_history)

    async def _step(self, chat_history: List[Message]) -> Tuple[AgentStep, List[Message]]:
        """Run a step of the agent.
        Args:
            chat_history: The chat history to provide to the agent.
        Returns:
            A tuple of an agent step and a list of new messages.
        """
        input_messages = self.memory.get_messages() + chat_history
        llm_resp = await self.run_llm(messages=input_messages)
        return await self._process_step(llm_resp, chat_history)

    async def _step_stream(
        self, chat_history: List[Message]
    ) -> AsyncIterator[Tuple[AgentStep, List[Message]]]:
        """Run a step of the agent in streaming mode.
        Args:
            chat_history: The chat history to provide to the agent.
        Returns:
            An async iterator that yields a tuple of an agent step and a list ofnew messages.
        """
        input_messages = self.memory.get_messages() + chat_history
        async for llm_resp in self.run_llm_stream(messages=input_messages):
            yield await self._process_step(llm_resp, chat_history)

    async def _run_stream(
        self, prompt: str, files: Optional[Sequence[File]] = None
    ) -> AsyncIterator[Tuple[AgentStep, List[Message]]]:
        """Run the agent with the given prompt and files in streaming mode.
        Args:
            prompt: The prompt for the agent to run.
            files: A list of files for the agent to use. If `None`, use an empty
                list.
        Returns:
            If `stream` is `False`, an agent response object. If `stream` is
            `True`, an async iterator that yields agent steps one by one.
        """
        chat_history: List[Message] = []
        steps_taken: List[AgentStep] = []

        run_input = await HumanMessage.create_with_files(
            prompt, files or [], include_file_urls=self.file_needs_url
        )

        num_steps_taken = 0
        chat_history.append(run_input)

        for tool in self._first_tools:
            curr_step, new_messages = await self._call_first_tools(chat_history, selected_tool=tool)
            if not isinstance(curr_step, EndStep):
                chat_history.extend(new_messages)
                num_steps_taken += 1
                steps_taken.append(curr_step)
            else:
                # If tool choice not work, skip this round
                _logger.warning(f"Selected tool [{tool.tool_name}] not work")

        is_finished = False
        new_messages = []
        end_step_msgs = []
        while is_finished is False:
            # IMPORTANT~! We use following code to get the response from LLM
            # When finish_reason is fuction_call, run_llm_stream return all info in one step, but
            # When finish_reason is normal chat, run_llm_stream return info in multiple steps.
            async for curr_step, new_messages in self._step_stream(chat_history):
                if isinstance(curr_step, ToolStep):
                    steps_taken.append(curr_step)
                    yield curr_step, new_messages

                elif isinstance(curr_step, PluginStep):
                    steps_taken.append(curr_step)
                    # 预留 调用了Plugin之后不结束的接口

                    # 此处为调用了Plugin之后直接结束的Plugin
                    curr_step = DEFAULT_FINISH_STEP
                    yield curr_step, new_messages

                elif isinstance(curr_step, EndStep):
                    is_finished = True
                    end_step_msgs.extend(new_messages)
                    yield curr_step, new_messages
                else:
                    raise RuntimeError("Invalid step type")
            chat_history.extend(new_messages)

        self.memory.add_message(run_input)
        end_step_msg = AIMessage(content="".join([item.content for item in end_step_msgs]))
        self.memory.add_message(end_step_msg)

    async def _process_step(self, llm_resp, chat_history) -> Tuple[AgentStep, List[Message]]:
        """Process and execute a step of the agent from LLM response.
        Args:
            llm_resp: The LLM response to convert.
            chat_history: The chat history to provide to the agent.
        Returns:
            A tuple of an agent step and a list of new messages.
        """
        new_messages: List[Message] = []
        output_message = llm_resp.message  # AIMessage
        new_messages.append(output_message)
        # handle function call
        if output_message.function_call is not None:
            tool_name = output_message.function_call["name"]
            tool_args = output_message.function_call["arguments"]
            tool_resp = await self.run_tool(tool_name=tool_name, tool_args=tool_args)
            new_messages.append(FunctionMessage(name=tool_name, content=tool_resp.json))
            return (
                ToolStep(
                    info=ToolInfo(tool_name=tool_name, tool_args=tool_args),
                    result=tool_resp.json,
                    input_files=tool_resp.input_files,
                    output_files=tool_resp.output_files,
                ),
                new_messages,
            )
        # handle plugin info with input/output files
        elif output_message.plugin_info is not None:
            file_manager = self.get_file_manager()
            return (
                PluginStep(
                    info=output_message.plugin_info,
                    result=output_message.content,
                    input_files=file_manager.sniff_and_extract_files_from_text(
                        chat_history[-1].content
                    ),  # TODO: make sure this is correct.
                    output_files=file_manager.sniff_and_extract_files_from_text(output_message.content),
                ),
                new_messages,
            )
        else:
            if output_message.clarify:
                # `clarify` and [`function_call`, `plugin`(directly end)] will not appear at the same time
                return EndStep(info=EndInfo(end_reason="CLARIFY"), result=None), new_messages
            return DEFAULT_FINISH_STEP, new_messages

    def _create_finished_response(
        self,
        chat_history: List[Message],
        steps: List[AgentStep],
        curr_step: EndStep,
    ) -> AgentResponse:
        last_message = chat_history[-1]
        return AgentResponse(
            text=last_message.content,
            chat_history=chat_history,
            steps=steps,
            status=curr_step.info["end_reason"],
        )

    def _create_stopped_response(
        self,
        chat_history: List[Message],
        steps: List[AgentStep],
    ) -> AgentResponse:
        return AgentResponse(
            text="Agent run stopped early.",
            chat_history=chat_history,
            steps=steps,
            status="STOPPED",
        )

__init__

__init__(llm: BaseERNIEBot, tools: Union[ToolManager, Iterable[BaseTool]], *, memory: Optional[Memory] = None, system: Optional[str] = None, callbacks: Optional[Union[CallbackManager, Iterable[CallbackHandler]]] = None, file_manager: Optional[FileManager] = None, plugins: Optional[List[str]] = None, max_steps: Optional[int] = None, first_tools: Optional[Sequence[BaseTool]] = []) -> None

Initialize a function agent.

Parameters:

Name	Type	Description	Default
llm	BaseERNIEBot	An LLM for the agent to use.	required
tools	Union[ToolManager, Iterable[BaseTool]]	A list of tools for the agent to use.	required
memory	Optional[Memory]	A memory object that equips the agent to remember chat history. If None, a WholeMemory object will be used.	None
system	Optional[str]	A message that tells the LLM how to interpret the conversations. If None, the system message contained in memory will be used.	None
callbacks	Optional[Union[CallbackManager, Iterable[CallbackHandler]]]	A list of callback handlers for the agent to use. If None, a default list of callbacks will be used.	None
file_manager	Optional[FileManager]	A file manager for the agent to interact with files. If None, a global file manager that can be shared among different components will be implicitly created and used.	None
plugins	Optional[List[str]]	A list of names of the plugins for the agent to use. If None, the agent will use a default list of plugins. Set plugins to [] to disable the use of plugins.	None
max_steps	Optional[int]	The maximum number of steps in each agent run. If None, use a default value.	None
first_tools	Optional[Sequence[BaseTool]]	Tools scheduled to be called sequentially at the beginning of each agent run.	[]

Raises:

Type	Description
ValueError	if max_steps is non-positive.
RuntimeError	if tools in first_tools but not in tools list.

Source code in erniebot-agent/src/erniebot_agent/agents/function_agent.py

def __init__(
    self,
    llm: BaseERNIEBot,
    tools: Union[ToolManager, Iterable[BaseTool]],
    *,
    memory: Optional[Memory] = None,
    system: Optional[str] = None,
    callbacks: Optional[Union[CallbackManager, Iterable[CallbackHandler]]] = None,
    file_manager: Optional[FileManager] = None,
    plugins: Optional[List[str]] = None,
    max_steps: Optional[int] = None,
    first_tools: Optional[Sequence[BaseTool]] = [],
) -> None:
    """Initialize a function agent.

    Args:
        llm: An LLM for the agent to use.
        tools: A list of tools for the agent to use.
        memory: A memory object that equips the agent to remember chat
            history. If `None`, a `WholeMemory` object will be used.
        system: A message that tells the LLM how to interpret the
            conversations. If `None`, the system message contained in
            `memory` will be used.
        callbacks: A list of callback handlers for the agent to use. If
            `None`, a default list of callbacks will be used.
        file_manager: A file manager for the agent to interact with files.
            If `None`, a global file manager that can be shared among
            different components will be implicitly created and used.
        plugins: A list of names of the plugins for the agent to use. If
            `None`, the agent will use a default list of plugins. Set
            `plugins` to `[]` to disable the use of plugins.
        max_steps: The maximum number of steps in each agent run. If `None`,
            use a default value.
        first_tools: Tools scheduled to be called sequentially at the
            beginning of each agent run.

    Raises:
        ValueError: if `max_steps` is non-positive.
        RuntimeError: if tools in first_tools but not in tools list.

    """
    super().__init__(
        llm=llm,
        tools=tools,
        memory=memory,
        system=system,
        callbacks=callbacks,
        file_manager=file_manager,
        plugins=plugins,
    )
    if max_steps is not None:
        if max_steps <= 0:
            raise ValueError("Invalid `max_steps` value")
        self.max_steps = max_steps
    else:
        self.max_steps = _MAX_STEPS

    if first_tools:
        self._first_tools = first_tools
        for tool in self._first_tools:
            if tool not in self.get_tools():
                raise RuntimeError("The tool in `first_tools` must be in the tools list.")
    else:
        self._first_tools = []

erniebot_agent.agents.callback

CallbackManager

The manager for callback handlers.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/callback_manager.py

@final
class CallbackManager(object):
    """The manager for callback handlers."""

    def __init__(self, handlers: Iterable[CallbackHandler]):
        """Initialize a callback manager.

        Args:
            handlers: An iterable of callback handlers.
        """
        super().__init__()
        self._handlers: List[CallbackHandler] = []
        self.set_handlers(handlers)

    @property
    def handlers(self) -> List[CallbackHandler]:
        """The list of callback handlers."""
        return self._handlers

    def add_handler(self, handler: CallbackHandler):
        """Add a callback handler."""
        self._handlers.append(handler)

    def remove_handler(self, handler: CallbackHandler):
        """Remove a callback handler."""
        self._handlers.remove(handler)

    def set_handlers(self, handlers: Iterable[CallbackHandler]):
        """Set the callback handlers."""
        self._handlers[:] = handlers

    def remove_all_handlers(self):
        """Remove all callback handlers."""
        self._handlers.clear()

    async def on_run_start(self, agent: BaseAgent, prompt: str) -> None:
        await self._handle_event(EventType.RUN_START, agent=agent, prompt=prompt)

    async def on_llm_start(self, agent: BaseAgent, llm: ChatModel, messages: List[Message]) -> None:
        await self._handle_event(EventType.LLM_START, agent=agent, llm=llm, messages=messages)

    async def on_llm_end(self, agent: BaseAgent, llm: ChatModel, response: LLMResponse) -> None:
        await self._handle_event(EventType.LLM_END, agent=agent, llm=llm, response=response)

    async def on_llm_error(self, agent: BaseAgent, llm: ChatModel, error: BaseException) -> None:
        await self._handle_event(EventType.LLM_ERROR, agent=agent, llm=llm, error=error)

    async def on_tool_start(self, agent: BaseAgent, tool: BaseTool, input_args: str) -> None:
        await self._handle_event(EventType.TOOL_START, agent=agent, tool=tool, input_args=input_args)

    async def on_tool_end(self, agent: BaseAgent, tool: BaseTool, response: ToolResponse) -> None:
        await self._handle_event(EventType.TOOL_END, agent=agent, tool=tool, response=response)

    async def on_tool_error(self, agent: BaseAgent, tool: BaseTool, error: BaseException) -> None:
        await self._handle_event(EventType.TOOL_ERROR, agent=agent, tool=tool, error=error)

    async def on_run_end(self, agent: BaseAgent, response: AgentResponse) -> None:
        await self._handle_event(EventType.RUN_END, agent=agent, response=response)

    async def on_run_error(self, agent: BaseAgent, error: BaseException) -> None:
        await self._handle_event(EventType.RUN_ERROR, agent=agent, error=error)

    async def _handle_event(self, event_type: EventType, *args: Any, **kwargs: Any) -> None:
        callback_name = "on_" + event_type.value
        for handler in self._handlers:
            callback = getattr(handler, callback_name, None)
            if not inspect.iscoroutinefunction(callback):
                raise RuntimeError("Callback must be a coroutine function.")
            await callback(*args, **kwargs)

handlers property

handlers: List[CallbackHandler]

The list of callback handlers.

__init__

__init__(handlers: Iterable[CallbackHandler])

Initialize a callback manager.

Parameters:

Name	Type	Description	Default
handlers	Iterable[CallbackHandler]	An iterable of callback handlers.	required

Source code in erniebot-agent/src/erniebot_agent/agents/callback/callback_manager.py

def __init__(self, handlers: Iterable[CallbackHandler]):
    """Initialize a callback manager.

    Args:
        handlers: An iterable of callback handlers.
    """
    super().__init__()
    self._handlers: List[CallbackHandler] = []
    self.set_handlers(handlers)

add_handler

add_handler(handler: CallbackHandler)

Add a callback handler.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/callback_manager.py

def add_handler(self, handler: CallbackHandler):
    """Add a callback handler."""
    self._handlers.append(handler)

remove_all_handlers

remove_all_handlers()

Remove all callback handlers.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/callback_manager.py

def remove_all_handlers(self):
    """Remove all callback handlers."""
    self._handlers.clear()

remove_handler

remove_handler(handler: CallbackHandler)

Remove a callback handler.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/callback_manager.py

def remove_handler(self, handler: CallbackHandler):
    """Remove a callback handler."""
    self._handlers.remove(handler)

set_handlers

set_handlers(handlers: Iterable[CallbackHandler])

Set the callback handlers.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/callback_manager.py

def set_handlers(self, handlers: Iterable[CallbackHandler]):
    """Set the callback handlers."""
    self._handlers[:] = handlers

CallbackHandler

The base class for callback handlers.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/base.py

class CallbackHandler(object):
    """The base class for callback handlers."""

    async def on_run_start(self, agent: BaseAgent, prompt: str) -> None:
        """Called when the agent starts running.

        Args:
            agent: The agent that is running.
            prompt: The prompt that the agent uses as input.
        """

    async def on_llm_start(self, agent: BaseAgent, llm: ChatModel, messages: List[Message]) -> None:
        """Called when the LLM starts running.

        Args:
            agent: The agent that is running.
            llm: The LLM that is running.
            messages: The messages that the LLM uses as input.
        """

    async def on_llm_end(self, agent: BaseAgent, llm: ChatModel, response: LLMResponse) -> None:
        """Called when the LLM successfully ends running.

        Args:
            agent: The agent that is running.
            llm: The LLM that is running.
            response: The response that the LLM returns.
        """

    async def on_llm_error(self, agent: BaseAgent, llm: ChatModel, error: BaseException) -> None:
        """Called when the LLM errors.

        Args:
            agent: The agent that is running.
            llm: The LLM that is running.
            error: The error that occured.
        """

    async def on_tool_start(self, agent: BaseAgent, tool: BaseTool, input_args: str) -> None:
        """Called when a tool starts running.

        Args:
            agent: The agent that is running.
            tool: The tool that is running.
            input_args: The input arguments that the tool uses.
        """

    async def on_tool_end(self, agent: BaseAgent, tool: BaseTool, response: ToolResponse) -> None:
        """Called when a tool successfully ends running.

        Args:
            agent: The agent that is running.
            tool: The tool that is running.
            response: The response that the tool returns.
        """

    async def on_tool_error(self, agent: BaseAgent, tool: BaseTool, error: BaseException) -> None:
        """Called when a tool errors.

        Args:
            agent: The agent that is running.
            tool: The tool that is running.
            error: The error that occured.
        """

    async def on_run_end(self, agent: BaseAgent, response: AgentResponse) -> None:
        """Called when the agent successfully ends running.

        Args:
            agent: The agent that is running.
            response: The response that the agent returns.
        """

    async def on_run_error(self, agent: BaseAgent, error: BaseException) -> None:
        """Called when the agent errors.

        Args:
            agent: The agent that is running.
            error: The error that occured.
        """

on_llm_end async

on_llm_end(agent: BaseAgent, llm: ChatModel, response: LLMResponse) -> None

Called when the LLM successfully ends running.

Parameters:

Name	Type	Description	Default
agent	BaseAgent	The agent that is running.	required
llm	ChatModel	The LLM that is running.	required
response	LLMResponse	The response that the LLM returns.	required

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/base.py

async def on_llm_end(self, agent: BaseAgent, llm: ChatModel, response: LLMResponse) -> None:
    """Called when the LLM successfully ends running.

    Args:
        agent: The agent that is running.
        llm: The LLM that is running.
        response: The response that the LLM returns.
    """

on_llm_error async

on_llm_error(agent: BaseAgent, llm: ChatModel, error: BaseException) -> None

Called when the LLM errors.

Parameters:

Name	Type	Description	Default
agent	BaseAgent	The agent that is running.	required
llm	ChatModel	The LLM that is running.	required
error	BaseException	The error that occured.	required

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/base.py

async def on_llm_error(self, agent: BaseAgent, llm: ChatModel, error: BaseException) -> None:
    """Called when the LLM errors.

    Args:
        agent: The agent that is running.
        llm: The LLM that is running.
        error: The error that occured.
    """

on_llm_start async

on_llm_start(agent: BaseAgent, llm: ChatModel, messages: List[Message]) -> None

Called when the LLM starts running.

Parameters:

Name	Type	Description	Default
agent	BaseAgent	The agent that is running.	required
llm	ChatModel	The LLM that is running.	required
messages	List[Message]	The messages that the LLM uses as input.	required

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/base.py

async def on_llm_start(self, agent: BaseAgent, llm: ChatModel, messages: List[Message]) -> None:
    """Called when the LLM starts running.

    Args:
        agent: The agent that is running.
        llm: The LLM that is running.
        messages: The messages that the LLM uses as input.
    """

on_run_end async

on_run_end(agent: BaseAgent, response: AgentResponse) -> None

Called when the agent successfully ends running.

Parameters:

Name	Type	Description	Default
agent	BaseAgent	The agent that is running.	required
response	AgentResponse	The response that the agent returns.	required

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/base.py

async def on_run_end(self, agent: BaseAgent, response: AgentResponse) -> None:
    """Called when the agent successfully ends running.

    Args:
        agent: The agent that is running.
        response: The response that the agent returns.
    """

on_run_error async

on_run_error(agent: BaseAgent, error: BaseException) -> None

Called when the agent errors.

Parameters:

Name	Type	Description	Default
agent	BaseAgent	The agent that is running.	required
error	BaseException	The error that occured.	required

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/base.py

async def on_run_error(self, agent: BaseAgent, error: BaseException) -> None:
    """Called when the agent errors.

    Args:
        agent: The agent that is running.
        error: The error that occured.
    """

on_run_start async

on_run_start(agent: BaseAgent, prompt: str) -> None

Called when the agent starts running.

Parameters:

Name	Type	Description	Default
agent	BaseAgent	The agent that is running.	required
prompt	str	The prompt that the agent uses as input.	required

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/base.py

async def on_run_start(self, agent: BaseAgent, prompt: str) -> None:
    """Called when the agent starts running.

    Args:
        agent: The agent that is running.
        prompt: The prompt that the agent uses as input.
    """

on_tool_end async

on_tool_end(agent: BaseAgent, tool: BaseTool, response: ToolResponse) -> None

Called when a tool successfully ends running.

Parameters:

Name	Type	Description	Default
agent	BaseAgent	The agent that is running.	required
tool	BaseTool	The tool that is running.	required
response	ToolResponse	The response that the tool returns.	required

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/base.py

async def on_tool_end(self, agent: BaseAgent, tool: BaseTool, response: ToolResponse) -> None:
    """Called when a tool successfully ends running.

    Args:
        agent: The agent that is running.
        tool: The tool that is running.
        response: The response that the tool returns.
    """

on_tool_error async

on_tool_error(agent: BaseAgent, tool: BaseTool, error: BaseException) -> None

Called when a tool errors.

Parameters:

Name	Type	Description	Default
agent	BaseAgent	The agent that is running.	required
tool	BaseTool	The tool that is running.	required
error	BaseException	The error that occured.	required

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/base.py

async def on_tool_error(self, agent: BaseAgent, tool: BaseTool, error: BaseException) -> None:
    """Called when a tool errors.

    Args:
        agent: The agent that is running.
        tool: The tool that is running.
        error: The error that occured.
    """

on_tool_start async

on_tool_start(agent: BaseAgent, tool: BaseTool, input_args: str) -> None

Called when a tool starts running.

Parameters:

Name	Type	Description	Default
agent	BaseAgent	The agent that is running.	required
tool	BaseTool	The tool that is running.	required
input_args	str	The input arguments that the tool uses.	required

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/base.py

async def on_tool_start(self, agent: BaseAgent, tool: BaseTool, input_args: str) -> None:
    """Called when a tool starts running.

    Args:
        agent: The agent that is running.
        tool: The tool that is running.
        input_args: The input arguments that the tool uses.
    """

LoggingHandler

A callback handler for logging.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/logging_handler.py

class LoggingHandler(CallbackHandler):
    """A callback handler for logging."""

    logger: logging.Logger

    def __init__(
        self,
        logger: Optional[logging.Logger] = None,
    ) -> None:
        """Initialize a logging handler.

        Args:
            logger: The logger to use. If `None`, a default logger will be used.
        """
        super().__init__()

        if logger is None:
            self.logger = _logger
        else:
            self.logger = logger

    async def on_run_start(self, agent: BaseAgent, prompt: str) -> None:
        """Called to log when the agent starts running."""
        self._agent_info(
            "%s is about to start running with input:\n%s",
            agent.__class__.__name__,
            ColoredContent(prompt, role="user"),
            subject="Run",
            state="Start",
        )

    async def on_llm_start(self, agent: BaseAgent, llm: ChatModel, messages: List[Message]) -> None:
        """Called to log when the LLM starts running."""
        # TODO: Prettier messages
        self._agent_info(
            "%s is about to start running with input:\n%s",
            llm.__class__.__name__,
            ColoredContent(messages[-1]),
            subject="LLM",
            state="Start",
        )

    async def on_llm_end(self, agent: BaseAgent, llm: ChatModel, response: LLMResponse) -> None:
        """Called to log when the LLM successfully ends running."""
        self._agent_info(
            "%s finished running with output:\n%s",
            llm.__class__.__name__,
            ColoredContent(response.message),
            subject="LLM",
            state="End",
        )

    async def on_tool_start(self, agent: BaseAgent, tool: BaseTool, input_args: str) -> None:
        """Called to log when a tool starts running."""
        js_inputs = to_pretty_json(input_args, from_json=True)
        self._agent_info(
            "%s is about to start running with input:\n%s",
            ColoredContent(tool.__class__.__name__, role="function"),
            ColoredContent(js_inputs, role="function"),
            subject="Tool",
            state="Start",
        )

    async def on_tool_end(self, agent: BaseAgent, tool: BaseTool, response: ToolResponse) -> None:
        """Called to log when a tool successfully ends running."""
        js_inputs = to_pretty_json(response.json, from_json=True)
        self._agent_info(
            "%s finished running with output:\n%s",
            ColoredContent(tool.__class__.__name__, role="function"),
            ColoredContent(js_inputs, role="function"),
            subject="Tool",
            state="End",
        )

    async def on_run_end(self, agent: BaseAgent, response: AgentResponse) -> None:
        """Called to log when the agent successfully ends running."""
        self._agent_info("%s finished running.", agent.__class__.__name__, subject="Run", state="End")

    def _agent_info(self, msg: str, *args, subject, state, **kwargs) -> None:
        msg = f"[{subject}][{state}] {msg}"
        self.logger.info(msg, *args, **kwargs)

__init__

__init__(logger: Optional[logging.Logger] = None) -> None

Initialize a logging handler.

Parameters:

Name	Type	Description	Default
logger	Optional[Logger]	The logger to use. If None, a default logger will be used.	None

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/logging_handler.py

def __init__(
    self,
    logger: Optional[logging.Logger] = None,
) -> None:
    """Initialize a logging handler.

    Args:
        logger: The logger to use. If `None`, a default logger will be used.
    """
    super().__init__()

    if logger is None:
        self.logger = _logger
    else:
        self.logger = logger

on_llm_end async

on_llm_end(agent: BaseAgent, llm: ChatModel, response: LLMResponse) -> None

Called to log when the LLM successfully ends running.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/logging_handler.py

async def on_llm_end(self, agent: BaseAgent, llm: ChatModel, response: LLMResponse) -> None:
    """Called to log when the LLM successfully ends running."""
    self._agent_info(
        "%s finished running with output:\n%s",
        llm.__class__.__name__,
        ColoredContent(response.message),
        subject="LLM",
        state="End",
    )

on_llm_start async

on_llm_start(agent: BaseAgent, llm: ChatModel, messages: List[Message]) -> None

Called to log when the LLM starts running.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/logging_handler.py

async def on_llm_start(self, agent: BaseAgent, llm: ChatModel, messages: List[Message]) -> None:
    """Called to log when the LLM starts running."""
    # TODO: Prettier messages
    self._agent_info(
        "%s is about to start running with input:\n%s",
        llm.__class__.__name__,
        ColoredContent(messages[-1]),
        subject="LLM",
        state="Start",
    )

on_run_end async

on_run_end(agent: BaseAgent, response: AgentResponse) -> None

Called to log when the agent successfully ends running.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/logging_handler.py

async def on_run_end(self, agent: BaseAgent, response: AgentResponse) -> None:
    """Called to log when the agent successfully ends running."""
    self._agent_info("%s finished running.", agent.__class__.__name__, subject="Run", state="End")

on_run_start async

on_run_start(agent: BaseAgent, prompt: str) -> None

Called to log when the agent starts running.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/logging_handler.py

async def on_run_start(self, agent: BaseAgent, prompt: str) -> None:
    """Called to log when the agent starts running."""
    self._agent_info(
        "%s is about to start running with input:\n%s",
        agent.__class__.__name__,
        ColoredContent(prompt, role="user"),
        subject="Run",
        state="Start",
    )

on_tool_end async

on_tool_end(agent: BaseAgent, tool: BaseTool, response: ToolResponse) -> None

Called to log when a tool successfully ends running.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/logging_handler.py

async def on_tool_end(self, agent: BaseAgent, tool: BaseTool, response: ToolResponse) -> None:
    """Called to log when a tool successfully ends running."""
    js_inputs = to_pretty_json(response.json, from_json=True)
    self._agent_info(
        "%s finished running with output:\n%s",
        ColoredContent(tool.__class__.__name__, role="function"),
        ColoredContent(js_inputs, role="function"),
        subject="Tool",
        state="End",
    )

on_tool_start async

on_tool_start(agent: BaseAgent, tool: BaseTool, input_args: str) -> None

Called to log when a tool starts running.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/logging_handler.py

async def on_tool_start(self, agent: BaseAgent, tool: BaseTool, input_args: str) -> None:
    """Called to log when a tool starts running."""
    js_inputs = to_pretty_json(input_args, from_json=True)
    self._agent_info(
        "%s is about to start running with input:\n%s",
        ColoredContent(tool.__class__.__name__, role="function"),
        ColoredContent(js_inputs, role="function"),
        subject="Tool",
        state="Start",
    )

erniebot_agent.agents.schema

LLMResponse dataclass

A response from an LLM.

Source code in erniebot-agent/src/erniebot_agent/agents/schema.py

@dataclass
class LLMResponse(object):
    """A response from an LLM."""

    message: AIMessage

ToolResponse dataclass

A response from a tool.

Source code in erniebot-agent/src/erniebot_agent/agents/schema.py

@dataclass
class ToolResponse(object):
    """A response from a tool."""

    json: str
    input_files: List[File]
    output_files: List[File]

AgentResponse dataclass

The final response from an agent.

Source code in erniebot-agent/src/erniebot_agent/agents/schema.py

@dataclass
class AgentResponse(object):
    """The final response from an agent."""

    text: str
    chat_history: List[Message]
    steps: List[AgentStep]
    status: Union[Literal["FINISHED"], Literal["STOPPED"]]

    @functools.cached_property  # lazy and prevent extra fime from multiple calls
    def annotations(self) -> Dict[str, List]:
        annotations = self._get_annotations()

        return annotations

    def _get_annotations(self) -> Dict[str, List]:
        # 1. split the text into parts and add file id to each part
        file_ids = protocol.extract_file_ids(self.text)

        places = []
        for file_id in file_ids:
            # remote file-id & local file-id may have different length.
            # TODO(shiyutang): in case of multiple same file_id
            places.append((self.text.index(file_id), len(file_id)))
        else:
            sorted(places, key=lambda x: x[0])

        split_text_list = []
        prev_idx = 0
        for place in places:
            file_start_index, file_len = place
            split_text_list.append(self.text[prev_idx:file_start_index])
            split_text_list.append(self.text[file_start_index : file_start_index + file_len])
            prev_idx = file_start_index + file_len
        else:
            split_text_list.append(self.text[prev_idx:])

        # 2. parse text to dict
        annotations: Dict = {"content_parts": []}

        for file_id in split_text_list:
            if file_id in file_ids:
                for step in self.steps:
                    if isinstance(step, AgentStepWithFiles):
                        for file in step.files:
                            if file_id == file.id:
                                file_meta = file.to_dict()
                                annotations["content_parts"].append(file_meta)
            else:
                annotations["content_parts"].append({"text": file_id})

        return annotations