#!/usr/bin/env python
# coding: utf-8

# # 自定义多动作代理
# 
# 本笔记本介绍如何创建自定义代理。
# 
# 代理包括两个部分：
# 
# - 工具：代理可用的工具。
# - 代理类本身：决定采取哪种行动。
# 
# 在本笔记本中，我们将介绍如何创建一个可以预测/一次执行多个步骤的自定义代理。

# In[1]:


# 导入所需的模块
from langchain.agents import AgentExecutor, BaseMultiActionAgent, Tool
from langchain_community.utilities import SerpAPIWrapper


# In[2]:


# 定义一个函数，函数名为random_word，接受一个字符串类型的参数query，返回一个字符串类型的结果
def random_word(query: str) -> str:
    # 打印一条消息，表示当前正在执行这个函数
    print("\nNow I'm doing this!")
    # 返回字符串"foo"
    return "foo"


# In[3]:


# 创建一个SerpAPIWrapper对象
search = SerpAPIWrapper()

# 创建一个包含两个Tool对象的列表
tools = [
    Tool(
        name="Search",  # 工具名称为"Search"
        func=search.run,  # 调用search对象的run方法
        description="useful for when you need to answer questions about current events",  # 描述为"在需要回答有关当前事件的问题时很有用"
    ),
    Tool(
        name="RandomWord",  # 工具名称为"RandomWord"
        func=random_word,  # 调用random_word函数
        description="call this to get a random word.",  # 描述为"调用此函数以获取一个随机单词"
    ),
]


# In[4]:


from typing import Any, List, Tuple, Union

from langchain_core.agents import AgentAction, AgentFinish


class FakeAgent(BaseMultiActionAgent):
    """Fake Custom Agent."""

    @property
    def input_keys(self):
        return ["input"]

    def plan(
        self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs: Any
    ) -> Union[List[AgentAction], AgentFinish]:
        """Given input, decided what to do.

        Args:
            intermediate_steps: Steps the LLM has taken to date,
                along with observations
            **kwargs: User inputs.

        Returns:
            Action specifying what tool to use.
        """
        if len(intermediate_steps) == 0:
            return [
                AgentAction(tool="Search", tool_input=kwargs["input"], log=""),
                AgentAction(tool="RandomWord", tool_input=kwargs["input"], log=""),
            ]
        else:
            return AgentFinish(return_values={"output": "bar"}, log="")

    async def aplan(
        self, intermediate_steps: List[Tuple[AgentAction, str]], **kwargs: Any
    ) -> Union[List[AgentAction], AgentFinish]:
        """Given input, decided what to do.

        Args:
            intermediate_steps: Steps the LLM has taken to date,
                along with observations
            **kwargs: User inputs.

        Returns:
            Action specifying what tool to use.
        """
        if len(intermediate_steps) == 0:
            return [
                AgentAction(tool="Search", tool_input=kwargs["input"], log=""),
                AgentAction(tool="RandomWord", tool_input=kwargs["input"], log=""),
            ]
        else:
            return AgentFinish(return_values={"output": "bar"}, log="")


# In[5]:


# 创建一个名为agent的FakeAgent对象
agent = FakeAgent()


# In[6]:


# 创建一个AgentExecutor对象，使用from_agent_and_tools方法，传入agent和tools参数，并设置verbose为True
agent_executor = AgentExecutor.from_agent_and_tools(
    agent=agent, tools=tools, verbose=True
)


# In[7]:


# 调用agent_executor的run方法，并传入参数"How many people live in canada as of 2023?"
agent_executor.run("How many people live in canada as of 2023?")


# In[ ]: