Skip to content

ToolBuilder

autogen.agentchat.contrib.captainagent.ToolBuilder #

ToolBuilder(corpus_root, retriever='all-mpnet-base-v2', type='default')
Source code in autogen/agentchat/contrib/captainagent/tool_retriever.py 
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
def __init__(self, corpus_root, retriever="all-mpnet-base-v2", type="default"):
    if type == "default":
        corpus_path = os.path.join(corpus_root, "tool_description.tsv")
        self.df = pd.read_csv(corpus_path, sep="\t")
        document_list = self.df["document_content"].tolist()
        self.TOOL_PROMPT = self.TOOL_PROMPT_DEFAULT
    else:
        self.TOOL_PROMPT = self.TOOL_PROMPT_USER_DEFINED
        # user defined tools, retrieve is actually not needed, just for consistency
        document_list = []
        for tool in corpus_root:
            document_list.append(tool.description)

    self.model = SentenceTransformer(retriever)
    self.embeddings = self.model.encode(document_list)
    self.type = type

TOOL_PROMPT_DEFAULT class-attribute instance-attribute #

TOOL_PROMPT_DEFAULT = "\n## Functions\nYou have access to the following functions. They can be accessed from the module called 'functions' by their function names.\nFor example, if there is a function called `foo` you could import it by writing `from functions import foo`\n{functions}\n"

TOOL_PROMPT_USER_DEFINED class-attribute instance-attribute #

TOOL_PROMPT_USER_DEFINED = '\n## Functions\nYou have access to the following functions. You can write python code to call these functions directly without importing them.\n{functions}\n'

df instance-attribute #

df = read_csv(corpus_path, sep='\t')

TOOL_PROMPT instance-attribute #

TOOL_PROMPT = TOOL_PROMPT_DEFAULT

model instance-attribute #

model = SentenceTransformer(retriever)

embeddings instance-attribute #

embeddings = encode(document_list)

type instance-attribute #

type = type

retrieve #

retrieve(query, top_k=3)
Source code in autogen/agentchat/contrib/captainagent/tool_retriever.py 
61
62
63
64
65
66
67
68
69
70
def retrieve(self, query, top_k=3):
    # Encode the query using the Sentence Transformer model
    query_embedding = self.model.encode([query])

    hits = util.semantic_search(query_embedding, self.embeddings, top_k=top_k)

    results = []
    for hit in hits[0]:
        results.append(self.df.iloc[hit["corpus_id"], 1])
    return results

bind #

bind(agent, functions)

Binds the function to the agent so that agent is aware of it.

Source code in autogen/agentchat/contrib/captainagent/tool_retriever.py 
72
73
74
75
76
77
def bind(self, agent: AssistantAgent, functions: str):
    """Binds the function to the agent so that agent is aware of it."""
    sys_message = agent.system_message
    sys_message += self.TOOL_PROMPT.format(functions=functions)
    agent.update_system_message(sys_message)
    return

bind_user_proxy #

bind_user_proxy(agent, tool_root)

Updates user proxy agent with a executor so that code executor can successfully execute function-related code. Returns an updated user proxy.

Source code in autogen/agentchat/contrib/captainagent/tool_retriever.py 
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def bind_user_proxy(self, agent: UserProxyAgent, tool_root: Union[str, list]):
    """Updates user proxy agent with a executor so that code executor can successfully execute function-related code.
    Returns an updated user proxy.
    """
    if isinstance(tool_root, str):
        # Find all the functions in the tool root
        functions = find_callables(tool_root)

        code_execution_config = agent._code_execution_config
        executor = LocalCommandLineCodeExecutor(
            timeout=code_execution_config.get("timeout", 180),
            work_dir=code_execution_config.get("work_dir", "coding"),
            functions=functions,
        )
        code_execution_config = {
            "executor": executor,
            "last_n_messages": code_execution_config.get("last_n_messages", 1),
        }
        updated_user_proxy = UserProxyAgent(
            name=agent.name,
            is_termination_msg=agent._is_termination_msg,
            code_execution_config=code_execution_config,
            human_input_mode="NEVER",
            default_auto_reply=agent._default_auto_reply,
        )
        return updated_user_proxy
    else:
        # second case: user defined tools
        code_execution_config = agent._code_execution_config
        executor = LocalExecutorWithTools(
            tools=tool_root,
            work_dir=code_execution_config.get("work_dir", "coding"),
        )
        code_execution_config = {
            "executor": executor,
            "last_n_messages": code_execution_config.get("last_n_messages", 1),
        }
        updated_user_proxy = UserProxyAgent(
            name=agent.name,
            is_termination_msg=agent._is_termination_msg,
            code_execution_config=code_execution_config,
            human_input_mode="NEVER",
            default_auto_reply=agent._default_auto_reply,
        )
        return updated_user_proxy