百度文心一言的例子

erniebot-agent/README.md

@@ -0,0 +1,44 @@
+<div align="center">
+
+<h1>ERNIE Bot Agent</h1>
+
+ERNIE Bot Agent 可以快速开发智能体。
+
+[![License](https://img.shields.io/badge/license-Apache%202-blue.svg)](LICENSE)
+![python version](https://img.shields.io/badge/python-3.8+-orange.svg)
+![support os](https://img.shields.io/badge/os-linux%2C%20win%2C%20mac-yellow.svg)
+
+</div>
+
+`ERNIE Bot Agent` 旨在为开发者提供快速搭建大模型Agent和应用的框架。该项目还在积极研发中，敬请期待我们后续的正式发版。
+
+## 主要功能
+
+### 大模型 Agent 框架
+
+`ERNIE Bot Agent` 将结合飞桨星河AI Studio社区，为开发者提供一站式的大模型Agent和应用搭建框架和平台。该项目还在积极研发中，敬请期待我们后续的正式发版。
+
+### 文心 LangChain 插件
+
+为了让大家更加高效、便捷地结合文心大模型与LangChain进行开发，`ERNIE Bot Agent`对`LangChain`框架进行了功能扩展，提供了基于文心大模型的大语言模型（LLM）组件、聊天模型（ChatModel）组件以及文本嵌入模型（Text Embedding Model）组件。详情请参见[使用范例Notebook](https://github.com/PaddlePaddle/ERNIE-Bot-SDK/blob/develop/erniebot-agent/examples/cookbook/how_to_use_langchain_extension.ipynb)。
+
+
+## 快速安装
+
+建议您可以使用pip快速安装 ERNIE Bot Agent 的最新稳定版。
+
+```shell
+pip install --upgrade erniebot-agent
+```
+
+如需使用develop版本，可以下载源码后执行如下命令安装
+
+```shell
+git clone https://github.com/PaddlePaddle/ERNIE-Bot-SDK.git
+cd ERNIE-Bot-SDK/erniebot-agent
+pip install .
+```
+
+## License
+
+ERNIE Bot Agent遵循Apache-2.0开源协议。

erniebot-agent/erniebot_agent/__init__.py

@@ -0,0 +1,19 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from erniebot_agent.utils.logging import logger, setup_logging
+
+__all__ = ["logger"]
+
+setup_logging()

erniebot-agent/erniebot_agent/agents/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

erniebot-agent/erniebot_agent/agents/base.py

@@ -0,0 +1,279 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+import inspect
+import json
+from typing import Any, Dict, List, Literal, Optional, Union
+
+from erniebot_agent.agents.callback.callback_manager import CallbackManager
+from erniebot_agent.agents.callback.default import get_default_callbacks
+from erniebot_agent.agents.callback.handlers.base import CallbackHandler
+from erniebot_agent.agents.schema import (
+    AgentFile,
+    AgentResponse,
+    LLMResponse,
+    ToolResponse,
+)
+from erniebot_agent.chat_models.base import ChatModel
+from erniebot_agent.file_io.file_manager import FileManager
+from erniebot_agent.file_io.protocol import is_local_file_id, is_remote_file_id
+from erniebot_agent.memory.base import Memory
+from erniebot_agent.messages import Message, SystemMessage
+from erniebot_agent.tools.base import Tool
+from erniebot_agent.tools.tool_manager import ToolManager
+from erniebot_agent.utils.logging import logger
+
+
+class BaseAgent(metaclass=abc.ABCMeta):
+    llm: ChatModel
+    memory: Memory
+
+    @abc.abstractmethod
+    async def async_run(self, prompt: str) -> AgentResponse:
+        raise NotImplementedError
+
+
+class Agent(BaseAgent):
+    def __init__(
+        self,
+        llm: ChatModel,
+        tools: Union[ToolManager, List[Tool]],
+        memory: Memory,
+        system_message: Optional[SystemMessage] = None,
+        *,
+        callbacks: Optional[Union[CallbackManager, List[CallbackHandler]]] = None,
+        file_manager: Optional[FileManager] = None,
+    ) -> None:
+        super().__init__()
+        self.llm = llm
+        self.memory = memory
+        # 1. Get system message exist in memory
+        # OR 2. overwrite by the system_message paased in the Agent.
+        if system_message:
+            self.system_message = system_message
+        else:
+            self.system_message = memory.get_system_message()
+        if isinstance(tools, ToolManager):
+            self._tool_manager = tools
+        else:
+            self._tool_manager = ToolManager(tools)
+        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
+
+    async def async_run(self, prompt: str) -> AgentResponse:
+        await self._callback_manager.on_run_start(agent=self, prompt=prompt)
+        agent_resp = await self._async_run(prompt)
+        await self._callback_manager.on_run_end(agent=self, response=agent_resp)
+        return agent_resp
+
+    def load_tool(self, tool: Tool) -> None:
+        self._tool_manager.add_tool(tool)
+
+    def unload_tool(self, tool: Tool) -> None:
+        self._tool_manager.remove_tool(tool)
+
+    def reset_memory(self) -> None:
+        self.memory.clear_chat_history()
+
+    def launch_gradio_demo(self, **launch_kwargs: Any):
+        # TODO: Unified optional dependencies management
+        try:
+            import gradio as gr
+        except ImportError:
+            raise ImportError(
+                "Could not import gradio, which is required for `launch_gradio_demo()`."
+                " Please run `pip install erniebot-agent[gradio]` to install the optional dependencies."
+            ) from None
+
+        raw_messages = []
+
+        def _pre_chat(text, history):
+            history.append([text, None])
+            return history, gr.update(value="", interactive=False), gr.update(interactive=False)
+
+        async def _chat(history):
+            prompt = history[-1][0]
+            if len(prompt) == 0:
+                raise gr.Error("Prompt should not be empty.")
+            response = await self.async_run(prompt)
+            history[-1][1] = response.text
+            raw_messages.extend(response.chat_history)
+            return (
+                history,
+                _messages_to_dicts(raw_messages),
+                _messages_to_dicts(self.memory.get_messages()),
+            )
+
+        def _post_chat():
+            return gr.update(interactive=True), gr.update(interactive=True)
+
+        def _clear():
+            raw_messages.clear()
+            self.reset_memory()
+            return None, None, None, None
+
+        def _messages_to_dicts(messages):
+            return [message.to_dict() for message in messages]
+
+        with gr.Blocks(
+            title="ERNIE Bot Agent Demo", theme=gr.themes.Soft(spacing_size="sm", text_size="md")
+        ) as demo:
+            with gr.Column():
+                chatbot = gr.Chatbot(
+                    label="Chat history",
+                    latex_delimiters=[
+                        {"left": "$$", "right": "$$", "display": True},
+                        {"left": "$", "right": "$", "display": False},
+                    ],
+                    bubble_full_width=False,
+                )
+                prompt_textbox = gr.Textbox(label="Prompt", placeholder="Write a prompt here...")
+                with gr.Row():
+                    submit_button = gr.Button("Submit")
+                    clear_button = gr.Button("Clear")
+                with gr.Accordion("Tools", open=False):
+                    attached_tools = self._tool_manager.get_tools()
+                    tool_descriptions = [tool.function_call_schema() for tool in attached_tools]
+                    gr.JSON(value=tool_descriptions)
+                with gr.Accordion("Raw messages", open=False):
+                    all_messages_json = gr.JSON(label="All messages")
+                    agent_memory_json = gr.JSON(label="Messges in memory")
+            prompt_textbox.submit(
+                _pre_chat,
+                inputs=[prompt_textbox, chatbot],
+                outputs=[chatbot, prompt_textbox, submit_button],
+            ).then(
+                _chat,
+                inputs=[chatbot],
+                outputs=[
+                    chatbot,
+                    all_messages_json,
+                    agent_memory_json,
+                ],
+            ).then(
+                _post_chat, outputs=[prompt_textbox, submit_button]
+            )
+            submit_button.click(
+                _pre_chat,
+                inputs=[prompt_textbox, chatbot],
+                outputs=[chatbot, prompt_textbox, submit_button],
+            ).then(
+                _chat,
+                inputs=[chatbot],
+                outputs=[
+                    chatbot,
+                    all_messages_json,
+                    agent_memory_json,
+                ],
+            ).then(
+                _post_chat, outputs=[prompt_textbox, submit_button]
+            )
+            clear_button.click(
+                _clear,
+                outputs=[
+                    chatbot,
+                    prompt_textbox,
+                    all_messages_json,
+                    agent_memory_json,
+                ],
+            )
+
+        demo.launch(**launch_kwargs)
+
+    @abc.abstractmethod
+    async def _async_run(self, prompt: str) -> AgentResponse:
+        raise NotImplementedError
+
+    async def _async_run_tool(self, tool_name: str, tool_args: str) -> ToolResponse:
+        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._async_run_tool_without_hooks(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
+
+    async def _async_run_llm(self, messages: List[Message], **opts: Any) -> LLMResponse:
+        await self._callback_manager.on_llm_start(agent=self, llm=self.llm, messages=messages)
+        try:
+            llm_resp = await self._async_run_llm_without_hooks(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
+
+    async def _async_run_tool_without_hooks(self, tool: Tool, tool_args: str) -> ToolResponse:
+        bnd_args = self._parse_tool_args(tool, tool_args)
+        # 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 = await self._sniff_and_extract_files_from_args(bnd_args.arguments, tool, "input")
+        tool_ret = await tool(*bnd_args.args, **bnd_args.kwargs)
+        output_files = await self._sniff_and_extract_files_from_args(tool_ret, tool, "output")
+        tool_ret_json = json.dumps(tool_ret, ensure_ascii=False)
+        return ToolResponse(json=tool_ret_json, files=input_files + output_files)
+
+    async def _async_run_llm_without_hooks(
+        self, messages: List[Message], functions=None, **opts: Any
+    ) -> LLMResponse:
+        llm_ret = await self.llm.async_chat(messages, functions=functions, stream=False, **opts)
+        return LLMResponse(message=llm_ret)
+
+    def _parse_tool_args(self, tool: Tool, tool_args: str) -> inspect.BoundArguments:
+        args_dict = json.loads(tool_args)
+        if not isinstance(args_dict, dict):
+            raise ValueError("`tool_args` cannot be interpreted as a dict.")
+        # TODO: Check types
+        sig = inspect.signature(tool.__call__)
+        bnd_args = sig.bind(**args_dict)
+        bnd_args.apply_defaults()
+        return bnd_args
+
+    async def _sniff_and_extract_files_from_args(
+        self, args: Dict[str, Any], tool: Tool, file_type: Literal["input", "output"]
+    ) -> List[AgentFile]:
+        agent_files: List[AgentFile] = []
+        for val in args.values():
+            if isinstance(val, str):
+                if is_local_file_id(val):
+                    if self.file_manager is None:
+                        logger.warning(
+                            f"A file is used by {repr(tool)}, but the agent has no file manager to fetch it."
+                        )
+                        continue
+                    file = self.file_manager.look_up_file_by_id(val)
+                    if file is None:
+                        raise RuntimeError(f"Unregistered ID {repr(val)} is used by {repr(tool)}.")
+                elif is_remote_file_id(val):
+                    if self.file_manager is None:
+                        logger.warning(
+                            f"A file is used by {repr(tool)}, but the agent has no file manager to fetch it."
+                        )
+                        continue
+                    file = self.file_manager.look_up_file_by_id(val)
+                    if file is None:
+                        file = await self.file_manager.retrieve_remote_file_by_id(val)
+                else:
+                    continue
+                agent_files.append(AgentFile(file=file, type=file_type, used_by=tool.tool_name))
+        return agent_files

erniebot-agent/erniebot_agent/agents/callback/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

erniebot-agent/erniebot_agent/agents/callback/callback_manager.py

@@ -0,0 +1,94 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import inspect
+from typing import TYPE_CHECKING, Any, List, Union, final
+
+from erniebot_agent.agents.callback.event import EventType
+from erniebot_agent.agents.callback.handlers.base import CallbackHandler
+from erniebot_agent.agents.schema import AgentResponse, LLMResponse, ToolResponse
+from erniebot_agent.chat_models.base import ChatModel
+from erniebot_agent.messages import Message
+from erniebot_agent.tools.base import Tool
+
+if TYPE_CHECKING:
+    from erniebot_agent.agents.base import Agent
+
+
+@final
+class CallbackManager(object):
+    def __init__(self, handlers: List[CallbackHandler]):
+        super().__init__()
+        self._handlers = handlers
+
+    @property
+    def handlers(self) -> List[CallbackHandler]:
+        return self._handlers
+
+    def add_handler(self, handler: CallbackHandler):
+        if handler in self._handlers:
+            raise RuntimeError(f"The callback handler {handler} is already registered.")
+        self._handlers.append(handler)
+
+    def remove_handler(self, handler):
+        try:
+            self._handlers.remove(handler)
+        except ValueError as e:
+            raise RuntimeError(f"The callback handler {handler} is not registered.") from e
+
+    def set_handlers(self, handlers: List[CallbackHandler]):
+        self._handlers = []
+        for handler in handlers:
+            self.add_handler(handler)
+
+    def remove_all_handlers(self):
+        self._handlers = []
+
+    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 TypeError("Callback must be a coroutine function.")
+            await callback(*args, **kwargs)
+
+    async def on_run_start(self, agent: Agent, prompt: str) -> None:
+        await self.handle_event(EventType.RUN_START, agent=agent, prompt=prompt)
+
+    async def on_llm_start(self, agent: Agent, 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: Agent, 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: Agent, 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: Agent, tool: Tool, 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: Agent, tool: Tool, response: ToolResponse) -> None:
+        await self.handle_event(EventType.TOOL_END, agent=agent, tool=tool, response=response)
+
+    async def on_tool_error(
+        self, agent: Agent, tool: Tool, 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: Agent, response: AgentResponse) -> None:
+        await self.handle_event(EventType.RUN_END, agent=agent, response=response)

erniebot-agent/erniebot_agent/agents/callback/default.py

@@ -0,0 +1,22 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List
+
+from erniebot_agent.agents.callback.handlers.base import CallbackHandler
+from erniebot_agent.agents.callback.handlers.logging_handler import LoggingHandler
+
+
+def get_default_callbacks() -> List[CallbackHandler]:
+    return [LoggingHandler()]

erniebot-agent/erniebot_agent/agents/callback/event.py

@@ -0,0 +1,26 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import enum
+
+
+class EventType(enum.Enum):
+    RUN_START = "run_start"
+    LLM_START = "llm_start"
+    LLM_END = "llm_end"
+    LLM_ERROR = "llm_error"
+    TOOL_START = "tool_start"
+    TOOL_END = "tool_end"
+    TOOL_ERROR = "tool_error"
+    RUN_END = "run_end"

erniebot-agent/erniebot_agent/agents/callback/handlers/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

erniebot-agent/erniebot_agent/agents/callback/handlers/base.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, List, Union
+
+from erniebot_agent.agents.schema import AgentResponse, LLMResponse, ToolResponse
+from erniebot_agent.chat_models.base import ChatModel
+from erniebot_agent.messages import Message
+from erniebot_agent.tools.base import Tool
+
+if TYPE_CHECKING:
+    from erniebot_agent.agents.base import Agent
+
+
+class CallbackHandler(object):
+    async def on_run_start(self, agent: Agent, prompt: str) -> None:
+        """"""
+
+    async def on_llm_start(self, agent: Agent, llm: ChatModel, messages: List[Message]) -> None:
+        """"""
+
+    async def on_llm_end(self, agent: Agent, llm: ChatModel, response: LLMResponse) -> None:
+        """"""
+
+    async def on_llm_error(
+        self, agent: Agent, llm: ChatModel, error: Union[Exception, KeyboardInterrupt]
+    ) -> None:
+        """"""
+
+    async def on_tool_start(self, agent: Agent, tool: Tool, input_args: str) -> None:
+        """"""
+
+    async def on_tool_end(self, agent: Agent, tool: Tool, response: ToolResponse) -> None:
+        """"""
+
+    async def on_tool_error(
+        self, agent: Agent, tool: Tool, error: Union[Exception, KeyboardInterrupt]
+    ) -> None:
+        """"""
+
+    async def on_run_end(self, agent: Agent, response: AgentResponse) -> None:
+        """"""

erniebot-agent/erniebot_agent/agents/callback/handlers/logging_handler.py

@@ -0,0 +1,107 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING, List, Optional, Union
+
+from erniebot_agent.agents.callback.handlers.base import CallbackHandler
+from erniebot_agent.agents.schema import AgentResponse, LLMResponse, ToolResponse
+from erniebot_agent.chat_models.base import ChatModel
+from erniebot_agent.messages import Message
+from erniebot_agent.tools.base import Tool
+from erniebot_agent.utils.json import to_pretty_json
+from erniebot_agent.utils.logging import logger as default_logger
+
+if TYPE_CHECKING:
+    from erniebot_agent.agents.base import Agent
+
+
+class LoggingHandler(CallbackHandler):
+    logger: logging.Logger
+
+    def __init__(self, logger: Optional[logging.Logger] = None) -> None:
+        super().__init__()
+        if logger is None:
+            self.logger = default_logger
+        else:
+            self.logger = logger
+
+    async def on_run_start(self, agent: Agent, prompt: str) -> None:
+        self.agent_info(
+            "%s is about to start running with input: %s\n",
+            agent.__class__.__name__,
+            prompt,
+            subject="Run",
+            state="Start",
+        )
+
+    async def on_llm_start(self, agent: Agent, llm: ChatModel, messages: List[Message]) -> None:
+        # TODO: Prettier messages
+        self.agent_info(
+            "%s is about to start running with input:\n%s\n",
+            llm.__class__.__name__,
+            messages,
+            subject="LLM",
+            state="Start",
+        )
+
+    async def on_llm_end(self, agent: Agent, llm: ChatModel, response: LLMResponse) -> None:
+        self.agent_info(
+            "%s finished running with output: %s\n",
+            llm.__class__.__name__,
+            response.message,
+            subject="LLM",
+            state="End",
+        )
+
+    async def on_llm_error(
+        self, agent: Agent, llm: ChatModel, error: Union[Exception, KeyboardInterrupt]
+    ) -> None:
+        pass
+
+    async def on_tool_start(self, agent: Agent, tool: Tool, input_args: str) -> None:
+        self.agent_info(
+            "%s is about to start running with input:\n%s\n",
+            tool.__class__.__name__,
+            to_pretty_json(input_args, from_json=True),
+            subject="Tool",
+            state="Start",
+        )
+
+    async def on_tool_end(self, agent: Agent, tool: Tool, response: ToolResponse) -> None:
+        self.agent_info(
+            "%s finished running with output:\n%s\n",
+            tool.__class__.__name__,
+            to_pretty_json(response.json, from_json=True),
+            subject="Tool",
+            state="End",
+        )
+
+    async def on_tool_error(
+        self, agent: Agent, tool: Tool, error: Union[Exception, KeyboardInterrupt]
+    ) -> None:
+        pass
+
+    async def on_run_end(self, agent: Agent, response: AgentResponse) -> None:
+        self.agent_info("%s finished running.\n", 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)
+
+    def agent_error(self, error: Union[Exception, KeyboardInterrupt], *args, subject, **kwargs) -> None:
+        error_msg = f"[{subject}][ERROR] {error}"
+        self.logger.error(error_msg, *args, **kwargs)

erniebot-agent/erniebot_agent/agents/functional_agent.py

@@ -0,0 +1,148 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Optional, Union
+
+from erniebot_agent.agents.base import Agent, ToolManager
+from erniebot_agent.agents.callback.callback_manager import CallbackManager
+from erniebot_agent.agents.callback.handlers.base import CallbackHandler
+from erniebot_agent.agents.schema import AgentAction, AgentFile, AgentResponse
+from erniebot_agent.chat_models.base import ChatModel
+from erniebot_agent.file_io.file_manager import FileManager
+from erniebot_agent.memory.base import Memory
+from erniebot_agent.messages import (
+    FunctionMessage,
+    HumanMessage,
+    Message,
+    SystemMessage,
+)
+from erniebot_agent.tools.base import Tool
+
+_MAX_STEPS = 5
+
+
+class FunctionalAgent(Agent):
+    def __init__(
+        self,
+        llm: ChatModel,
+        tools: Union[ToolManager, List[Tool]],
+        memory: Memory,
+        system_message: Optional[SystemMessage] = None,
+        *,
+        callbacks: Optional[Union[CallbackManager, List[CallbackHandler]]] = None,
+        file_manager: Optional[FileManager] = None,
+        max_steps: Optional[int] = None,
+    ) -> None:
+        super().__init__(
+            llm=llm,
+            tools=tools,
+            memory=memory,
+            system_message=system_message,
+            callbacks=callbacks,
+            file_manager=file_manager,
+        )
+        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 _async_run(self, prompt: str) -> AgentResponse:
+        chat_history: List[Message] = []
+        actions_taken: List[AgentAction] = []
+        files_involved: List[AgentFile] = []
+        ask = HumanMessage(content=prompt)
+
+        num_steps_taken = 0
+        next_step_input: Message = ask
+        while num_steps_taken < self.max_steps:
+            curr_step_output = await self._async_step(
+                next_step_input, chat_history, actions_taken, files_involved
+            )
+            if curr_step_output is None:
+                response = self._create_finished_response(chat_history, actions_taken, files_involved)
+                self.memory.add_message(chat_history[0])
+                self.memory.add_message(chat_history[-1])
+                return response
+            num_steps_taken += 1
+            next_step_input = curr_step_output
+        response = self._create_stopped_response(chat_history, actions_taken, files_involved)
+        return response
+
+    async def _async_step(
+        self,
+        step_input,
+        chat_history: List[Message],
+        actions: List[AgentAction],
+        files: List[AgentFile],
+    ) -> Optional[Message]:
+        maybe_action = await self._async_plan(step_input, chat_history)
+        if isinstance(maybe_action, AgentAction):
+            action: AgentAction = maybe_action
+            tool_resp = await self._async_run_tool(tool_name=action.tool_name, tool_args=action.tool_args)
+            actions.append(action)
+            files.extend(tool_resp.files)
+            return FunctionMessage(name=action.tool_name, content=tool_resp.json)
+        else:
+            return None
+
+    async def _async_plan(
+        self, input_message: Message, chat_history: List[Message]
+    ) -> Optional[AgentAction]:
+        chat_history.append(input_message)
+        messages = self.memory.get_messages() + chat_history
+        llm_resp = await self._async_run_llm(
+            messages=messages,
+            functions=self._tool_manager.get_tool_schemas(),
+            system=self.system_message.content if self.system_message is not None else None,
+        )
+        output_message = llm_resp.message
+        chat_history.append(output_message)
+        if output_message.function_call is not None:
+            return AgentAction(
+                tool_name=output_message.function_call["name"],  # type: ignore
+                tool_args=output_message.function_call["arguments"],
+            )
+        else:
+            return None
+
+    def _create_finished_response(
+        self,
+        chat_history: List[Message],
+        actions: List[AgentAction],
+        files: List[AgentFile],
+    ) -> AgentResponse:
+        last_message = chat_history[-1]
+        return AgentResponse(
+            text=last_message.content,
+            chat_history=chat_history,
+            actions=actions,
+            files=files,
+            status="FINISHED",
+        )
+
+    def _create_stopped_response(
+        self,
+        chat_history: List[Message],
+        actions: List[AgentAction],
+        files: List[AgentFile],
+    ) -> AgentResponse:
+        return AgentResponse(
+            text="Agent run stopped early.",
+            chat_history=chat_history,
+            actions=actions,
+            files=files,
+            status="STOPPED",
+        )

erniebot-agent/erniebot_agent/agents/schema.py

@@ -0,0 +1,93 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass
+from typing import List, Optional, Tuple, Union
+
+from erniebot_agent.file_io.base import File
+from erniebot_agent.messages import AIMessage, Message
+from typing_extensions import Literal
+
+
+@dataclass
+class AgentAction(object):
+    """An action for an agent to execute."""
+
+    tool_name: str
+    tool_args: str
+
+
+@dataclass
+class AgentPlan(object):
+    """A plan that contains a list of actions."""
+
+    actions: List[AgentAction]
+
+
+@dataclass
+class LLMResponse(object):
+    """A response from an LLM."""
+
+    message: AIMessage
+
+
+@dataclass
+class ToolResponse(object):
+    """A response from a tool."""
+
+    json: str
+    files: List["AgentFile"]
+
+
+@dataclass
+class AgentResponse(object):
+    """The final response from an agent."""
+
+    text: str
+    chat_history: List[Message]
+    actions: List[AgentAction]
+    files: List["AgentFile"]
+    status: Union[Literal["FINISHED"], Literal["STOPPED"]]
+
+    def get_last_output_file(self) -> Optional[File]:
+        for agent_file in self.files[::-1]:
+            if agent_file.type == "output":
+                return agent_file.file
+        else:
+            return None
+
+    def get_output_files(self) -> List[File]:
+        return [agent_file.file for agent_file in self.files if agent_file.type == "output"]
+
+    def get_tool_input_output_files(self, tool_name: str) -> Tuple[List[File], List[File]]:
+        input_files: List[File] = []
+        output_files: List[File] = []
+        for agent_file in self.files:
+            if agent_file.used_by == tool_name:
+                if agent_file.type == "input":
+                    input_files.append(agent_file.file)
+                elif agent_file.type == "output":
+                    output_files.append(agent_file.file)
+                else:
+                    raise RuntimeError("File type is neither input nor output.")
+        return input_files, output_files
+
+
+@dataclass
+class AgentFile(object):
+    """A file that is used by an agent."""
+
+    file: File
+    type: Literal["input", "output"]
+    used_by: str

erniebot-agent/erniebot_agent/chat_models/__init__.py

@@ -0,0 +1,17 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .erniebot import ERNIEBot
+
+__all__ = ["ERNIEBot"]

erniebot-agent/erniebot_agent/chat_models/base.py

@@ -0,0 +1,60 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import ABCMeta, abstractmethod
+from typing import Any, AsyncIterator, List, Literal, Union, overload
+
+from erniebot_agent.messages import AIMessage, AIMessageChunk, Message
+
+
+class ChatModel(metaclass=ABCMeta):
+    """The base class of chat-optimized LLM."""
+
+    def __init__(self, model: str):
+        self.model = model
+
+    @overload
+    async def async_chat(
+        self, messages: List[Message], *, stream: Literal[False] = ..., **kwargs: Any
+    ) -> AIMessage:
+        ...
+
+    @overload
+    async def async_chat(
+        self, messages: List[Message], *, stream: Literal[True], **kwargs: Any
+    ) -> AsyncIterator[AIMessageChunk]:
+        ...
+
+    @overload
+    async def async_chat(
+        self, messages: List[Message], *, stream: bool, **kwargs: Any
+    ) -> Union[AIMessage, AsyncIterator[AIMessageChunk]]:
+        ...
+
+    @abstractmethod
+    async def async_chat(
+        self, messages: List[Message], *, stream: bool = False, **kwargs: Any
+    ) -> Union[AIMessage, AsyncIterator[AIMessageChunk]]:
+        """Asynchronously chats with the LLM.
+
+        Args:
+            messages (List[Message]): A list of messages.
+            stream (bool): Whether to use streaming generation. Defaults to False.
+            **kwargs: Arbitrary keyword arguments.
+
+        Returns:
+            If stream is False, returns a single message.
+            If stream is True, returns an asynchronous iterator of message chunks.
+        """
+        raise NotImplementedError

erniebot-agent/erniebot_agent/chat_models/erniebot.py

@@ -0,0 +1,135 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import (
+    Any,
+    AsyncIterator,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Type,
+    TypeVar,
+    Union,
+    overload,
+)
+
+from erniebot_agent.chat_models.base import ChatModel
+from erniebot_agent.messages import AIMessage, AIMessageChunk, FunctionCall, Message
+
+import erniebot
+from erniebot.response import EBResponse
+
+_T = TypeVar("_T", AIMessage, AIMessageChunk)
+
+
+class ERNIEBot(ChatModel):
+    def __init__(
+        self, model: str, api_type: Optional[str] = None, access_token: Optional[str] = None
+    ) -> None:
+        """Initializes an instance of the `ERNIEBot` class.
+
+        Args:
+            model (str): The model name. It should be "ernie-bot", "ernie-bot-turbo", "ernie-bot-8k", or
+                "ernie-bot-4".
+            api_type (Optional[str]): The API type for erniebot. It should be "aistudio" or "qianfan".
+            access_token (Optional[str]): The access token for erniebot.
+        """
+        super().__init__(model=model)
+        self.api_type = api_type
+        self.access_token = access_token
+
+    @overload
+    async def async_chat(
+        self,
+        messages: List[Message],
+        *,
+        stream: Literal[False] = ...,
+        functions: Optional[List[dict]] = ...,
+        **kwargs: Any,
+    ) -> AIMessage:
+        ...
+
+    @overload
+    async def async_chat(
+        self,
+        messages: List[Message],
+        *,
+        stream: Literal[True],
+        functions: Optional[List[dict]] = ...,
+        **kwargs: Any,
+    ) -> AsyncIterator[AIMessageChunk]:
+        ...
+
+    @overload
+    async def async_chat(
+        self, messages: List[Message], *, stream: bool, functions: Optional[List[dict]] = ..., **kwargs: Any
+    ) -> Union[AIMessage, AsyncIterator[AIMessageChunk]]:
+        ...
+
+    async def async_chat(
+        self,
+        messages: List[Message],
+        *,
+        stream: bool = False,
+        functions: Optional[List[dict]] = None,
+        **kwargs: Any,
+    ) -> Union[AIMessage, AsyncIterator[AIMessageChunk]]:
+        """Asynchronously chats with the ERNIE Bot model.
+
+        Args:
+            messages (List[Message]): A list of messages.
+            stream (bool): Whether to use streaming generation. Defaults to False.
+            functions (Optional[List[dict]]): The function definitions to be used by the model.
+                Defaults to None.
+            **kwargs: Keyword arguments, such as `top_p`, `temperature`, `penalty_score`, and `system`.
+
+        Returns:
+            If `stream` is False, returns a single message.
+            If `stream` is True, returns an asynchronous iterator of message chunks.
+        """
+        cfg_dict: Dict[str, Any] = {"model": self.model, "_config_": {}}
+        if self.api_type is not None:
+            cfg_dict["_config_"]["api_type"] = self.api_type
+        if self.access_token is not None:
+            cfg_dict["_config_"]["access_token"] = self.access_token
+
+        # TODO: process system message
+        cfg_dict["messages"] = [m.to_dict() for m in messages]
+        if functions is not None:
+            cfg_dict["functions"] = functions
+
+        name_list = ["top_p", "temperature", "penalty_score", "system"]
+        for name in name_list:
+            if name in kwargs:
+                cfg_dict[name] = kwargs[name]
+
+        # TODO: Improve this when erniebot typing issue is fixed.
+        response: Any = await erniebot.ChatCompletion.acreate(stream=stream, **cfg_dict)
+        if isinstance(response, EBResponse):
+            return self.convert_response_to_output(response, AIMessage)
+        else:
+            return (self.convert_response_to_output(resp, AIMessageChunk) async for resp in response)
+
+    @staticmethod
+    def convert_response_to_output(response: EBResponse, output_type: Type[_T]) -> _T:
+        if hasattr(response, "function_call"):
+            function_call = FunctionCall(
+                name=response.function_call["name"],
+                thoughts=response.function_call["thoughts"],
+                arguments=response.function_call["arguments"],
+            )
+            return output_type(content="", function_call=function_call, token_usage=response.usage)
+        else:
+            return output_type(content=response.result, function_call=None, token_usage=response.usage)

erniebot-agent/erniebot_agent/extensions/langchain/chat_models/__init__.py

@@ -0,0 +1 @@
+from .erniebot import ErnieBotChat

erniebot-agent/erniebot_agent/extensions/langchain/chat_models/erniebot.py

@@ -0,0 +1,356 @@
+from __future__ import annotations
+
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Type,
+    Union,
+)
+
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForLLMRun,
+    CallbackManagerForLLMRun,
+)
+from langchain.chat_models.base import BaseChatModel
+from langchain.llms.base import create_base_retry_decorator
+from langchain.pydantic_v1 import Field, root_validator
+from langchain.schema import ChatGeneration, ChatResult
+from langchain.schema.messages import (
+    AIMessage,
+    AIMessageChunk,
+    BaseMessage,
+    ChatMessage,
+    FunctionMessage,
+    HumanMessage,
+    SystemMessage,
+)
+from langchain.schema.output import ChatGenerationChunk
+from langchain.utils import get_from_dict_or_env
+
+_MessageDict = Dict[str, Any]
+
+
+class ErnieBotChat(BaseChatModel):
+    """ERNIE Bot Chat large language models API.
+
+    To use, you should have the ``erniebot`` python package installed, and the
+    environment variable ``EB_ACCESS_TOKEN`` set with your AI Studio access token.
+
+    Example:
+        .. code-block:: python
+            from erniebot_agent.extensions.langchain.chat_models import ErnieBotChat
+            erniebot_chat = ErnieBotChat(model="ernie-bot")
+    """
+
+    client: Any = None
+    max_retries: int = 6
+    """Maximum number of retries to make when generating."""
+    aistudio_access_token: Optional[str] = None
+    """AI Studio access token."""
+    streaming: Optional[bool] = False
+    """Whether to stream the results or not."""
+    model: str = "ernie-bot"
+    """Model to use."""
+    top_p: Optional[float] = 0.8
+    """Parameter of nucleus sampling that affects the diversity of generated content."""
+    temperature: Optional[float] = 0.95
+    """Sampling temperature to use."""
+    penalty_score: Optional[float] = 1
+    """Penalty assigned to tokens that have been generated."""
+    request_timeout: Optional[int] = 60
+    """How many seconds to wait for the server to send data before giving up."""
+    model_kwargs: Dict[str, Any] = Field(default_factory=dict)
+    """Holds any model parameters valid for `create` call not explicitly specified."""
+
+    ernie_client_id: Optional[str] = None
+    ernie_client_secret: Optional[str] = None
+    """For raising deprecation warnings."""
+
+    @property
+    def _default_params(self) -> Dict[str, Any]:
+        """Get the default parameters for calling ERNIE Bot API."""
+        normal_params = {
+            "model": self.model,
+            "top_p": self.top_p,
+            "temperature": self.temperature,
+            "penalty_score": self.penalty_score,
+            "request_timeout": self.request_timeout,
+        }
+        return {**normal_params, **self.model_kwargs}
+
+    @property
+    def _identifying_params(self) -> Dict[str, Any]:
+        return self._default_params
+
+    @property
+    def _invocation_params(self) -> Dict[str, Any]:
+        """Get the parameters used to invoke the model."""
+        auth_cfg: Dict[str, Optional[str]] = {
+            "api_type": "aistudio",
+            "access_token": self.aistudio_access_token,
+        }
+        return {**{"_config_": auth_cfg}, **self._default_params}
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm."""
+        return "erniebot"
+
+    @root_validator()
+    def validate_enviroment(cls, values: Dict) -> Dict:
+        values["aistudio_access_token"] = get_from_dict_or_env(
+            values,
+            "aistudio_access_token",
+            "EB_ACCESS_TOKEN",
+        )
+
+        try:
+            import erniebot
+
+            values["client"] = erniebot.ChatCompletion
+        except ImportError:
+            raise ImportError(
+                "Could not import erniebot python package. Please install it with `pip install erniebot`."
+            )
+        return values
+
+    def _generate(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        if self.streaming:
+            chunks = self._stream(messages, stop=stop, run_manager=run_manager, **kwargs)
+            generation: Optional[ChatGenerationChunk] = None
+            for chunk in chunks:
+                if generation is None:
+                    generation = chunk
+                else:
+                    generation += chunk
+            assert generation is not None
+            return ChatResult(generations=[generation])
+        else:
+            params = self._invocation_params
+            params.update(kwargs)
+            params["messages"] = self._convert_messages_to_dicts(messages)
+            system_prompt = self._build_system_prompt_from_messages(messages)
+            if system_prompt is not None:
+                params["system"] = system_prompt
+            params["stream"] = False
+            response = _create_completion_with_retry(self, run_manager=run_manager, **params)
+            return self._build_chat_result_from_response(response)
+
+    async def _agenerate(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> ChatResult:
+        if self.streaming:
+            chunks = self._astream(messages, stop=stop, run_manager=run_manager, **kwargs)
+            generation: Optional[ChatGenerationChunk] = None
+            async for chunk in chunks:
+                if generation is None:
+                    generation = chunk
+                else:
+                    generation += chunk
+            assert generation is not None
+            return ChatResult(generations=[generation])
+        else:
+            params = self._invocation_params
+            params.update(kwargs)
+            params["messages"] = self._convert_messages_to_dicts(messages)
+            system_prompt = self._build_system_prompt_from_messages(messages)
+            if system_prompt is not None:
+                params["system"] = system_prompt
+            params["stream"] = False
+            response = await _acreate_completion_with_retry(self, run_manager=run_manager, **params)
+            return self._build_chat_result_from_response(response)
+
+    def _stream(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> Iterator[ChatGenerationChunk]:
+        if stop is not None:
+            raise TypeError("Currently, `stop` is not supported when streaming is enabled.")
+        params = self._invocation_params
+        params.update(kwargs)
+        params["messages"] = self._convert_messages_to_dicts(messages)
+        system_prompt = self._build_system_prompt_from_messages(messages)
+        if system_prompt is not None:
+            params["system"] = system_prompt
+        params["stream"] = True
+        for resp in _create_completion_with_retry(self, run_manager=run_manager, **params):
+            chunk = self._build_chunk_from_response(resp)
+            yield chunk
+            if run_manager:
+                run_manager.on_llm_new_token(chunk.text, chunk=chunk)
+
+    async def _astream(
+        self,
+        messages: List[BaseMessage],
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[ChatGenerationChunk]:
+        if stop is not None:
+            raise TypeError("Currently, `stop` is not supported when streaming is enabled.")
+        params = self._invocation_params
+        params.update(kwargs)
+        params["messages"] = self._convert_messages_to_dicts(messages)
+        system_prompt = self._build_system_prompt_from_messages(messages)
+        if system_prompt is not None:
+            params["system"] = system_prompt
+        params["stream"] = True
+        async for resp in await _acreate_completion_with_retry(self, run_manager=run_manager, **params):
+            chunk = self._build_chunk_from_response(resp)
+            yield chunk
+            if run_manager:
+                await run_manager.on_llm_new_token(chunk.text, chunk=chunk)
+
+    def _build_chat_result_from_response(self, response: Mapping[str, Any]) -> ChatResult:
+        message_dict = self._build_dict_from_response(response)
+        generation = ChatGeneration(
+            message=self._convert_dict_to_message(message_dict),
+            generation_info=dict(finish_reason="stop"),
+        )
+        token_usage = response.get("usage", {})
+        llm_output = {"token_usage": token_usage, "model_name": self.model}
+        return ChatResult(generations=[generation], llm_output=llm_output)
+
+    def _build_chunk_from_response(self, response: Mapping[str, Any]) -> ChatGenerationChunk:
+        message_dict = self._build_dict_from_response(response)
+        message = self._convert_dict_to_message(message_dict)
+        msg_chunk = AIMessageChunk(
+            content=message.content,
+            additional_kwargs=message.additional_kwargs,
+        )
+        return ChatGenerationChunk(message=msg_chunk)
+
+    def _build_dict_from_response(self, response: Mapping[str, Any]) -> _MessageDict:
+        message_dict: _MessageDict = {"role": "assistant"}
+        if "function_call" in response:
+            message_dict["content"] = None
+            message_dict["function_call"] = response["function_call"]
+        else:
+            message_dict["content"] = response["result"]
+        return message_dict
+
+    def _build_system_prompt_from_messages(self, messages: List[BaseMessage]) -> Optional[str]:
+        system_message_content_list: List[str] = []
+        for msg in messages:
+            if isinstance(msg, SystemMessage):
+                if isinstance(msg.content, str):
+                    system_message_content_list.append(msg.content)
+                else:
+                    raise TypeError
+        if len(system_message_content_list) > 0:
+            return "\n".join(system_message_content_list)
+        else:
+            return None
+
+    def _convert_messages_to_dicts(self, messages: List[BaseMessage]) -> List[dict]:
+        erniebot_messages = []
+        for msg in messages:
+            if isinstance(msg, SystemMessage):
+                # Ignore system messages, as we handle them elsewhere.
+                continue
+            eb_msg = self._convert_message_to_dict(msg)
+            erniebot_messages.append(eb_msg)
+        return erniebot_messages
+
+    @staticmethod
+    def _convert_dict_to_message(message_dict: _MessageDict) -> BaseMessage:
+        role = message_dict["role"]
+        if role == "user":
+            return HumanMessage(content=message_dict["content"])
+        elif role == "assistant":
+            content = message_dict["content"] or ""
+            if message_dict.get("function_call"):
+                additional_kwargs = {"function_call": dict(message_dict["function_call"])}
+            else:
+                additional_kwargs = {}
+            return AIMessage(content=content, additional_kwargs=additional_kwargs)
+        elif role == "function":
+            return FunctionMessage(content=message_dict["content"], name=message_dict["name"])
+        else:
+            return ChatMessage(content=message_dict["content"], role=role)
+
+    @staticmethod
+    def _convert_message_to_dict(message: BaseMessage) -> _MessageDict:
+        message_dict: _MessageDict
+        if isinstance(message, ChatMessage):
+            message_dict = {"role": message.role, "content": message.content}
+        elif isinstance(message, HumanMessage):
+            message_dict = {"role": "user", "content": message.content}
+        elif isinstance(message, AIMessage):
+            message_dict = {"role": "assistant", "content": message.content}
+            if "function_call" in message.additional_kwargs:
+                message_dict["function_call"] = message.additional_kwargs["function_call"]
+                if message_dict["content"] == "":
+                    message_dict["content"] = None
+        elif isinstance(message, FunctionMessage):
+            message_dict = {
+                "role": "function",
+                "content": message.content,
+                "name": message.name,
+            }
+        else:
+            raise TypeError(f"Got unknown type {message}")
+
+        return message_dict
+
+
+def _create_completion_with_retry(
+    llm: ErnieBotChat,
+    run_manager: Optional[CallbackManagerForLLMRun] = None,
+    **kwargs: Any,
+) -> Any:
+    retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)
+
+    @retry_decorator
+    def _client_create(**kwargs: Any) -> Any:
+        return llm.client.create(**kwargs)
+
+    return _client_create(**kwargs)
+
+
+async def _acreate_completion_with_retry(
+    llm: ErnieBotChat,
+    run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+    **kwargs: Any,
+) -> Any:
+    retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)
+
+    @retry_decorator
+    async def _client_acreate(**kwargs: Any) -> Any:
+        return await llm.client.acreate(**kwargs)
+
+    return await _client_acreate(**kwargs)
+
+
+def _create_retry_decorator(
+    llm: ErnieBotChat,
+    run_manager: Optional[Union[AsyncCallbackManagerForLLMRun, CallbackManagerForLLMRun]] = None,
+) -> Callable[[Any], Any]:
+    import erniebot
+
+    errors: List[Type[BaseException]] = [
+        erniebot.errors.TimeoutError,
+        erniebot.errors.RequestLimitError,
+    ]
+    return create_base_retry_decorator(
+        error_types=errors, max_retries=llm.max_retries, run_manager=run_manager
+    )

erniebot-agent/erniebot_agent/extensions/langchain/embeddings/__init__.py

@@ -0,0 +1 @@
+from .ernie import ErnieEmbeddings

erniebot-agent/erniebot_agent/extensions/langchain/embeddings/ernie.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from typing import Any, Dict, List, Optional
+
+from langchain.pydantic_v1 import BaseModel, root_validator
+from langchain.schema.embeddings import Embeddings
+from langchain.utils import get_from_dict_or_env
+
+
+class ErnieEmbeddings(BaseModel, Embeddings):
+    """ERNIE embedding models.
+
+    To use, you should have the ``erniebot`` python package installed, and the
+    environment variable ``EB_ACCESS_TOKEN`` set with your AI Studio access token.
+
+    Example:
+        .. code-block:: python
+            from erniebot_agent.extensions.langchain.embeddings import ErnieEmbeddings
+            ernie_embeddings = ErnieEmbeddings()
+    """
+
+    client: Any = None
+    max_retries: int = 6
+    """Maximum number of retries to make when generating."""
+    chunk_size: int = 16
+    """Chunk size to use when the input is a list of texts."""
+    aistudio_access_token: Optional[str] = None
+    """AI Studio access token."""
+    model: str = "ernie-text-embedding"
+    """Model to use."""
+    request_timeout: Optional[int] = 60
+    """How many seconds to wait for the server to send data before giving up."""
+
+    ernie_client_id: Optional[str] = None
+    ernie_client_secret: Optional[str] = None
+    """For raising deprecation warnings."""
+
+    @root_validator()
+    def validate_environment(cls, values: Dict) -> Dict:
+        values["aistudio_access_token"] = get_from_dict_or_env(
+            values,
+            "aistudio_access_token",
+            "EB_ACCESS_TOKEN",
+        )
+
+        try:
+            import erniebot
+
+            values["client"] = erniebot.Embedding
+        except ImportError:
+            raise ImportError(
+                "Could not import erniebot python package. Please install it with `pip install erniebot`."
+            )
+        return values
+
+    def embed_query(self, text: str) -> List[float]:
+        resp = self.embed_documents([text])
+        return resp[0]
+
+    async def aembed_query(self, text: str) -> List[float]:
+        embeddings = await self.aembed_documents([text])
+        return embeddings[0]
+
+    def embed_documents(self, texts: List[str]) -> List[List[float]]:
+        text_in_chunks = [texts[i : i + self.chunk_size] for i in range(0, len(texts), self.chunk_size)]
+        lst = []
+        for chunk in text_in_chunks:
+            resp = self.client.create(_config_=self._get_auth_config(), input=chunk, model=self.model)
+            lst.extend([res["embedding"] for res in resp["data"]])
+        return lst
+
+    async def aembed_documents(self, texts: List[str]) -> List[List[float]]:
+        text_in_chunks = [texts[i : i + self.chunk_size] for i in range(0, len(texts), self.chunk_size)]
+        lst = []
+        for chunk in text_in_chunks:
+            resp = await self.client.acreate(_config_=self._get_auth_config(), input=chunk, model=self.model)
+            for res in resp["data"]:
+                lst.extend([res["embedding"]])
+        return lst
+
+    def _get_auth_config(self) -> dict:
+        return {"api_type": "aistudio", "access_token": self.aistudio_access_token}

erniebot-agent/erniebot_agent/extensions/langchain/llms/__init__.py

@@ -0,0 +1 @@
+from .erniebot import ErnieBot

erniebot-agent/erniebot_agent/extensions/langchain/llms/erniebot.py

@@ -0,0 +1,239 @@
+from __future__ import annotations
+
+from typing import (
+    Any,
+    AsyncIterator,
+    Callable,
+    Dict,
+    Iterator,
+    List,
+    Mapping,
+    Optional,
+    Type,
+    Union,
+)
+
+from langchain.callbacks.manager import (
+    AsyncCallbackManagerForLLMRun,
+    CallbackManagerForLLMRun,
+)
+from langchain.llms.base import LLM, create_base_retry_decorator
+from langchain.llms.utils import enforce_stop_tokens
+from langchain.pydantic_v1 import Field, root_validator
+from langchain.schema.output import GenerationChunk
+from langchain.utils import get_from_dict_or_env
+
+
+class ErnieBot(LLM):
+    """ERNIE Bot large language models.
+
+    To use, you should have the ``erniebot`` python package installed, and the
+    environment variable ``EB_ACCESS_TOKEN`` set with your AI Studio access token.
+
+    Example:
+        .. code-block:: python
+
+            from erniebot_agent.extensions.langchain.llms import ErnieBot
+            erniebot = ErnieBot(model="ernie-bot")
+    """
+
+    client: Any = None
+    max_retries: int = 6
+    """Maximum number of retries to make when generating."""
+    aistudio_access_token: Optional[str] = None
+    """AI Studio access token."""
+    streaming: Optional[bool] = False
+    """Whether to stream the results or not."""
+    model: str = "ernie-bot"
+    """Model to use."""
+    top_p: Optional[float] = 0.8
+    """Parameter of nucleus sampling that affects the diversity of generated content."""
+    temperature: Optional[float] = 0.95
+    """Sampling temperature to use."""
+    penalty_score: Optional[float] = 1
+    """Penalty assigned to tokens that have been generated."""
+    request_timeout: Optional[int] = 60
+    """How many seconds to wait for the server to send data before giving up."""
+    model_kwargs: Dict[str, Any] = Field(default_factory=dict)
+    """Holds any model parameters valid for `create` call not explicitly specified."""
+
+    @property
+    def _default_params(self) -> Dict[str, Any]:
+        """Get the default parameters for calling ERNIE Bot API."""
+        normal_params = {
+            "model": self.model,
+            "top_p": self.top_p,
+            "temperature": self.temperature,
+            "penalty_score": self.penalty_score,
+            "request_timeout": self.request_timeout,
+        }
+        return {**normal_params, **self.model_kwargs}
+
+    @property
+    def _identifying_params(self) -> Dict[str, Any]:
+        return self._default_params
+
+    @property
+    def _invocation_params(self) -> Dict[str, Any]:
+        """Get the parameters used to invoke the model."""
+        auth_cfg: Dict[str, Optional[str]] = {
+            "api_type": "aistudio",
+            "access_token": self.aistudio_access_token,
+        }
+        return {**{"_config_": auth_cfg}, **self._default_params}
+
+    @property
+    def _llm_type(self) -> str:
+        """Return type of llm."""
+        return "erniebot"
+
+    @root_validator()
+    def validate_enviroment(cls, values: Dict) -> Dict:
+        values["aistudio_access_token"] = get_from_dict_or_env(
+            values,
+            "aistudio_access_token",
+            "EB_ACCESS_TOKEN",
+        )
+
+        try:
+            import erniebot
+
+            values["client"] = erniebot.ChatCompletion
+        except ImportError:
+            raise ImportError(
+                "Could not import erniebot python package. Please install it with `pip install erniebot`."
+            )
+        return values
+
+    def _call(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        if self.streaming:
+            text = ""
+            for chunk in self._stream(prompt, stop, run_manager, **kwargs):
+                text += chunk.text
+            return text
+        else:
+            params = self._invocation_params
+            params.update(kwargs)
+            params["messages"] = [self._build_user_message_from_prompt(prompt)]
+            params["stream"] = False
+            response = _create_completion_with_retry(self, run_manager=run_manager, **params)
+            text = response["result"]
+            if stop is not None:
+                text = enforce_stop_tokens(text, stop)
+            return text
+
+    async def _acall(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> str:
+        if self.streaming:
+            text = ""
+            async for chunk in self._astream(prompt, stop, run_manager, **kwargs):
+                text += chunk.text
+            return text
+        else:
+            params = self._invocation_params
+            params.update(kwargs)
+            params["messages"] = [self._build_user_message_from_prompt(prompt)]
+            params["stream"] = False
+            response = await _acreate_completion_with_retry(self, run_manager=run_manager, **params)
+            text = response["result"]
+            if stop is not None:
+                text = enforce_stop_tokens(text, stop)
+            return text
+
+    def _stream(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[CallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> Iterator[GenerationChunk]:
+        if stop is not None:
+            raise TypeError("Currently, `stop` is not supported when streaming is enabled.")
+        params = self._invocation_params
+        params.update(kwargs)
+        params["messages"] = [self._build_user_message_from_prompt(prompt)]
+        params["stream"] = True
+        for resp in _create_completion_with_retry(self, run_manager=run_manager, **params):
+            chunk = self._build_chunk_from_response(resp)
+            yield chunk
+            if run_manager:
+                run_manager.on_llm_new_token(chunk.text, chunk=chunk)
+
+    async def _astream(
+        self,
+        prompt: str,
+        stop: Optional[List[str]] = None,
+        run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[GenerationChunk]:
+        if stop is not None:
+            raise TypeError("Currently, `stop` is not supported when streaming is enabled.")
+        params = self._invocation_params
+        params.update(kwargs)
+        params["messages"] = [self._build_user_message_from_prompt(prompt)]
+        params["stream"] = True
+        async for resp in await _acreate_completion_with_retry(self, run_manager=run_manager, **params):
+            chunk = self._build_chunk_from_response(resp)
+            yield chunk
+            if run_manager:
+                await run_manager.on_llm_new_token(chunk.text, chunk=chunk)
+
+    def _build_chunk_from_response(self, response: Mapping[str, Any]) -> GenerationChunk:
+        return GenerationChunk(text=response["result"])
+
+    def _build_user_message_from_prompt(self, prompt: str) -> Dict[str, str]:
+        return {"role": "user", "content": prompt}
+
+
+def _create_completion_with_retry(
+    llm: ErnieBot,
+    run_manager: Optional[CallbackManagerForLLMRun] = None,
+    **kwargs: Any,
+) -> Any:
+    retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)
+
+    @retry_decorator
+    def _client_create(**kwargs: Any) -> Any:
+        return llm.client.create(**kwargs)
+
+    return _client_create(**kwargs)
+
+
+async def _acreate_completion_with_retry(
+    llm: ErnieBot,
+    run_manager: Optional[AsyncCallbackManagerForLLMRun] = None,
+    **kwargs: Any,
+) -> Any:
+    retry_decorator = _create_retry_decorator(llm, run_manager=run_manager)
+
+    @retry_decorator
+    async def _client_acreate(**kwargs: Any) -> Any:
+        return await llm.client.acreate(**kwargs)
+
+    return await _client_acreate(**kwargs)
+
+
+def _create_retry_decorator(
+    llm: ErnieBot,
+    run_manager: Optional[Union[AsyncCallbackManagerForLLMRun, CallbackManagerForLLMRun]] = None,
+) -> Callable[[Any], Any]:
+    import erniebot
+
+    errors: List[Type[BaseException]] = [
+        erniebot.errors.TimeoutError,
+        erniebot.errors.RequestLimitError,
+    ]
+    return create_base_retry_decorator(
+        error_types=errors, max_retries=llm.max_retries, run_manager=run_manager
+    )

erniebot-agent/erniebot_agent/file_io/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

erniebot-agent/erniebot_agent/file_io/base.py

@@ -0,0 +1,46 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+
+
+class File(metaclass=abc.ABCMeta):
+    def __init__(self, id: str, filename: str, created_at: int) -> None:
+        super().__init__()
+        self.id = id
+        self.filename = filename
+        self.created_at = created_at
+
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, File):
+            return self.id == other.id
+        else:
+            return False
+
+    def __repr__(self) -> str:
+        attrs_str = self._get_attrs_str()
+        return f"<{self.__class__.__name__} {attrs_str}>"
+
+    @abc.abstractmethod
+    async def read_contents(self) -> bytes:
+        raise NotImplementedError
+
+    def _get_attrs_str(self) -> str:
+        return ", ".join(
+            [
+                f"id: {repr(self.id)}",
+                f"filename: {repr(self.filename)}",
+                f"created_at: {repr(self.created_at)}",
+            ]
+        )

erniebot-agent/erniebot_agent/file_io/file_manager.py

@@ -0,0 +1,138 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import pathlib
+import uuid
+from typing import Literal, Optional, Union, overload
+
+import anyio
+from erniebot_agent.file_io.base import File
+from erniebot_agent.file_io.file_registry import FileRegistry, get_file_registry
+from erniebot_agent.file_io.local_file import LocalFile, create_local_file_from_path
+from erniebot_agent.file_io.remote_file import RemoteFile, RemoteFileClient
+from erniebot_agent.utils.temp_file import create_tracked_temp_dir
+from typing_extensions import TypeAlias
+
+_PathType: TypeAlias = Union[str, os.PathLike]
+
+
+class FileManager(object):
+    _remote_file_client: Optional[RemoteFileClient]
+
+    def __init__(
+        self,
+        remote_file_client: Optional[RemoteFileClient] = None,
+        *,
+        auto_register: bool = True,
+        save_dir: Optional[_PathType] = None,
+    ) -> None:
+        super().__init__()
+        if remote_file_client is not None:
+            self._remote_file_client = remote_file_client
+        else:
+            self._remote_file_client = None
+        self._auto_register = auto_register
+        if save_dir is not None:
+            self._save_dir = pathlib.Path(save_dir)
+        else:
+            # This can be done lazily, but we need to be careful about race conditions.
+            self._save_dir = create_tracked_temp_dir()
+
+        self._file_registry = get_file_registry()
+
+    @property
+    def registry(self) -> FileRegistry:
+        return self._file_registry
+
+    @property
+    def remote_file_client(self) -> RemoteFileClient:
+        if self._remote_file_client is None:
+            raise AttributeError("No remote file client is set.")
+        else:
+            return self._remote_file_client
+
+    @overload
+    async def create_file_from_path(
+        self, file_path: _PathType, *, file_type: Literal["local"] = ...
+    ) -> LocalFile:
+        ...
+
+    @overload
+    async def create_file_from_path(
+        self, file_path: _PathType, *, file_type: Literal["remote"]
+    ) -> RemoteFile:
+        ...
+
+    async def create_file_from_path(
+        self, file_path: _PathType, *, file_type: Literal["local", "remote"] = "local"
+    ) -> Union[LocalFile, RemoteFile]:
+        file: Union[LocalFile, RemoteFile]
+        if file_type == "local":
+            file = await self.create_local_file_from_path(file_path)
+        elif file_type == "remote":
+            file = await self.create_remote_file_from_path(file_path)
+        else:
+            raise ValueError(f"Unsupported file type: {file_type}")
+        return file
+
+    async def create_local_file_from_path(self, file_path: _PathType) -> LocalFile:
+        file = create_local_file_from_path(pathlib.Path(file_path))
+        self._file_registry.register_file(file)
+        return file
+
+    async def create_remote_file_from_path(self, file_path: _PathType) -> RemoteFile:
+        file = await self.remote_file_client.upload_file(pathlib.Path(file_path))
+        if self._auto_register:
+            self._file_registry.register_file(file)
+        return file
+
+    @overload
+    async def create_file_from_bytes(
+        self, file_contents: bytes, filename: str, *, file_type: Literal["local"] = ...
+    ) -> LocalFile:
+        ...
+
+    @overload
+    async def create_file_from_bytes(
+        self, file_contents: bytes, filename: str, *, file_type: Literal["remote"]
+    ) -> RemoteFile:
+        ...
+
+    async def create_file_from_bytes(
+        self, file_contents: bytes, filename: str, *, file_type: Literal["local", "remote"] = "local"
+    ) -> Union[LocalFile, RemoteFile]:
+        # Can we do this with in-memory files?
+        file_path = self._fs_create_file(
+            prefix=pathlib.PurePath(filename).stem, suffix=pathlib.PurePath(filename).suffix
+        )
+        async with await anyio.open_file(file_path, "wb") as f:
+            await f.write(file_contents)
+        file = await self.create_file_from_path(file_path, file_type=file_type)
+        return file
+
+    async def retrieve_remote_file_by_id(self, file_id: str) -> RemoteFile:
+        file = await self.remote_file_client.retrieve_file(file_id)
+        if self._auto_register:
+            self._file_registry.register_file(file)
+        return file
+
+    def look_up_file_by_id(self, file_id: str) -> Optional[File]:
+        return self._file_registry.look_up_file(file_id)
+
+    def _fs_create_file(self, prefix: Optional[str] = None, suffix: Optional[str] = None) -> pathlib.Path:
+        filename = f"{prefix or ''}{str(uuid.uuid4())}{suffix or ''}"
+        file_path = self._save_dir / filename
+        file_path.touch()
+        return file_path

erniebot-agent/erniebot_agent/file_io/file_registry.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import threading
+from typing import Dict, List, Optional
+
+from erniebot_agent.file_io.base import File
+from erniebot_agent.utils.misc import Singleton
+
+
+class FileRegistry(metaclass=Singleton):
+    def __init__(self) -> None:
+        super().__init__()
+        self._id_to_file: Dict[str, File] = {}
+        self._lock = threading.Lock()
+
+    def register_file(self, file: File) -> None:
+        file_id = file.id
+        with self._lock:
+            # Re-registering an existing file is allowed.
+            # We simply update the registry.
+            self._id_to_file[file_id] = file
+
+    def unregister_file(self, file: File) -> None:
+        file_id = file.id
+        with self._lock:
+            if file_id not in self._id_to_file:
+                raise RuntimeError(f"ID {repr(file_id)} is not registered.")
+            self._id_to_file.pop(file_id)
+
+    def look_up_file(self, file_id: str) -> Optional[File]:
+        with self._lock:
+            return self._id_to_file.get(file_id, None)
+
+    def list_files(self) -> List[File]:
+        with self._lock:
+            return list(self._id_to_file.values())
+
+
+_file_registry = FileRegistry()
+
+
+def get_file_registry() -> FileRegistry:
+    return _file_registry

erniebot-agent/erniebot_agent/file_io/local_file.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pathlib
+import time
+import uuid
+
+import anyio
+from erniebot_agent.file_io.base import File
+from erniebot_agent.file_io.protocol import (
+    build_local_file_id_from_uuid,
+    is_local_file_id,
+)
+from erniebot_agent.utils.logging import logger
+
+
+class LocalFile(File):
+    def __init__(self, id: str, filename: str, created_at: int, path: pathlib.Path) -> None:
+        if not is_local_file_id(id):
+            raise ValueError("Invalid file ID: {id}")
+        super().__init__(id=id, filename=filename, created_at=created_at)
+        self.path = path
+
+    async def read_contents(self) -> bytes:
+        return await anyio.Path(self.path).read_bytes()
+
+    def _get_attrs_str(self) -> str:
+        attrs_str = super()._get_attrs_str()
+        attrs_str += f", path: {repr(self.path)}"
+        return attrs_str
+
+
+def create_local_file_from_path(file_path: pathlib.Path) -> LocalFile:
+    if not file_path.exists():
+        logger.warn("File %s does not exist.", file_path)
+    file_id = _generate_local_file_id()
+    filename = file_path.name
+    created_at = int(time.time())
+    file = LocalFile(id=file_id, filename=filename, created_at=created_at, path=file_path)
+    return file
+
+
+def _generate_local_file_id():
+    return build_local_file_id_from_uuid(str(uuid.uuid1()))

erniebot-agent/erniebot_agent/file_io/protocol.py

@@ -0,0 +1,57 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import re
+from typing import List
+
+_LOCAL_FILE_ID_PREFIX = "file-local-"
+_REMOTE_FILE_ID_PREFIX = "file-remote-"
+_UUID_PATTERN = r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}"
+_LOCAL_FILE_ID_PATTERN = _LOCAL_FILE_ID_PREFIX + _UUID_PATTERN
+_REMOTE_FILE_ID_PATTERN = _REMOTE_FILE_ID_PREFIX + _UUID_PATTERN
+
+_compiled_local_file_id_pattern = re.compile(_LOCAL_FILE_ID_PATTERN)
+_compiled_remote_file_id_pattern = re.compile(_REMOTE_FILE_ID_PATTERN)
+
+
+def build_local_file_id_from_uuid(uuid: str) -> str:
+    return _LOCAL_FILE_ID_PREFIX + uuid
+
+
+def build_remote_file_id_from_uuid(uuid: str) -> str:
+    return _REMOTE_FILE_ID_PREFIX + uuid
+
+
+def is_file_id(str_: str) -> bool:
+    return is_local_file_id(str_) or is_remote_file_id(str_)
+
+
+def is_local_file_id(str_: str) -> bool:
+    return _compiled_local_file_id_pattern.fullmatch(str_) is not None
+
+
+def is_remote_file_id(str_: str) -> bool:
+    return _compiled_remote_file_id_pattern.fullmatch(str_) is not None
+
+
+def extract_file_ids(str_: str) -> List[str]:
+    return extract_local_file_ids(str_) + extract_remote_file_ids(str_)
+
+
+def extract_local_file_ids(str_: str) -> List[str]:
+    return _compiled_local_file_id_pattern.findall(str_)
+
+
+def extract_remote_file_ids(str_: str) -> List[str]:
+    return _compiled_remote_file_id_pattern.findall(str_)

erniebot-agent/erniebot_agent/file_io/remote_file.py

@@ -0,0 +1,153 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+import asyncio
+import functools
+import pathlib
+import time
+import uuid
+from typing import ClassVar, Dict, List
+
+import anyio
+from baidubce.auth.bce_credentials import BceCredentials
+from baidubce.bce_client_configuration import BceClientConfiguration
+from baidubce.services.bos.bos_client import BosClient
+from erniebot_agent.file_io.base import File
+from erniebot_agent.file_io.protocol import (
+    build_remote_file_id_from_uuid,
+    is_remote_file_id,
+)
+
+
+class RemoteFile(File):
+    def __init__(self, id: str, filename: str, created_at: int, client: "RemoteFileClient") -> None:
+        if not is_remote_file_id(id):
+            raise ValueError("Invalid file ID: {id}")
+        super().__init__(id=id, filename=filename, created_at=created_at)
+        self._client = client
+
+    async def read_contents(self) -> bytes:
+        file_contents = await self._client.retrieve_file_contents(self.id)
+        return file_contents
+
+    async def delete(self) -> None:
+        await self._client.delete_file(self.id)
+
+
+class RemoteFileClient(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    async def upload_file(self, file_path: pathlib.Path) -> RemoteFile:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def retrieve_file(self, file_id: str) -> RemoteFile:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def retrieve_file_contents(self, file_id: str) -> bytes:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def list_files(self) -> List[RemoteFile]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    async def delete_file(self, file_id: str) -> None:
+        raise NotImplementedError
+
+
+class BOSFileClient(RemoteFileClient):
+    _ENDPOINT: ClassVar[str] = "bj.bcebos.com"
+
+    def __init__(self, ak: str, sk: str, bucket_name: str, prefix: str) -> None:
+        super().__init__()
+        self.bucket_name = bucket_name
+        self.prefix = prefix
+        config = BceClientConfiguration(credentials=BceCredentials(ak, sk), endpoint=self._ENDPOINT)
+        self._bos_client = BosClient(config=config)
+
+    async def upload_file(self, file_path: pathlib.Path) -> RemoteFile:
+        file_id = self._generate_file_id()
+        filename = file_path.name
+        created_at = int(time.time())
+        user_metadata: Dict[str, str] = {"id": file_id, "filename": filename, "created_at": str(created_at)}
+        async with await anyio.open_file(file_path, mode="rb") as f:
+            data = await f.read()
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(
+            None,
+            functools.partial(
+                self._bos_client.put_object_from_string,
+                bucket=self.bucket_name,
+                key=self._get_key(file_id),
+                data=data,
+                user_metadata=user_metadata,
+            ),
+        )
+        return RemoteFile(
+            id=file_id,
+            filename=filename,
+            created_at=created_at,
+            client=self,
+        )
+
+    async def retrieve_file(self, file_id: str) -> RemoteFile:
+        loop = asyncio.get_running_loop()
+        response = await loop.run_in_executor(
+            None,
+            functools.partial(
+                self._bos_client.get_object_meta_data, self.bucket_name, self._get_key(file_id)
+            ),
+        )
+        user_metadata = {
+            "id": response.metadata.bce_meta_id,
+            "filename": response.metadata.bce_meta_filename,
+            "created_at": int(response.metadata.bce_meta_created_at),
+        }
+        if file_id != user_metadata["id"]:
+            raise RuntimeError("`file_id` is not the same as the one in metadata.")
+
+        return RemoteFile(
+            id=user_metadata["id"],
+            filename=user_metadata["filename"],
+            created_at=user_metadata["created_at"],
+            client=self,
+        )
+
+    async def retrieve_file_contents(self, file_id: str) -> bytes:
+        loop = asyncio.get_running_loop()
+        result = await loop.run_in_executor(
+            None,
+            functools.partial(
+                self._bos_client.get_object_as_string, self.bucket_name, self._get_key(file_id)
+            ),
+        )
+        return result
+
+    async def list_files(self) -> List[RemoteFile]:
+        raise RuntimeError(f"`{self.__class__.__name__}.list_files` is not supported.")
+
+    async def delete_file(self, file_id: str) -> None:
+        loop = asyncio.get_running_loop()
+        await loop.run_in_executor(
+            None, functools.partial(self._bos_client.delete_object, self.bucket_name, self._get_key(file_id))
+        )
+
+    def _get_key(self, file_id: str) -> str:
+        return self.prefix + file_id
+
+    @staticmethod
+    def _generate_file_id() -> str:
+        return build_remote_file_id_from_uuid(str(uuid.uuid1()))

erniebot-agent/erniebot_agent/memory/__init__.py

@@ -0,0 +1,18 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .base import Memory
+from .limit_token_memory import LimitTokensMemory
+from .sliding_window_memory import SlidingWindowMemory
+from .whole_memory import WholeMemory

erniebot-agent/erniebot_agent/memory/base.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import List, Optional, Union
+
+from erniebot_agent.messages import AIMessage, Message, SystemMessage
+
+
+class MessageManager:
+    """
+    Messages Manager.
+    """
+
+    def __init__(self) -> None:
+        self.messages: List[Message] = []
+        self._system_message: Union[SystemMessage, None] = None
+
+    @property
+    def system_message(self) -> Optional[Message]:
+        """
+        The message manager have only one system message.
+
+        return: Message or None
+        """
+        return self._system_message
+
+    @system_message.setter
+    def system_message(self, message: SystemMessage) -> None:
+        if self._system_message is not None:
+            Warning("system message has been set, the previous one will be replaced")
+
+        self._system_message = message
+
+    def add_messages(self, messages: List[Message]) -> None:
+        self.messages.extend(messages)
+
+    def add_message(self, message: Message) -> None:
+        if isinstance(message, SystemMessage):
+            self.system_message = message
+        else:
+            self.messages.append(message)
+
+    def pop_message(self) -> Message:
+        return self.messages.pop(0)
+
+    def clear_messages(self) -> None:
+        self.messages = []
+
+    def update_last_message_token_count(self, token_count: int):
+        if token_count == 0:
+            self.messages[-1].token_count = len(self.messages[-1].content)
+        else:
+            self.messages[-1].token_count = token_count
+
+    def retrieve_messages(self) -> List[Message]:
+        return self.messages
+
+
+class Memory:
+    """The base class of memory"""
+
+    def __init__(self):
+        self.msg_manager = MessageManager()
+
+    def add_messages(self, messages: List[Message]):
+        for message in messages:
+            self.add_message(message)
+
+    def add_message(self, message: Message):
+        if isinstance(message, AIMessage):
+            self.msg_manager.update_last_message_token_count(message.query_tokens_count)
+        self.msg_manager.add_message(message)
+
+    def get_messages(self) -> List[Message]:
+        return self.msg_manager.retrieve_messages()
+
+    def get_system_message(self) -> SystemMessage:
+        return self.msg_manager.system_message
+
+    def clear_chat_history(self):
+        self.msg_manager.clear_messages()
+
+
+class WholeMemory(Memory):
+    """The memory include all the messages"""
+
+    def __init__(self):
+        super().__init__()

erniebot-agent/erniebot_agent/memory/limit_token_memory.py

@@ -0,0 +1,59 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+from erniebot_agent.memory import Memory
+from erniebot_agent.messages import AIMessage, Message
+
+
+class LimitTokensMemory(Memory):
+    """This class controls max tokens less than max_token_limit.
+    If tokens >= max_token_limit, pop message from memory.
+    """
+
+    def __init__(self, max_token_limit=6000):
+        super().__init__()
+        self.max_token_limit = max_token_limit
+        self.mem_token_count = 0
+
+        assert (
+            max_token_limit is None
+        ) or max_token_limit > 0, "max_token_limit should be None or positive integer, \
+                but got {max_token_limit}".format(
+            max_token_limit=max_token_limit
+        )
+
+    def add_message(self, message: Message):
+        super().add_message(message)
+        # TODO(shiyutang): 仅在添加AIMessage时截断会导致HumanMessage传入到LLM时可能长度超限
+        # 最优方案为每条message产生时确定token_count，从而在每次加入message时都进行prune_message
+        if isinstance(message, AIMessage):
+            self.prune_message()
+
+    def prune_message(self):
+        self.mem_token_count += self.msg_manager.messages[-1].token_count
+        self.mem_token_count += self.msg_manager.messages[-2].token_count  # add human message token length
+        if self.max_token_limit is not None:
+            while self.mem_token_count > self.max_token_limit:
+                deleted_message = self.msg_manager.pop_message()
+                self.mem_token_count -= deleted_message.token_count
+            else:
+                # if delete all
+                if len(self.get_messages()) == 0:
+                    raise RuntimeError(
+                        "The messsage is now empty. \
+                            It indicates {} which takes up {} tokens and exeeded {} tokens.".format(
+                            deleted_message, len(deleted_message.content), self.max_token_limit
+                        )
+                    )

erniebot-agent/erniebot_agent/memory/sliding_window_memory.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from erniebot_agent.memory import Memory
+from erniebot_agent.messages import Message
+
+
+class SlidingWindowMemory(Memory):
+    """This class controls max number of messages."""
+
+    def __init__(self, max_num_message: int):
+        super().__init__()
+        self.max_num_message = max_num_message
+
+        assert (isinstance(max_num_message, int)) and (
+            max_num_message > 0
+        ), "max_num_message should be positive integer, but got {max_token_limit}".format(
+            max_token_limit=max_num_message
+        )
+
+    def add_message(self, message: Message):
+        super().add_message(message=message)
+        self.prune_message()
+
+    def prune_message(self):
+        while len(self.get_messages()) > self.max_num_message:
+            self.msg_manager.pop_message()
+            # `messages` must have an odd number of elements.
+            if len(self.get_messages()) % 2 == 0:
+                self.msg_manager.pop_message()

erniebot-agent/erniebot_agent/memory/whole_memory.py

@@ -0,0 +1,19 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from erniebot_agent.memory.base import Memory
+
+
+class WholeMemory(Memory):
+    """"""

erniebot-agent/erniebot_agent/messages.py

@@ -0,0 +1,124 @@
+#
+# Licensed under the Apache License, Version 2.0 (the "License"
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License
+
+from typing import Dict, List, Optional, TypedDict
+
+import erniebot.utils.token_helper as token_helper
+
+
+class Message:
+    """The base class of a message."""
+
+    def __init__(self, role: str, content: str, token_count: Optional[int] = None):
+        self.role = role
+        self.content = content
+        self._token_count = token_count
+        self._param_names = ["role", "content"]
+
+    @property
+    def token_count(self):
+        """Get the number of tokens of the message."""
+        if self._token_count is None:
+            raise AttributeError("The token count of the message has not been set.")
+        return self._token_count
+
+    @token_count.setter
+    def token_count(self, token_count: int):
+        """Set the number of tokens of the message."""
+        if self._token_count is not None:
+            raise AttributeError("The token count of the message can only be set once.")
+        self._token_count = token_count
+
+    def to_dict(self) -> Dict[str, str]:
+        res = {}
+        for name in self._param_names:
+            res[name] = getattr(self, name)
+        return res
+
+    def __str__(self) -> str:
+        return f"<{self._get_attrs_str()}>"
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__} {self._get_attrs_str()}>"
+
+    def _get_attrs_str(self) -> str:
+        parts: List[str] = []
+        for name in self._param_names:
+            value = getattr(self, name)
+            if value is not None and value != "":
+                parts.append(f"{name}: {repr(value)}")
+        if self._token_count is not None:
+            parts.append(f"token_count: {self._token_count}")
+        return ", ".join(parts)
+
+
+class SystemMessage(Message):
+    """A message from a human to set system information."""
+
+    def __init__(self, content: str):
+        super().__init__(role="system", content=content, token_count=len(content))
+
+
+class HumanMessage(Message):
+    """A message from a human."""
+
+    def __init__(self, content: str):
+        super().__init__(role="user", content=content)
+
+
+class FunctionCall(TypedDict):
+    name: str
+    thoughts: str
+    arguments: str
+
+
+class TokenUsage(TypedDict):
+    prompt_tokens: int
+    completion_tokens: int
+
+
+class AIMessage(Message):
+    """A message from an assistant."""
+
+    def __init__(
+        self,
+        content: str,
+        function_call: Optional[FunctionCall],
+        token_usage: Optional[TokenUsage] = None,
+    ):
+        if token_usage is None:
+            prompt_tokens = 0
+            completion_tokens = token_helper.approx_num_tokens(content)
+        else:
+            prompt_tokens, completion_tokens = self._parse_token_count(token_usage)
+        super().__init__(role="assistant", content=content, token_count=completion_tokens)
+        self.function_call = function_call
+        self.query_tokens_count = prompt_tokens
+        self._param_names = ["role", "content", "function_call"]
+
+    def _parse_token_count(self, token_usage: TokenUsage):
+        """Parse the token count information from LLM."""
+        return token_usage["prompt_tokens"], token_usage["completion_tokens"]
+
+
+class FunctionMessage(Message):
+    """A message from a human, containing the result of a function call."""
+
+    def __init__(self, name: str, content: str):
+        super().__init__(role="function", content=content)
+        self.name = name
+        self._param_names = ["role", "name", "content"]
+
+
+class AIMessageChunk(AIMessage):
+    """A message chunk from an assistant."""

erniebot-agent/erniebot_agent/prompt/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .base import BasePromptTemplate
+from .prompt_template import PromptTemplate

erniebot-agent/erniebot_agent/prompt/base.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from abc import ABC, abstractmethod
+from typing import List, Optional
+
+
+class BasePromptTemplate(ABC):
+    def __init__(self, input_variables: Optional[List[str]]):
+        self.input_variables: Optional[List[str]] = input_variables
+
+    @abstractmethod
+    def format(self, **kwargs):
+        raise NotImplementedError
+
+    @abstractmethod
+    def format_as_message(self, message_class, **kwargs):
+        raise NotImplementedError

erniebot-agent/erniebot_agent/prompt/prompt_template.py

@@ -0,0 +1,80 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, List, Optional
+
+from erniebot_agent.messages import HumanMessage
+from erniebot_agent.prompt import BasePromptTemplate
+from jinja2 import Environment, meta
+
+
+def jinja2_formatter(template: str, **kwargs: Any) -> str:
+    """Format a template using jinja2."""
+    try:
+        from jinja2 import Template
+    except ImportError:
+        raise ImportError(
+            "jinja2 not installed, which is needed to use the jinja2_formatter. "
+            "Please install it with `pip install jinja2`."
+        )
+
+    return Template(template).render(**kwargs)
+
+
+class PromptTemplate(BasePromptTemplate):
+    """format the prompt for llm input."""
+
+    def __init__(
+        self, template: str, name: Optional[str] = None, input_variables: Optional[List[str]] = None
+    ):
+        super().__init__(input_variables)
+        self.name = name
+        self.template = template
+        self.validate_template = True if input_variables is not None else False  # todo: 验证模板是否正确
+
+    def format(self, **kwargs) -> str:
+        if self.validate_template:
+            error = self._validate_template()
+            if error:
+                raise KeyError("The input_variables of PromptTemplate and template are not match! " + error)
+        return jinja2_formatter(self.template, **kwargs)
+
+    def _validate_template(self):
+        """
+        Validate that the input variables are valid for the template.
+
+        Args:
+            template: The template string.
+            input_variables: The input variables.
+        """
+        input_variables_set = set(self.input_variables)
+        env = Environment()
+        ast = env.parse(self.template)
+        valid_variables = meta.find_undeclared_variables(ast)
+
+        missing_variables = valid_variables - input_variables_set
+        extra_variables = input_variables_set - valid_variables
+
+        Error_message = ""
+        if missing_variables:
+            Error_message += f"The missing input variables: {missing_variables} "
+
+        if extra_variables:
+            Error_message += f"The extra input variables: {extra_variables}"
+
+        return Error_message
+
+    def format_as_message(self, **kwargs):
+        prompt = self.format(**kwargs)
+        return HumanMessage(content=prompt)

erniebot-agent/erniebot_agent/retrieval/baizhong_search.py

@@ -0,0 +1,296 @@
+import base64
+import json
+from concurrent.futures import ThreadPoolExecutor
+from typing import Any, Dict, List, Optional
+
+import requests
+from erniebot_agent.utils.exception import BaizhongError
+from erniebot_agent.utils.logging import logger
+from tqdm import tqdm
+
+from .document import Document
+
+
+class BaizhongSearch:
+    def __init__(
+        self,
+        base_url: str,
+        project_name: Optional[str] = None,
+        remark: Optional[str] = None,
+        project_id: Optional[int] = None,
+        max_seq_length: int = 512,
+    ) -> None:
+        self.base_url = base_url
+        self.max_seq_length = max_seq_length
+        if project_id is not None:
+            logger.info(f"Loading existing project with `project_id={project_id}`")
+            self.project_id = project_id
+        elif project_name is not None:
+            logger.info("Creating new project and schema")
+            self.index = self.create_project(project_name, remark)
+            logger.info("Project creation succeeded")
+            self.project_id = self.index["result"]["projectId"]
+            self.create_schema()
+            logger.info("Schema creation succeeded")
+        else:
+            raise BaizhongError("You must provide either a `project_name` or a `project_id`.")
+
+    def create_project(self, project_name: str, remark: Optional[str] = None):
+        """
+        Create a project using the Baizhong API.
+
+        Returns:
+            dict: A dictionary containing information about the created project.
+
+        Raises:
+            BaizhongError: If the API request fails, this exception is raised with details about the error.
+        """
+        json_data = {
+            "projectName": project_name,
+            "remark": remark,
+        }
+        res = requests.post(f"{self.base_url}/baizhong/web-api/v2/project/add", json=json_data)
+        if res.status_code == 200:
+            result = res.json()
+            if result["errCode"] != 0:
+                raise BaizhongError(message=result["errMsg"], error_code=result["errCode"])
+            return result
+        else:
+            raise BaizhongError(message=f"request error: {res.text}", error_code=res.status_code)
+
+    def create_schema(self):
+        """
+        Create a schema for a project using the Baizhong API.
+
+        Returns:
+            dict: A dictionary containing information about the created schema.
+
+        Raises:
+            BaizhongError: If the API request fails, this exception is raised with details about the error.
+        """
+        json_data = {
+            "projectId": self.project_id,
+            "schemaJson": {
+                "paraSize": self.max_seq_length,
+                "dataSegmentationMod": "neisou",
+                "storeType": "ElasticSearch",
+                "properties": {
+                    "title": {"type": "text", "shortindex": True},
+                    "content_se": {"type": "text", "longindex": True},
+                },
+            },
+        }
+        res = requests.post(f"{self.base_url}/baizhong/web-api/v2/project-schema/create", json=json_data)
+        if res.status_code == 200:
+            result = res.json()
+            if result["errCode"] != 0:
+                raise BaizhongError(message=result["errMsg"], error_code=result["errCode"])
+            return res.json()
+        else:
+            raise BaizhongError(message=f"request error: {res.text}", error_code=res.status_code)
+
+    def update_schema(
+        self,
+    ):
+        """
+        Update the schema for a project using the Baizhong API.
+
+        Returns:
+            dict: A dictionary containing information about the updated schema.
+
+        Raises:
+            BaizhongError: If the API request fails, this exception is raised with details about the error.
+        """
+        json_data = {
+            "projectId": self.project_id,
+            "schemaJson": {
+                "paraSize": self.max_seq_length,
+                "dataSegmentationMod": "neisou",
+                "storeType": "ElasticSearch",
+                "properties": {
+                    "title": {"type": "text", "shortindex": True},
+                    "content_se": {"type": "text", "longindex": True},
+                },
+            },
+        }
+        res = requests.post(f"{self.base_url}/baizhong/web-api/v2/project-schema/update", json=json_data)
+        status_code = res.status_code
+        if status_code == 200:
+            result = res.json()
+            if result["errCode"] != 0:
+                raise BaizhongError(message=result["errMsg"], error_code=result["errCode"])
+            return result
+        else:
+            raise BaizhongError(message=f"request error: {res.text}", error_code=res.status_code)
+
+    def search(self, query: str, top_k: int = 10, filters: Optional[Dict[str, Any]] = None):
+        """
+        Perform a search using the Baizhong common search API.
+
+        Args:
+            query (str): The search query.
+            top_k (int, optional): The number of top results to retrieve (default is 10).
+            filters (Optional[Dict[str, Any]], optional): Additional filters to apply to the search query
+            (default is None).
+
+        Returns:
+            List[Dict[str, Any]]: A list of dictionaries containing search results.
+
+        Raises:
+            BaizhongError: If the API request fails, this exception is raised with details about the error.
+        """
+        json_data = {
+            "query": query,
+            "projectId": self.project_id,
+            "size": top_k,
+        }
+        if filters is not None:
+            filterConditions = {"filterConditions": {"bool": {"filter": {"match": filters}}}}
+            json_data.update(filterConditions)
+        res = requests.post(f"{self.base_url}/baizhong/common-search/v2/search", json=json_data)
+        if res.status_code == 200:
+            result = res.json()
+            if result["errCode"] != 0:
+                raise BaizhongError(message=result["errMsg"], error_code=result["errCode"])
+            list_data = []
+            for item in result["hits"]:
+                content = item["_source"]["doc"]
+                content = base64.b64decode(content).decode("utf-8")
+                json_data = json.loads(content)
+                list_data.append(json_data)
+            return list_data
+        else:
+            raise BaizhongError(message=f"request error: {res.text}", error_code=res.status_code)
+
+    def add_documents(self, documents: List[Document], batch_size: int = 1, thread_count: int = 1):
+        """
+        Add a batch of documents to the Baizhong system using multi-threading.
+
+        Args:
+            documents (List[Document]): A list of Document objects to be added.
+            batch_size (int, optional): The size of each batch of documents (defaults to 1).
+            thread_count (int, optional): The number of threads to use for concurrent document addition
+            (defaults to 1).
+
+        Returns:
+            List[Union[None, Exception]]: A list of results from the document addition process.
+
+        Note:
+            This function uses multi-threading to improve the efficiency of adding a large number of
+            documents.
+
+        """
+        if type(documents[0]) == Document:
+            list_dicts = [item.to_dict() for item in documents]
+        all_data = []
+        for i in tqdm(range(0, len(list_dicts), batch_size)):
+            batch_data = list_dicts[i : i + batch_size]
+            all_data.append(batch_data)
+        with ThreadPoolExecutor(max_workers=thread_count) as executor:
+            res = executor.map(self._add_documents, all_data)
+        return list(res)
+
+    def get_document_by_id(self, doc_id):
+        """
+        Retrieve a document from the Baizhong system by its ID.
+
+        Args:
+            doc_id: The ID of the document to retrieve.
+
+        Returns:
+            dict: A dictionary containing information about the retrieved document.
+
+        Raises:
+            BaizhongError: If the API request fails, this exception is raised with details about the error.
+        """
+        json_data = {"projectId": self.project_id, "followIndexFlag": True, "dataBody": [doc_id]}
+        res = requests.post(f"{self.base_url}/baizhong/data-api/v2/flush/get", json=json_data)
+        if res.status_code == 200:
+            result = res.json()
+            if result["errCode"] != 0:
+                raise BaizhongError(message=result["errMsg"], error_code=result["errCode"])
+            return result
+        else:
+            raise BaizhongError(message=f"request error: {res.text}", error_code=res.status_code)
+
+    def delete_documents(
+        self,
+        ids: Optional[List[str]] = None,
+    ):
+        """
+        Delete documents from the Baizhong system.
+
+        Args:
+            ids (Optional[List[str]], optional): A list of document IDs to delete. If not provided,
+            all documents will be deleted.
+
+        Returns:
+            dict: A dictionary containing information about the deletion process.
+
+        Raises:
+            NotImplementedError: If the deletion of all documents is attempted, this exception is raised
+            as it is not yet implemented.
+            BaizhongError: If the API request fails, this exception is raised with details about the error.
+        """
+        json_data: Dict[str, Any] = {"projectId": self.project_id, "followIndexFlag": True}
+        if ids is not None:
+            json_data["dataBody"] = ids
+        else:
+            # TODO: delete all documents
+            raise NotImplementedError
+        res = requests.post(f"{self.base_url}/baizhong/data-api/v2/flush/delete", json=json_data)
+        if res.status_code == 200:
+            result = res.json()
+            if result["errCode"] != 0:
+                raise BaizhongError(message=result["errMsg"], error_code=result["errCode"])
+            return result
+        else:
+            raise BaizhongError(message=f"request error: {res.text}", error_code=res.status_code)
+
+    def _add_documents(self, documents: List[Dict[str, Any]]):
+        """
+        Internal method to add a batch of documents to the Baizhong system.
+
+        Args:
+            documents (List[Dict[str, Any]]): A list of dictionaries representing documents to be added.
+
+        Returns:
+            dict: A dictionary containing information about the document addition process.
+
+        Raises:
+            BaizhongError: If the API request fails, this exception is raised with details about the error.
+        """
+        json_data = {"projectId": self.project_id, "followIndexFlag": True, "dataBody": documents}
+        res = requests.post(f"{self.base_url}/baizhong/data-api/v2/flush/add", json=json_data)
+        if res.status_code == 200:
+            result = res.json()
+            if result["errCode"] != 0:
+                raise BaizhongError(message=result["errMsg"], error_code=result["errCode"])
+            return result
+        else:
+            # TODO(wugaosheng): retry 3 times
+            raise BaizhongError(message=f"request error: {res.text}", error_code=res.status_code)
+
+    @classmethod
+    def delete_project(cls, project_id: int):
+        """
+        Class method to delete a project using the Baizhong API.
+
+        Args:
+            project_id (int): The ID of the project to be deleted.
+
+        Returns:
+            dict: A dictionary containing information about the deletion process.
+
+        Raises:
+            BaizhongError: If the API request fails, this exception is raised with details about the error.
+        """
+        json_data = {"projectId": project_id}
+        res = requests.post(f"{cls.base_url}/baizhong/web-api/v2/project/delete", json=json_data)
+        if res.status_code == 200:
+            result = res.json()
+            if result["errCode"] != 0:
+                raise BaizhongError(message=result["errMsg"], error_code=result["errCode"])
+            return res.json()
+        else:
+            raise BaizhongError(message=f"request error: {res.text}", error_code=res.status_code)

erniebot-agent/erniebot_agent/retrieval/document.py

@@ -0,0 +1,123 @@
+import hashlib
+import json
+from typing import Any, Dict, Optional, Union
+
+from pydantic import BaseConfig
+from pydantic.dataclasses import dataclass
+
+BaseConfig.arbitrary_types_allowed = True
+
+
+@dataclass
+class Document:
+    id: str
+    title: str
+    content_se: str
+    meta: Dict[str, Any]
+
+    def __init__(
+        self,
+        content_se: str,
+        title: str,
+        id: Optional[str] = None,
+        meta: Optional[Dict[str, Any]] = None,
+    ):
+        self.content_se = content_se
+        self.title = title
+        self.id = id or self._get_id()
+        self.meta = meta or {}
+
+    @classmethod
+    def _get_id(cls, content_se=None) -> str:
+        md5_bytes = content_se.encode(encoding="UTF-8")
+        md5_string = hashlib.md5(md5_bytes).hexdigest()
+        return md5_string
+
+    def to_dict(self, field_map: Optional[Dict[str, Any]] = None) -> Dict:
+        """
+        Convert Document to dict. An optional field_map can be supplied to
+        change the names of the keys in the resulting dict.
+        This way you can work with standardized Document objects in erniebot-agent,
+        but adjust the format that they are serialized / stored in other places
+        (e.g. elasticsearch)
+        Example:
+
+        ```python
+            doc = Document(content="some text", content_type="text")
+            doc.to_dict(field_map={"custom_content_field": "content"})
+
+            # Returns {"custom_content_field": "some text"}
+        ```
+
+        :param field_map: Dict with keys being the custom target keys and values
+        being the standard Document attributes
+        :return: dict with content of the Document
+        """
+        if not field_map:
+            field_map = {}
+
+        inv_field_map = {v: k for k, v in field_map.items()}
+        _doc: Dict[str, str] = {}
+        for k, v in self.__dict__.items():
+            # Exclude internal fields (Pydantic, ...) fields from the conversion process
+            if k.startswith("__"):
+                continue
+            k = k if k not in inv_field_map else inv_field_map[k]
+            _doc[k] = v
+        return _doc
+
+    @classmethod
+    def from_dict(cls, dict: Dict[str, Any], field_map: Optional[Dict[str, Any]] = None):
+        """
+        Create Document from dict. An optional `field_map` parameter
+        can be supplied to adjust for custom names of the keys in the
+        input dict. This way you can work with standardized Document
+        objects in erniebot-agent, but adjust the format that
+        they are serialized / stored in other places (e.g. elasticsearch).
+
+        Example:
+
+        ```python
+            my_dict = {"custom_content_field": "some text", "content_type": "text"}
+            Document.from_dict(my_dict, field_map={"custom_content_field": "content"})
+        ```
+
+        :param field_map: Dict with keys being the custom target keys and values
+        being the standard Document attributes
+        :return: A Document object
+        """
+        if not field_map:
+            field_map = {}
+
+        _doc = dict.copy()
+        init_args = ["content_se", "meta", "id", "title"]
+        if "meta" not in _doc.keys():
+            _doc["meta"] = {}
+        if "id" not in _doc.keys():
+            _doc["id"] = cls._get_id(_doc["content_se"])
+        # copy additional fields into "meta"
+        for k, v in _doc.items():
+            # Exclude internal fields (Pydantic, ...) fields from the conversion process
+            if k.startswith("__"):
+                continue
+            if k not in init_args and k not in field_map:
+                _doc["meta"][k] = v
+        # remove additional fields from top level
+        _new_doc = {}
+        for k, v in _doc.items():
+            if k in init_args:
+                _new_doc[k] = v
+            elif k in field_map:
+                k = field_map[k]
+                _new_doc[k] = v
+        return cls(**_new_doc)
+
+    @classmethod
+    def from_json(cls, data: Union[str, Dict[str, Any]], field_map: Optional[Dict[str, Any]] = None):
+        if not field_map:
+            field_map = {}
+        if isinstance(data, str):
+            dict_data = json.loads(data)
+        else:
+            dict_data = data
+        return cls.from_dict(dict_data, field_map=field_map)

erniebot-agent/erniebot_agent/tools/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .image_generation_tool import ImageGenerationTool

erniebot-agent/erniebot_agent/tools/baizhong_tool.py

@@ -0,0 +1,65 @@
+from __future__ import annotations
+
+from typing import Any, List, Optional, Type
+
+from erniebot_agent.messages import AIMessage, HumanMessage
+from erniebot_agent.tools.schema import ToolParameterView
+from pydantic import Field
+
+from .base import Tool
+
+
+class BaizhongSearchToolInputView(ToolParameterView):
+    query: str = Field(description="Query")
+    top_k: int = Field(description="Number of results to return")
+
+
+class SearchResponseDocument(ToolParameterView):
+    id: str = Field(description="text id")
+    title: str = Field(description="title")
+    document: str = Field(description="content")
+
+
+class BaizhongSearchToolOutputView(ToolParameterView):
+    documents: List[SearchResponseDocument] = Field(description="research results")
+
+
+class BaizhongSearchTool(Tool):
+    description: str = "aurora search tool"
+    input_type: Type[ToolParameterView] = BaizhongSearchToolInputView
+    ouptut_type: Type[ToolParameterView] = BaizhongSearchToolOutputView
+
+    def __init__(self, description, db, input_type=None, output_type=None, examples=None) -> None:
+        super().__init__()
+        self.db = db
+        self.description = description
+        if input_type is not None:
+            self.input_type = input_type
+        if output_type is not None:
+            self.ouptut_type = output_type
+        if examples is not None:
+            self.few_shot_examples = examples
+
+    async def __call__(self, query: str, top_k: int = 10, filters: Optional[dict[str, Any]] = None):
+        res = self.db.search(query, top_k, filters)
+        return res
+
+    @property
+    def examples(
+        self,
+    ) -> List[Any]:
+        few_shot_objects: List[Any] = []
+        for item in self.few_shot_examples:
+            few_shot_objects.append(HumanMessage(item["user"]))
+            few_shot_objects.append(
+                AIMessage(
+                    "",
+                    function_call={
+                        "name": self.tool_name,
+                        "thoughts": item["thoughts"],
+                        "arguments": item["arguments"],
+                    },
+                )
+            )
+
+        return few_shot_objects

erniebot-agent/erniebot_agent/tools/base.py

@@ -0,0 +1,428 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import json
+import os
+import tempfile
+from abc import ABC, abstractmethod
+from dataclasses import asdict, dataclass, field
+from typing import Any, Dict, List, Optional, Type
+
+import requests
+from erniebot_agent.messages import AIMessage, FunctionCall, HumanMessage, Message
+from erniebot_agent.tools.schema import (
+    Endpoint,
+    EndpointInfo,
+    RemoteToolView,
+    ToolParameterView,
+    scrub_dict,
+)
+from erniebot_agent.utils.http import url_file_exists
+from erniebot_agent.utils.logging import logger
+from openapi_spec_validator import validate
+from openapi_spec_validator.readers import read_from_filename
+from yaml import safe_dump
+
+import erniebot
+
+
+def validate_openapi_yaml(yaml_file: str) -> bool:
+    """do validation on the yaml file
+
+    Args:
+        yaml_file (str): the path of yaml file
+
+    Returns:
+        bool: whether yaml file is valid
+    """
+    yaml_dict = read_from_filename(yaml_file)[0]
+    try:
+        validate(yaml_dict)
+        return True
+    except Exception as e:  # type: ignore
+        logger.error(e)
+        return False
+
+
+class BaseTool(ABC):
+    @abstractmethod
+    async def __call__(self, *args: Any, **kwds: Any) -> Any:
+        raise NotImplementedError
+
+    @abstractmethod
+    def function_call_schema(self) -> dict:
+        raise NotImplementedError
+
+
+class Tool(BaseTool, ABC):
+    description: str
+    name: Optional[str] = None
+    input_type: Optional[Type[ToolParameterView]] = None
+    ouptut_type: Optional[Type[ToolParameterView]] = None
+
+    def __str__(self) -> str:
+        return "<name: {0}, description: {1}>".format(self.name, self.description)
+
+    def __repr__(self):
+        return self.__str__()
+
+    @property
+    def tool_name(self):
+        return self.name or self.__class__.__name__
+
+    @abstractmethod
+    async def __call__(self, *args: Any, **kwds: Any) -> Dict[str, Any]:
+        """the body of tools
+
+        Returns:
+            Any:
+        """
+        raise NotImplementedError
+
+    def function_call_schema(self) -> dict:
+        inputs = {
+            "name": self.tool_name,
+            "description": self.description,
+            "examples": [example.to_dict() for example in self.examples],
+        }
+        if self.input_type is not None:
+            inputs["parameters"] = self.input_type.function_call_schema()
+        if self.ouptut_type is not None:
+            inputs["responses"] = self.ouptut_type.function_call_schema()
+
+        return scrub_dict(inputs) or {}
+
+    @property
+    def examples(self) -> List[Message]:
+        return []
+
+
+class RemoteTool(BaseTool):
+    def __init__(
+        self,
+        tool_view: RemoteToolView,
+        server_url: str,
+        headers: dict,
+        examples: Optional[List[Message]] = None,
+    ) -> None:
+        self.tool_view = tool_view
+        self.server_url = server_url
+        self.headers = headers
+        self.examples = examples
+
+    def __str__(self) -> str:
+        return "<name: {0}, server_url: {1}, description: {2}>".format(
+            self.tool_name, self.server_url, self.tool_view.description
+        )
+
+    def __repr__(self):
+        return self.__str__()
+
+    @property
+    def tool_name(self):
+        return self.tool_view.name
+
+    async def __call__(self, **tool_arguments: Dict[str, Any]) -> Any:
+        url = self.server_url + self.tool_view.uri
+
+        if self.tool_view.method == "get":
+            response = requests.get(url, params=tool_arguments, headers=self.headers)
+        elif self.tool_view.method == "post":
+            response = requests.post(url, json=tool_arguments, headers=self.headers)
+        elif self.tool_view.method == "put":
+            response = requests.put(url, json=tool_arguments, headers=self.headers)
+        elif self.tool_view.method == "delete":
+            response = requests.delete(url, json=tool_arguments, headers=self.headers)
+        else:
+            raise ValueError(f"method<{self.tool_view.method}> is invalid")
+
+        if response.status_code != 200:
+            raise ValueError(f"the resource is invalid, the error message is: {response.text}")
+
+        return response.json()
+
+    def function_call_schema(self) -> dict:
+        schema = self.tool_view.function_call_schema()
+        if self.examples is not None:
+            schema["examples"] = [example.to_dict() for example in self.examples]
+
+        return schema or {}
+
+
+@dataclass
+class RemoteToolkit:
+    """RemoteToolkit can be converted by openapi.yaml and endpoint"""
+
+    openapi: str
+    info: EndpointInfo
+    servers: List[Endpoint]
+    paths: List[RemoteToolView]
+
+    component_schemas: dict[str, Type[ToolParameterView]]
+    headers: dict
+    examples: List[Message] = field(default_factory=list)
+
+    def __getitem__(self, tool_name: str):
+        return self.get_tool(tool_name)
+
+    def get_tools(self) -> List[RemoteTool]:
+        return [
+            RemoteTool(
+                path, self.servers[0].url, self.headers, examples=self.get_examples_by_name(path.name)
+            )
+            for path in self.paths
+        ]
+
+    def get_examples_by_name(self, tool_name: str) -> List[Message]:
+        """get examples by tool-name
+
+        Args:
+            tool_name (str): the name of the tool
+
+        Returns:
+            List[Message]: the messages
+        """
+        # 1. split messages
+        tool_examples: List[List[Message]] = []
+        examples: List[Message] = []
+        for example in self.examples:
+            if isinstance(example, HumanMessage):
+                if len(examples) == 0:
+                    examples.append(example)
+                else:
+                    tool_examples.append(examples)
+                    examples = [example]
+            else:
+                examples.append(example)
+
+        if len(examples) > 0:
+            tool_examples.append(examples)
+
+        final_exampels: List[Message] = []
+        # 2. find the target tool examples or empty messages
+        for examples in tool_examples:
+            tool_names = [
+                example.function_call.get("name", None)
+                for example in examples
+                if isinstance(example, AIMessage) and example.function_call is not None
+            ]
+            tool_names = [name for name in tool_names if name]
+
+            if tool_name in tool_names:
+                final_exampels.extend(examples)
+
+        return final_exampels
+
+    def get_tool(self, tool_name: str) -> RemoteTool:
+        paths = [path for path in self.paths if path.name == tool_name]
+        assert len(paths) == 1, f"tool<{tool_name}> not found in paths"
+        return RemoteTool(
+            paths[0], self.servers[0].url, self.headers, examples=self.get_examples_by_name(tool_name)
+        )
+
+    def to_openapi_dict(self) -> dict:
+        """convert plugin schema to openapi spec dict"""
+        spec_dict = {
+            "openapi": self.openapi,
+            "info": asdict(self.info),
+            "servers": [asdict(server) for server in self.servers],
+            "paths": {tool_view.uri: tool_view.to_openapi_dict() for tool_view in self.paths},
+            "components": {
+                "schemas": {
+                    uri: parameters_view.to_openapi_dict()
+                    for uri, parameters_view in self.component_schemas.items()
+                }
+            },
+        }
+        return scrub_dict(spec_dict, remove_empty_dict=True) or {}
+
+    def to_openapi_file(self, file: str):
+        """generate openapi configuration file
+
+        Args:
+            file (str): the path of the openapi yaml file
+        """
+        spec_dict = self.to_openapi_dict()
+        with open(file, "w+", encoding="utf-8") as f:
+            safe_dump(spec_dict, f, indent=4)
+
+    @classmethod
+    def from_openapi_dict(
+        cls, openapi_dict: Dict[str, Any], access_token: Optional[str] = None
+    ) -> RemoteToolkit:
+        info = EndpointInfo(**openapi_dict["info"])
+        servers = [Endpoint(**server) for server in openapi_dict.get("servers", [])]
+
+        # components
+        component_schemas = openapi_dict["components"]["schemas"]
+        fields = {}
+        for schema_name, schema in component_schemas.items():
+            parameter_view = ToolParameterView.from_openapi_dict(schema_name, schema)
+            fields[schema_name] = parameter_view
+
+        # paths
+        paths = []
+        for path, path_info in openapi_dict.get("paths", {}).items():
+            for method, path_method_info in path_info.items():
+                paths.append(
+                    RemoteToolView.from_openapi_dict(
+                        uri=path,
+                        method=method,
+                        path_info=path_method_info,
+                        parameters_views=fields,
+                    )
+                )
+
+        return RemoteToolkit(
+            openapi=openapi_dict["openapi"],
+            info=info,
+            servers=servers,
+            paths=paths,
+            component_schemas=fields,
+            headers=cls._get_authorization_headers(access_token),
+        )  # type: ignore
+
+    @classmethod
+    def from_openapi_file(cls, file: str, access_token: Optional[str] = None) -> RemoteToolkit:
+        """only support openapi v3.0.1
+
+        Args:
+            file (str): the path of openapi yaml file
+            access_token (Optional[str]): the path of openapi yaml file
+        """
+        if not validate_openapi_yaml(file):
+            raise ValueError(f"invalid openapi yaml file: {file}")
+
+        spec_dict, _ = read_from_filename(file)
+        return cls.from_openapi_dict(spec_dict, access_token=access_token)
+
+    @classmethod
+    def _get_authorization_headers(cls, access_token: Optional[str]) -> dict:
+        if access_token is None:
+            access_token = erniebot.access_token
+
+        headers = {"Content-Type": "application/json"}
+        if access_token is None:
+            logger.warning("access_token is NOT provided, this may cause 403 HTTP error..")
+        else:
+            headers["Authorization"] = f"token {access_token}"
+        return headers
+
+    @classmethod
+    def from_url(cls, url: str, access_token: Optional[str] = None) -> RemoteToolkit:
+        # 1. download openapy.yaml file to temp directory
+        if not url.endswith("/"):
+            url += "/"
+        openapi_yaml_url = url + ".well-known/openapi.yaml"
+
+        with tempfile.TemporaryDirectory() as temp_dir:
+            response = requests.get(openapi_yaml_url, headers=cls._get_authorization_headers(access_token))
+            if response.status_code != 200:
+                raise ValueError(f"the resource is invalid, the error message is: {response.text}")
+
+            file_content = response.content.decode("utf-8")
+            if not file_content.strip():
+                raise ValueError(f"the content is empty from: {openapi_yaml_url}")
+
+            file_path = os.path.join(temp_dir, "openapi.yaml")
+            with open(file_path, "w+", encoding="utf-8") as f:
+                f.write(file_content)
+
+            toolkit = RemoteToolkit.from_openapi_file(file_path, access_token=access_token)
+            for server in toolkit.servers:
+                server.url = url
+
+            toolkit.examples = cls.load_remote_examples_yaml(url, access_token)
+
+        return toolkit
+
+    @classmethod
+    def load_remote_examples_yaml(cls, url: str, access_token: Optional[str] = None) -> List[Message]:
+        """load remote examples by url: url/.well-known/examples.yaml
+
+        Args:
+            url (str): the base url of the remote toolkit
+        """
+        if not url.endswith("/"):
+            url += "/"
+        examples_yaml_url = url + ".well-known/examples.yaml"
+        if not url_file_exists(examples_yaml_url, cls._get_authorization_headers(access_token)):
+            return []
+
+        examples = []
+        with tempfile.TemporaryDirectory() as temp_dir:
+            response = requests.get(examples_yaml_url, headers=cls._get_authorization_headers(access_token))
+            if response.status_code != 200:
+                raise ValueError(
+                    f"Invalid resource, status_code: {response.status_code}, error message: {response.text}"
+                )
+
+            file_content = response.content.decode("utf-8")
+            if not file_content.strip():
+                raise ValueError(f"the content is empty from: {examples_yaml_url}")
+
+            file_path = os.path.join(temp_dir, "examples.yaml")
+            with open(file_path, "w+", encoding="utf-8") as f:
+                f.write(file_content)
+
+            examples = cls.load_examples_yaml(file_path)
+
+        return examples
+
+    @classmethod
+    def load_examples_dict(cls, examples_dict: Dict[str, Any]) -> List[Message]:
+        messages: List[Message] = []
+        for examples in examples_dict["examples"]:
+            examples = examples["context"]
+            for example in examples:
+                if "user" == example["role"]:
+                    messages.append(HumanMessage(example["content"]))
+                elif "bot" in example["role"]:
+                    plugin = example["plugin"]
+                    if "operationId" in plugin:
+                        function_call: FunctionCall = {
+                            "name": plugin["operationId"],
+                            "thoughts": plugin["thoughts"],
+                            "arguments": json.dumps(plugin["requestArguments"], ensure_ascii=False),
+                        }
+                    else:
+                        function_call = {
+                            "name": "",
+                            "thoughts": plugin["thoughts"],
+                            "arguments": "{}",
+                        }  # type: ignore
+                    messages.append(AIMessage("", function_call=function_call))
+                else:
+                    raise ValueError(f"invald role: <{example['role']}>")
+        return messages
+
+    @classmethod
+    def load_examples_yaml(cls, file: str) -> List[Message]:
+        """load examples from yaml file
+
+        Args:
+            file (str): the path of examples file
+
+        Returns:
+            List[Message]: the list of messages
+        """
+        content: dict = read_from_filename(file)[0]
+        if len(content) == 0 or "examples" not in content:
+            raise ValueError("invalid examples configuration file")
+        return cls.load_examples_dict(content)
+
+    def function_call_schemas(self) -> List[dict]:
+        return [tool.function_call_schema() for tool in self.get_tools()]

erniebot-agent/erniebot_agent/tools/calculator_tool.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from typing import Dict, List, Type
+
+from erniebot_agent.messages import AIMessage, HumanMessage, Message
+from erniebot_agent.tools.schema import ToolParameterView
+from pydantic import Field
+
+from .base import Tool
+
+
+class CalculatorToolInputView(ToolParameterView):
+    math_formula: str = Field(description='标准的数学公式，例如："2+3"、"3 - 4 * 6", "(3 + 4) * (6 + 4)" 等。 ')
+
+
+class CalculatorToolOutputView(ToolParameterView):
+    formula_result: float = Field(description="数学公式计算的结果")
+
+
+class CalculatorTool(Tool):
+    description: str = "CalculatorTool用于执行数学公式计算"
+    input_type: Type[ToolParameterView] = CalculatorToolInputView
+    ouptut_type: Type[ToolParameterView] = CalculatorToolOutputView
+
+    async def __call__(self, math_formula: str) -> Dict[str, float]:
+        return {"formula_result": eval(math_formula)}
+
+    @property
+    def examples(self) -> List[Message]:
+        return [
+            HumanMessage("请告诉我三加六等于多少？"),
+            AIMessage(
+                "",
+                function_call={
+                    "name": self.tool_name,
+                    "thoughts": f"用户想知道3加6等于多少，我可以使用{self.tool_name}工具来计算公式，其中`math_formula`字段的内容为：'3+6'。",
+                    "arguments": '{"math_formula": "3+6"}',
+                },
+                token_usage={
+                    "prompt_tokens": 5,
+                    "completion_tokens": 7,
+                },  # TODO: Functional AIMessage will not add in the memory, will it add token_usage?
+            ),
+            HumanMessage("一加八再乘以5是多少？"),
+            AIMessage(
+                "",
+                function_call={
+                    "name": self.tool_name,
+                    "thoughts": f"用户想知道1加8再乘5等于多少，我可以使用{self.tool_name}工具来计算公式，"
+                    "其中`math_formula`字段的内容为：'(1+8)*5'。",
+                    "arguments": '{"math_formula": "(1+8)*5"}',
+                },
+                token_usage={"prompt_tokens": 5, "completion_tokens": 7},  # For test only
+            ),
+            HumanMessage("我想知道十二除以四再加五等于多少？"),
+            AIMessage(
+                "",
+                function_call={
+                    "name": self.tool_name,
+                    "thoughts": f"用户想知道12除以4再加5等于多少，我可以使用{self.tool_name}工具来计算公式，"
+                    "其中`math_formula`字段的内容为：'12/4+5'。",
+                    "arguments": '{"math_formula": "12/4+5"}',
+                },
+                token_usage={"prompt_tokens": 5, "completion_tokens": 7},  # For test only
+            ),
+        ]

erniebot-agent/erniebot_agent/tools/current_time_tool.py

@@ -0,0 +1,47 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Dict, List, Type
+
+from erniebot_agent.messages import AIMessage, HumanMessage, Message
+from erniebot_agent.tools.schema import ToolParameterView
+from pydantic import Field
+
+from .base import Tool
+
+
+class CurrentTimeToolOutputView(ToolParameterView):
+    current_time: str = Field(description="当前时间")
+
+
+class CurrentTimeTool(Tool):
+    description: str = "CurrentTimeTool 用于获取当前时间"
+    ouptut_type: Type[ToolParameterView] = CurrentTimeToolOutputView
+
+    async def __call__(self) -> Dict[str, str]:
+        return {"current_time": datetime.strftime(datetime.now(), "%Y年%m月%d号 %点:%分:%秒")}
+
+    @property
+    def examples(self) -> List[Message]:
+        return [
+            HumanMessage("现在几点钟了"),
+            AIMessage(
+                "",
+                function_call={
+                    "name": self.tool_name,
+                    "thoughts": f"用户想知道现在几点了，我可以使用{self.tool_name}来获取当前时间，并从其中获得当前小时时间。",
+                    "arguments": "{}",
+                },
+                token_usage={"prompt_tokens": 5, "completion_tokens": 7},  # For test only
+            ),
+            HumanMessage("现在是什么时候？"),
+            AIMessage(
+                "",
+                function_call={
+                    "name": self.tool_name,
+                    "thoughts": f"用户想知道现在几点了，我可以使用{self.tool_name}来获取当前时间",
+                    "arguments": "{}",
+                },
+                token_usage={"prompt_tokens": 5, "completion_tokens": 7},  # For test only
+            ),
+        ]

erniebot-agent/erniebot_agent/tools/image_generation_tool.py

@@ -0,0 +1,117 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import os
+import uuid
+from typing import Any, Dict, List, Optional, Type
+
+from erniebot_agent.messages import AIMessage, HumanMessage, Message
+from erniebot_agent.tools.base import Tool
+from erniebot_agent.tools.schema import ToolParameterView
+from erniebot_agent.utils.common import download_file, get_cache_dir
+from pydantic import Field
+
+import erniebot
+
+
+class ImageGenerationInputView(ToolParameterView):
+    prompt: str = Field(description="描述图像内容、风格的文本。例如：生成一张月亮的照片，月亮很圆。")
+    width: int = Field(description="生成图片的宽度")
+    height: int = Field(description="生成图片的高度")
+    image_num: int = Field(description="生成图片的数量")
+
+
+class ImageGenerationOutputView(ToolParameterView):
+    image_path: str = Field(description="图片在本地机器上的保存路径")
+
+
+class ImageGenerationTool(Tool):
+    description: str = "AI作图、生成图片、画图的工具"
+    input_type: Type[ToolParameterView] = ImageGenerationInputView
+    ouptut_type: Type[ToolParameterView] = ImageGenerationOutputView
+
+    def __init__(
+        self,
+        yinian_access_token: Optional[str] = None,
+        yinian_ak: Optional[str] = None,
+        yinian_sk: Optional[str] = None,
+    ) -> None:
+        self.config: Dict[str, Optional[Any]]
+        if yinian_access_token is not None:
+            self.config = {"api_type": "yinian", "access_token": yinian_access_token}
+        elif yinian_ak is not None and yinian_sk is not None:
+            self.config = {"api_type": "yinian", "ak": yinian_ak, "sk": yinian_sk}
+        else:
+            raise ValueError("Please set the yinian_access_token, or set yinian_ak and yinian_sk")
+
+    async def __call__(
+        self,
+        prompt: str,
+        width: int = 512,
+        height: int = 512,
+        image_num: int = 1,
+    ) -> Dict[str, List[str]]:
+        response = erniebot.Image.create(
+            model="ernie-vilg-v2",
+            prompt=prompt,
+            width=width,
+            height=height,
+            image_num=image_num,
+            _config_=self.config,
+        )
+
+        image_path = []
+        cache_dir = get_cache_dir()
+        for item in response["data"]["sub_task_result_list"]:
+            image_url = item["final_image_list"][0]["img_url"]
+            save_path = os.path.join(cache_dir, f"img_{uuid.uuid1()}.png")
+            download_file(image_url, save_path)
+            image_path.append(save_path)
+        return {"image_path": image_path}
+
+    @property
+    def examples(self) -> List[Message]:
+        return [
+            HumanMessage("画一张小狗的图片，图像高度512，图像宽度512"),
+            AIMessage(
+                "",
+                function_call={
+                    "name": "ImageGenerationTool",
+                    "thoughts": "用户需要我生成一张小狗的图片，图像高度为512，宽度为512。我可以使用ImageGenerationTool工具来满足用户的需求。",
+                    "arguments": '{"prompt":"画一张小狗的图片，图像高度512，图像宽度512",'
+                    '"width":512,"height":512,"image_num":1}',
+                },
+            ),
+            HumanMessage("生成两张天空的图片"),
+            AIMessage(
+                "",
+                function_call={
+                    "name": self.tool_name,
+                    "thoughts": "用户想要生成两张天空的图片，我需要调用ImageGenerationTool工具的call接口，"
+                    "并设置prompt为'生成两张天空的图片'，width和height可以默认为512，image_num默认为2。",
+                    "arguments": '{"prompt":"生成两张天空的图片","width":512,"height":512,"image_num":2}',
+                },
+            ),
+            HumanMessage("使用AI作图工具，生成1张小猫的图片，高度和高度是1024"),
+            AIMessage(
+                "",
+                function_call={
+                    "name": self.tool_name,
+                    "thoughts": "用户需要生成一张小猫的图片，高度和宽度都是1024。我可以使用ImageGenerationTool工具来满足用户的需求。",
+                    "arguments": '{"prompt":"生成一张小猫的照片。","width":1024,"height":1024,"image_num":1}',
+                },
+            ),
+        ]

erniebot-agent/erniebot_agent/tools/schema.py

@@ -0,0 +1,415 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Type, get_args
+
+from erniebot_agent.utils.logging import logger
+from pydantic import BaseModel, Field, create_model
+from pydantic.fields import FieldInfo
+
+INVALID_FIELD_NAME = "__invalid_field_name__"
+
+
+def is_optional_type(type: Type):
+    args = get_args(type)
+    if len(args) == 0:
+        return False
+
+    return len([arg for arg in args if arg is None.__class__]) > 0
+
+
+def get_typing_list_type(type):
+    """get typing.List[T] element type
+
+    Args:
+        type (typing.List): Generics type
+    """
+    # 1. checking list type
+    if getattr(type, "_name", None) != "List":
+        return None
+
+    arg_type = get_args(type)[0]
+    return json_type(arg_type)
+
+
+def json_type(type: Optional[Type[object]] = None):
+    if type is None:
+        return "object"
+
+    mapping = {
+        int: "integer",
+        str: "string",
+        list: "array",
+        List: "array",
+        float: "number",
+        ToolParameterView: "object",
+    }
+
+    if inspect.isclass(type) and issubclass(type, ToolParameterView):
+        return "object"
+
+    if getattr(type, "_name", None) == "List":
+        return "array"
+
+    if type not in mapping:
+        args = [arg for arg in get_args(type) if arg is not None.__class__]
+        if len(args) > 1 or len(args) == 0:
+            raise ValueError(
+                "only support simple type: FieldType=int/str/float/ToolParameterView, "
+                "so the target type should be one of: FieldType, List[FieldType], "
+                f"Optional[FieldType], but receive {type}"
+            )
+        type = args[0]
+
+    if type in mapping:
+        return mapping[type]
+
+    if inspect.isclass(type) and issubclass(type, ToolParameterView):
+        return "object"
+
+    return str(type)
+
+
+def python_type_from_json_type(json_type_dict: dict) -> Type[object]:
+    simple_types = {"integer": int, "string": str, "number": float, "object": ToolParameterView}
+    if json_type_dict["type"] in simple_types:
+        return simple_types[json_type_dict["type"]]
+
+    assert (
+        json_type_dict["type"] == "array"
+    ), f"only support simple_types<{','.join(simple_types)}> and array type"
+    assert "type" in json_type_dict["items"], "<items> field must be defined when 'type'=array"
+
+    json_type_value = json_type_dict["items"]["type"]
+    if json_type_value == "string":
+        return List[str]
+    if json_type_value == "integer":
+        return List[int]
+    if json_type_value == "number":
+        return List[float]
+    if json_type_value == "object":
+        return List[ToolParameterView]
+
+    raise ValueError(f"unsupported data type: {json_type_value}")
+
+
+def scrub_dict(d: dict, remove_empty_dict: bool = False) -> Optional[dict]:
+    """remove empty Value node,
+
+        function_call_schema: require
+
+    Args:
+        d (dict): the instance of dictionary
+        remove_empty_dict (bool): whether remove empty dict
+
+    Returns:
+        dict: the dictionary data after slimming down
+    """
+    if type(d) is dict:
+        result = {}
+        for k, v in d.items():
+            v = scrub_dict(v, remove_empty_dict)
+            if v is not None:
+                result[k] = v
+
+        if len(result) == 0:
+            if not remove_empty_dict:
+                return {}
+            return None
+
+        return result
+    elif isinstance(d, list):
+        return [scrub_dict(item, remove_empty_dict) for item in d]  # type: ignore
+    else:
+        return d
+
+
+class OpenAPIProperty(BaseModel):
+    type: str
+    description: Optional[str] = None
+    required: Optional[List[str]] = None
+    items: dict = Field(default_factory=dict)
+    properties: dict = Field(default_factory=dict)
+
+
+def get_field_openapi_property(field_info: FieldInfo) -> OpenAPIProperty:
+    """convert pydantic FieldInfo instance to OpenAPIProperty value
+
+    Args:
+        field_info (FieldInfo): the field instance
+
+    Returns:
+        OpenAPIProperty: the converted OpenAPI Property
+    """
+    typing_list_type = get_typing_list_type(field_info.annotation)
+    if typing_list_type is not None:
+        field_type = "array"
+    elif is_optional_type(field_info.annotation):
+        field_type = json_type(get_args(field_info.annotation)[0])
+    else:
+        field_type = json_type(field_info.annotation)
+
+    property = {
+        "type": field_type,
+        "description": field_info.description,
+    }
+
+    if property["type"] == "array":
+        if typing_list_type == "object":
+            list_type: Type[ToolParameterView] = get_args(field_info.annotation)[0]
+            property["items"] = list_type.to_openapi_dict()
+        else:
+            property["items"] = {"type": typing_list_type}
+    elif property["type"] == "object":
+        if is_optional_type(field_info.annotation):
+            field_type_class: Type[ToolParameterView] = get_args(field_info.annotation)[0]
+        else:
+            field_type_class = field_info.annotation
+
+        openapi_dict = field_type_class.to_openapi_dict()
+        property.update(openapi_dict)
+
+    property["description"] = property.get("description", "")
+    return OpenAPIProperty(**property)
+
+
+class ToolParameterView(BaseModel):
+    @classmethod
+    def from_openapi_dict(cls, name, schema: dict) -> Type[ToolParameterView]:
+        """parse openapi component schemas to ParameterView
+        Args:
+            response_or_returns (dict): the content of status code
+
+        Returns:
+            _type_: _description_
+        """
+
+        # TODO(wj-Mcat): to load Optional field
+        fields = {}
+        for field_name, field_dict in schema.get("properties", {}).items():
+            field_type = python_type_from_json_type(field_dict)
+
+            if field_type is List[ToolParameterView]:
+                SubParameterView: Type[ToolParameterView] = ToolParameterView.from_openapi_dict(
+                    field_name, field_dict["items"]
+                )
+                field_type = List[SubParameterView]  # type: ignore
+
+            # TODO(wj-Mcat): remove supporting for `summary` field
+            if "summary" in field_dict:
+                description = field_dict["summary"]
+                logger.info("`summary` field will be deprecated, please use `description`")
+
+                if "description" in field_dict:
+                    logger.info("`description` field will be used instead of `summary`")
+                    description = field_dict["description"]
+            else:
+                description = field_dict.get("description", None)
+
+            description = description or ""
+
+            field = FieldInfo(annotation=field_type, description=description)
+
+            # TODO(wj-Mcat): to handle list field required & not-required
+            # if get_typing_list_type(field_type) is not None:
+            #     field.default_factory = list
+
+            fields[field_name] = (field_type, field)
+
+        return create_model("OpenAPIParameterView", __base__=ToolParameterView, **fields)
+
+    @classmethod
+    def to_openapi_dict(cls) -> dict:
+        """convert ParametersView to openapi spec dict
+
+        Returns:
+            dict: schema of openapi
+        """
+
+        required_names, properties = [], {}
+        for field_name, field_info in cls.model_fields.items():
+            if field_info.is_required() and not is_optional_type(field_info.annotation):
+                required_names.append(field_name)
+
+            properties[field_name] = dict(get_field_openapi_property(field_info))
+
+        result = {
+            "type": "object",
+            "properties": properties,
+        }
+        if len(required_names) > 0:
+            result["required"] = required_names
+        result = scrub_dict(result, remove_empty_dict=True)  # type: ignore
+        return result or {}
+
+    @classmethod
+    def function_call_schema(cls) -> dict:
+        """get function_call schame
+
+        Returns:
+            dict: the schema of function_call
+        """
+        return cls.to_openapi_dict()
+
+    @classmethod
+    def from_dict(cls, field_map: Dict[str, Any]):
+        """
+        Class method to create a Pydantic model dynamically based on a dictionary.
+
+        Args:
+            field_map (Dict[str, Any]): A dictionary mapping field names to their corresponding type
+            and description.
+
+        Returns:
+            PydanticModel: A dynamically created Pydantic model with fields specified by the
+            input dictionary.
+
+        Note:
+            This method is used to create a Pydantic model dynamically based on the provided dictionary,
+            where each field's type and description are specified in the input.
+
+        """
+        fields = {}
+        for field_name, field_dict in field_map.items():
+            field_type = field_dict["type"]
+            description = field_dict["description"]
+            field = FieldInfo(annotation=field_type, description=description)
+            fields[field_name] = (field_type, field)
+        return create_model(cls.__name__, __base__=ToolParameterView, **fields)
+
+
+@dataclass
+class RemoteToolView:
+    uri: str
+    method: str
+    name: str
+    description: str
+    parameters: Optional[Type[ToolParameterView]] = None
+    parameters_description: Optional[str] = None
+    returns: Optional[Type[ToolParameterView]] = None
+    returns_description: Optional[str] = None
+
+    returns_ref_uri: Optional[str] = None
+    parameters_ref_uri: Optional[str] = None
+
+    def to_openapi_dict(self):
+        result = {
+            "operationId": self.name,
+            "description": self.description,
+        }
+        if self.returns is not None:
+            response = {
+                "200": {
+                    "description": self.returns_description,
+                    "content": {
+                        "application/json": {
+                            "schema": {"$ref": "#/components/schemas/" + (self.returns_ref_uri or "")}
+                        }
+                    },
+                }
+            }
+            result["responses"] = response
+
+        if self.parameters is not None:
+            parameters = {
+                "required": True,
+                "content": {
+                    "application/json": {
+                        "schema": {"$ref": "#/components/schemas/" + (self.parameters_ref_uri or "")}
+                    }
+                },
+            }
+            result["requestBody"] = parameters
+        return {self.method: result}
+
+    @staticmethod
+    def from_openapi_dict(
+        uri: str, method: str, path_info: dict, parameters_views: dict[str, Type[ToolParameterView]]
+    ) -> RemoteToolView:
+        """construct RemoteToolView from openapi spec-dict info
+
+        Args:
+            uri (str): the url path of remote tool
+            method (str): http method: one of [get, post, put, delete]
+            path_info (dict): the spec info of remote tool
+            parameters_views (dict[str, ParametersView]):
+                the dict of parameters views which are the schema of input/output of tool
+
+        Returns:
+            RemoteToolView: the instance of remote tool view
+        """
+        parameters_ref_uri, returns_ref_uri = None, None
+        parameters, parameters_description = None, None
+        if "requestBody" in path_info:
+            request_ref = path_info["requestBody"]["content"]["application/json"]["schema"]["$ref"]
+            parameters_ref_uri = request_ref.split("/")[-1]
+            assert parameters_ref_uri in parameters_views
+            parameters = parameters_views[parameters_ref_uri]
+            parameters_description = path_info["requestBody"].get("description", None)
+
+        returns, returns_description = None, None
+        if "responses" in path_info:
+            response_ref = list(path_info["responses"].values())[0]["content"]["application/json"]["schema"][
+                "$ref"
+            ]
+            returns_ref_uri = response_ref.split("/")[-1]
+            assert returns_ref_uri in parameters_views
+            returns = parameters_views[returns_ref_uri]
+            returns_description = list(path_info["responses"].values())[0].get("description", None)
+
+        return RemoteToolView(
+            name=path_info["operationId"],
+            parameters=parameters,
+            parameters_description=parameters_description,
+            returns=returns,
+            returns_description=returns_description,
+            description=path_info.get("description", path_info.get("summary", None)),
+            method=method,
+            uri=uri,
+            # save ref id info
+            returns_ref_uri=returns_ref_uri,
+            parameters_ref_uri=parameters_ref_uri,
+        )
+
+    def function_call_schema(self):
+        inputs = {
+            "name": self.name,
+            "description": self.description,
+            # TODO(wj-Mcat): read examples from openapi.yaml
+            # "examples": [example.to_dict() for example in self.examples],
+        }
+        if self.parameters is not None:
+            inputs["parameters"] = self.parameters.function_call_schema()  # type: ignore
+        else:
+            inputs["parameters"] = {"type": "object", "properties": {}}
+
+        if self.returns is not None:
+            inputs["responses"] = self.returns.function_call_schema()  # type: ignore
+        return scrub_dict(inputs) or {}
+
+
+@dataclass
+class Endpoint:
+    url: str
+
+
+@dataclass
+class EndpointInfo:
+    title: str
+    description: str
+    version: str

erniebot-agent/erniebot_agent/tools/tool_manager.py

@@ -0,0 +1,69 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+from typing import Dict, List, final
+
+from erniebot_agent.tools.base import Tool
+
+
+@final
+class ToolManager(object):
+    """A `ToolManager` instance manages tools for an agent.
+
+    This implementation is based on `ToolsManager` in
+    https://github.com/deepset-ai/haystack/blob/main/haystack/agents/base.py
+    """
+
+    def __init__(self, tools: List[Tool]) -> None:
+        super().__init__()
+        self._tools: Dict[str, Tool] = {}
+        for tool in tools:
+            self.add_tool(tool)
+
+    def __getitem__(self, tool_name: str) -> Tool:
+        return self.get_tool(tool_name)
+
+    def add_tool(self, tool: Tool) -> None:
+        tool_name = tool.tool_name
+        if tool_name in self._tools:
+            raise RuntimeError(f"Name {repr(tool_name)} is already registered.")
+        self._tools[tool_name] = tool
+
+    def remove_tool(self, tool: Tool) -> None:
+        tool_name = tool.tool_name
+        if tool_name not in self._tools:
+            raise RuntimeError(f"Name {repr(tool_name)} is not registered.")
+        if self._tools[tool_name] is not tool:
+            raise RuntimeError(f"The tool with the registered name {repr(tool_name)} is not the given tool.")
+        self._tools.pop(tool_name)
+
+    def get_tool(self, tool_name: str) -> Tool:
+        if tool_name not in self._tools:
+            raise RuntimeError(f"Name {repr(tool_name)} is not registered.")
+        return self._tools[tool_name]
+
+    def get_tools(self) -> List[Tool]:
+        return list(self._tools.values())
+
+    def get_tool_names(self) -> str:
+        return ", ".join(self._tools.keys())
+
+    def get_tool_names_with_descriptions(self) -> str:
+        return "\n".join(
+            f"{name}:{json.dumps(tool.function_call_schema())}" for name, tool in self._tools.items()
+        )
+
+    def get_tool_schemas(self):
+        return [tool.function_call_schema() for tool in self._tools.values()]