Skip to content

Agent Module

erniebot_agent.agents

Agent

The base class for agents.

Typically, this is the class that a custom agent class should inherit from. A class inheriting 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 
 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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
class Agent(GradioMixin, BaseAgent[BaseERNIEBot]):
    """The base class for agents.

    Typically, this is the class that a custom agent class should inherit from.
    A class inheriting 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, List[BaseTool]],
        *,
        memory: Optional[Memory] = None,
        system_message: Optional[SystemMessage] = None,
        callbacks: Optional[Union[CallbackManager, List[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: A list of 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_message: 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.
        """
        super().__init__()
        self.llm = llm
        if isinstance(tools, ToolManager):
            self._tool_manager = tools
        else:
            self._tool_manager = ToolManager(tools)

        if memory is None:
            self.memory = WholeMemory()
        else:
            self.memory = memory

        if system_message:
            self.system_message = system_message
        else:
            self.system_message = self.memory.get_system_message()
        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
        self._plugins = plugins
        if plugins is not None:
            raise NotImplementedError("The use of plugins is not supported yet.")
        self._init_file_needs_url()

    @final
    async def run(self, prompt: str, files: Optional[List[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)
        agent_resp = await self._run(prompt, files)
        await self._callback_manager.on_run_end(agent=self, response=agent_resp)
        return agent_resp

    @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
        await self._callback_manager.on_tool_end(agent=self, tool=tool, response=tool_resp)
        return tool_resp

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

        Args:
            messages: The input messages.
            **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, **opts)
        except (Exception, KeyboardInterrupt) as e:
            await self._callback_manager.on_llm_error(agent=self, llm=self.llm, error=e)
            raise
        await self._callback_manager.on_llm_end(agent=self, llm=self.llm, response=llm_resp)
        return llm_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_tools(self) -> List[BaseTool]:
        """Get the tools that the agent can choose from."""
        return self._tool_manager.get_tools()

    def get_tool(self, tool_name: str) -> BaseTool:
        """Get the tool by its name.

        Args:
            tool_name: the tool name of the tool to get.
        """
        return self._tool_manager.get_tool(tool_name)

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

    async def get_file_manager(self) -> FileManager:
        if self._file_manager is None:
            file_manager = await GlobalFileManagerHandler().get()
        else:
            file_manager = self._file_manager
        return file_manager

    @abc.abstractmethod
    async def _run(self, prompt: str, files: Optional[List[File]] = None) -> AgentResponse:
        """Run the agent asynchronously without invoking callbacks.

        This method is called in `run`.
        """
        raise NotImplementedError

    async def _run_tool(self, tool: BaseTool, tool_args: str) -> ToolResponse:
        """Run the given tool asynchronously without invoking callbacks.

        This method is called in `run_tool`.
        """
        parsed_tool_args = self._parse_tool_args(tool_args)
        file_manager = await 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_list(list(parsed_tool_args.values()))
        tool_ret = await tool(**parsed_tool_args)
        if isinstance(tool_ret, dict):
            output_files = file_manager.sniff_and_extract_files_from_list(list(tool_ret.values()))
        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)

    async def _run_llm(self, messages: List[Message], functions=None, **opts: Any) -> LLMResponse:
        """Run the LLM asynchronously without invoking callbacks.

        This method is called in `run_llm`.
        """
        llm_ret = await self.llm.chat(messages, functions=functions, stream=False, **opts)
        return LLMResponse(message=llm_ret)

    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: List[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 = await 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__(llm, tools, *, memory=None, system_message=None, callbacks=None, file_manager=None, plugins=None)

Initialize an agent.

Parameters:

Name Type Description Default
llm BaseERNIEBot

An LLM for the agent to use.

 required
tools Union[ToolManager, List[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 not specified, a new WholeMemory object will be instantiated.

 None
system_message Optional[SystemMessage]

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, List[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 
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
def __init__(
    self,
    llm: BaseERNIEBot,
    tools: Union[ToolManager, List[BaseTool]],
    *,
    memory: Optional[Memory] = None,
    system_message: Optional[SystemMessage] = None,
    callbacks: Optional[Union[CallbackManager, List[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: A list of 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_message: 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.
    """
    super().__init__()
    self.llm = llm
    if isinstance(tools, ToolManager):
        self._tool_manager = tools
    else:
        self._tool_manager = ToolManager(tools)

    if memory is None:
        self.memory = WholeMemory()
    else:
        self.memory = memory

    if system_message:
        self.system_message = system_message
    else:
        self.system_message = self.memory.get_system_message()
    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
    self._plugins = plugins
    if plugins is not None:
        raise NotImplementedError("The use of plugins is not supported yet.")
    self._init_file_needs_url()

get_tool(tool_name)

Get the tool by its name.

Parameters:

Name Type Description Default
tool_name str

the tool name of the tool to get.

 required
Source code in erniebot-agent/src/erniebot_agent/agents/agent.py 
179
180
181
182
183
184
185
def get_tool(self, tool_name: str) -> BaseTool:
    """Get the tool by its name.

    Args:
        tool_name: the tool name of the tool to get.
    """
    return self._tool_manager.get_tool(tool_name)

get_tools()

Get the tools that the agent can choose from.

Source code in erniebot-agent/src/erniebot_agent/agents/agent.py 
175
176
177
def get_tools(self) -> List[BaseTool]:
    """Get the tools that the agent can choose from."""
    return self._tool_manager.get_tools()

load_tool(tool)

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 
157
158
159
160
161
162
163
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()

Clear the chat history.

Source code in erniebot-agent/src/erniebot_agent/agents/agent.py 
187
188
189
def reset_memory(self) -> None:
    """Clear the chat history."""
    self.memory.clear_chat_history()

run(prompt, files=None) async

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[List[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 
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
@final
async def run(self, prompt: str, files: Optional[List[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)
    agent_resp = await self._run(prompt, files)
    await self._callback_manager.on_run_end(agent=self, response=agent_resp)
    return agent_resp

run_llm(messages, **opts) async

Run the LLM asynchronously.

Parameters:

Name Type Description Default
messages List[Message]

The input messages.

 required
**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 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
@final
async def run_llm(self, messages: List[Message], **opts: Any) -> LLMResponse:
    """Run the LLM asynchronously.

    Args:
        messages: The input messages.
        **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, **opts)
    except (Exception, KeyboardInterrupt) as e:
        await self._callback_manager.on_llm_error(agent=self, llm=self.llm, error=e)
        raise
    await self._callback_manager.on_llm_end(agent=self, llm=self.llm, response=llm_resp)
    return llm_resp

run_tool(tool_name, tool_args) async

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 
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
@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
    await self._callback_manager.on_tool_end(agent=self, tool=tool, response=tool_resp)
    return tool_resp

unload_tool(tool)

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 
165
166
167
168
169
170
171
172
173
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 
 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
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
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, List[BaseTool]],
        *,
        memory: Optional[Memory] = None,
        system_message: Optional[SystemMessage] = None,
        callbacks: Optional[Union[CallbackManager, List[CallbackHandler]]] = None,
        file_manager: Optional[FileManager] = None,
        plugins: Optional[List[str]] = None,
        max_steps: Optional[int] = None,
    ) -> 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.
            system_message: 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.
        """
        super().__init__(
            llm=llm,
            tools=tools,
            memory=memory,
            system_message=system_message,
            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

    async def _run(self, prompt: str, files: Optional[List[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)
        while num_steps_taken < self.max_steps:
            curr_step, new_messages = await self._step(chat_history)
            chat_history.extend(new_messages)
            if not isinstance(curr_step, NoActionStep):
                steps_taken.append(curr_step)
            if isinstance(curr_step, (NoActionStep, PluginStep)):  # plugin with action
                response = self._create_finished_response(chat_history, steps_taken)
                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 _step(self, chat_history: List[Message]) -> Tuple[AgentStep, List[Message]]:
        new_messages: List[Message] = []
        input_messages = self.memory.get_messages() + chat_history
        llm_resp = await self.run_llm(
            messages=input_messages,
            functions=self._tool_manager.get_tool_schemas(),
            system=self.system_message.content if self.system_message is not None else None,
            plugins=self._plugins,
        )
        output_message = llm_resp.message  # AIMessage
        new_messages.append(output_message)
        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,
            )
        elif output_message.plugin_info is not None:
            file_manager = await 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:
            return NO_ACTION_STEP, new_messages

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

    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__(llm, tools, *, memory=None, system_message=None, callbacks=None, file_manager=None, plugins=None, max_steps=None)

Initialize a function agent.

Parameters:

Name Type Description Default
llm BaseERNIEBot

An LLM for the agent to use.

 required
tools Union[ToolManager, List[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.

 None
system_message Optional[SystemMessage]

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, List[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
Source code in erniebot-agent/src/erniebot_agent/agents/function_agent.py 
 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
def __init__(
    self,
    llm: BaseERNIEBot,
    tools: Union[ToolManager, List[BaseTool]],
    *,
    memory: Optional[Memory] = None,
    system_message: Optional[SystemMessage] = None,
    callbacks: Optional[Union[CallbackManager, List[CallbackHandler]]] = None,
    file_manager: Optional[FileManager] = None,
    plugins: Optional[List[str]] = None,
    max_steps: Optional[int] = None,
) -> 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.
        system_message: 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.
    """
    super().__init__(
        llm=llm,
        tools=tools,
        memory=memory,
        system_message=system_message,
        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

erniebot_agent.agents.callback

CallbackManager

The manager for callback handlers.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/callback_manager.py 
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
@final
class CallbackManager(object):
    """The manager for callback handlers."""

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

        Args:
            handlers: A list 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: List[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: Union[Exception, KeyboardInterrupt]
    ) -> 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: Union[Exception, KeyboardInterrupt]
    ) -> 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 _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: List[CallbackHandler] property

The list of callback handlers.

__init__(handlers)

Initialize a callback manager.

Parameters:

Name Type Description Default
handlers List[CallbackHandler]

A list of callback handlers.

 required
Source code in erniebot-agent/src/erniebot_agent/agents/callback/callback_manager.py 
31
32
33
34
35
36
37
38
39
def __init__(self, handlers: List[CallbackHandler]):
    """Initialize a callback manager.

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

add_handler(handler)

Add a callback handler.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/callback_manager.py 
46
47
48
def add_handler(self, handler: CallbackHandler):
    """Add a callback handler."""
    self._handlers.append(handler)

remove_all_handlers()

Remove all callback handlers.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/callback_manager.py 
58
59
60
def remove_all_handlers(self):
    """Remove all callback handlers."""
    self._handlers.clear()

remove_handler(handler)

Remove a callback handler.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/callback_manager.py 
50
51
52
def remove_handler(self, handler: CallbackHandler):
    """Remove a callback handler."""
    self._handlers.remove(handler)

set_handlers(handlers)

Set the callback handlers.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/callback_manager.py 
54
55
56
def set_handlers(self, handlers: List[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 
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
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 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: Union[Exception, KeyboardInterrupt]
    ) -> 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 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: Union[Exception, KeyboardInterrupt]
    ) -> 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 ends running.

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

on_llm_end(agent, llm, response) async

Called when the LLM 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 
44
45
46
47
48
49
50
51
async def on_llm_end(self, agent: BaseAgent, llm: ChatModel, response: LLMResponse) -> None:
    """Called when the LLM 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(agent, llm, error) async

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 Union[Exception, KeyboardInterrupt]

The error that occured.

 required
Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/base.py 
53
54
55
56
57
58
59
60
61
62
async def on_llm_error(
    self, agent: BaseAgent, llm: ChatModel, error: Union[Exception, KeyboardInterrupt]
) -> 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(agent, llm, messages) async

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 
35
36
37
38
39
40
41
42
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(agent, response) async

Called when the agent 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 
93
94
95
96
97
98
99
async def on_run_end(self, agent: BaseAgent, response: AgentResponse) -> None:
    """Called when the agent ends running.

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

on_run_start(agent, prompt) async

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 
27
28
29
30
31
32
33
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(agent, tool, response) async

Called when a tool 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 
73
74
75
76
77
78
79
80
async def on_tool_end(self, agent: BaseAgent, tool: BaseTool, response: ToolResponse) -> None:
    """Called when a tool 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(agent, tool, error) async

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 Union[Exception, KeyboardInterrupt]

The error that occured.

 required
Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/base.py 
82
83
84
85
86
87
88
89
90
91
async def on_tool_error(
    self, agent: BaseAgent, tool: BaseTool, error: Union[Exception, KeyboardInterrupt]
) -> 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(agent, tool, input_args) async

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 
64
65
66
67
68
69
70
71
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 
 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
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 = default_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 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 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 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__(logger=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 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
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 = default_logger
    else:
        self.logger = logger

on_llm_end(agent, llm, response) async

Called to log when the LLM ends running.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/logging_handler.py 
72
73
74
75
76
77
78
79
80
async def on_llm_end(self, agent: BaseAgent, llm: ChatModel, response: LLMResponse) -> None:
    """Called to log when the LLM 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(agent, llm, messages) async

Called to log when the LLM starts running.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/logging_handler.py 
61
62
63
64
65
66
67
68
69
70
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(agent, response) async

Called to log when the agent ends running.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/logging_handler.py 
104
105
106
async def on_run_end(self, agent: BaseAgent, response: AgentResponse) -> None:
    """Called to log when the agent ends running."""
    self._agent_info("%s finished running.", agent.__class__.__name__, subject="Run", state="End")

on_run_start(agent, prompt) async

Called to log when the agent starts running.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/logging_handler.py 
51
52
53
54
55
56
57
58
59
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(agent, tool, response) async

Called to log when a tool ends running.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/logging_handler.py 
 93
 94
 95
 96
 97
 98
 99
100
101
102
async def on_tool_end(self, agent: BaseAgent, tool: BaseTool, response: ToolResponse) -> None:
    """Called to log when a tool 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(agent, tool, input_args) async

Called to log when a tool starts running.

Source code in erniebot-agent/src/erniebot_agent/agents/callback/handlers/logging_handler.py 
82
83
84
85
86
87
88
89
90
91
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 
68
69
70
71
72
@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 
75
76
77
78
79
80
81
@dataclass
class ToolResponse(object):
    """A response from a tool."""

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

AgentResponse dataclass

The final response from an agent.

Source code in erniebot-agent/src/erniebot_agent/agents/schema.py 
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
@dataclass
class AgentResponse(object):
    """The final response from an agent."""

    text: str
    chat_history: Sequence[Message]
    steps: Sequence[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