2020from sentence_transformers import SentenceTransformer , util
2121
2222from autogen import AssistantAgent , UserProxyAgent
23- from autogen .coding import LocalCommandLineCodeExecutor
23+ from autogen .coding import CodeExecutor , CodeExtractor , LocalCommandLineCodeExecutor , MarkdownCodeExtractor
2424from autogen .coding .base import CodeBlock , CodeResult
25- from autogen .function_utils import load_basemodels_if_needed
25+ from autogen .function_utils import get_function_schema , load_basemodels_if_needed
2626from autogen .tools import Tool
2727
2828
2929class ToolBuilder :
30- TOOL_USING_PROMPT = """# Functions
31- You have access to the following functions. They can be accessed from the module called 'functions' by their function names.
30+ TOOL_PROMPT_DEFAULT = """\n # # Functions
31+ You have access to the following functions. They can be accessed from the module called 'functions' by their function names.
3232For example, if there is a function called `foo` you could import it by writing `from functions import foo`
3333{functions}
34+ """
35+ TOOL_PROMPT_USER_DEFINED = """\n ## Functions
36+ You have access to the following functions. You can write python code to call these functions directly without importing them.
37+ {functions}
3438"""
3539
36- def __init__ (self , corpus_path , retriever = "all-mpnet-base-v2" ):
37-
38- self .df = pd .read_csv (corpus_path , sep = "\t " )
39- document_list = self .df ["document_content" ].tolist ()
40+ def __init__ (self , corpus_root , retriever = "all-mpnet-base-v2" , type = "default" ):
41+ if type == "default" :
42+ corpus_path = os .path .join (corpus_root , "tool_description.tsv" )
43+ self .df = pd .read_csv (corpus_path , sep = "\t " )
44+ document_list = self .df ["document_content" ].tolist ()
45+ self .TOOL_PROMPT = self .TOOL_PROMPT_DEFAULT
46+ else :
47+ self .TOOL_PROMPT = self .TOOL_PROMPT_USER_DEFINED
48+ # user defined tools, retrieve is actually not needed, just for consistency
49+ document_list = []
50+ for tool in corpus_root :
51+ document_list .append (tool .description )
4052
4153 self .model = SentenceTransformer (retriever )
4254 self .embeddings = self .model .encode (document_list )
55+ self .type = type
4356
4457 def retrieve (self , query , top_k = 3 ):
4558 # Encode the query using the Sentence Transformer model
@@ -55,39 +68,59 @@ def retrieve(self, query, top_k=3):
5568 def bind (self , agent : AssistantAgent , functions : str ):
5669 """Binds the function to the agent so that agent is aware of it."""
5770 sys_message = agent .system_message
58- sys_message += self .TOOL_USING_PROMPT .format (functions = functions )
71+ sys_message += self .TOOL_PROMPT .format (functions = functions )
5972 agent .update_system_message (sys_message )
6073 return
6174
62- def bind_user_proxy (self , agent : UserProxyAgent , tool_root : str ):
75+ def bind_user_proxy (self , agent : UserProxyAgent , tool_root : Union [ str , list ] ):
6376 """
6477 Updates user proxy agent with a executor so that code executor can successfully execute function-related code.
6578 Returns an updated user proxy.
6679 """
67- # Find all the functions in the tool root
68- functions = find_callables (tool_root )
69-
70- code_execution_config = agent ._code_execution_config
71- executor = LocalCommandLineCodeExecutor (
72- timeout = code_execution_config .get ("timeout" , 180 ),
73- work_dir = code_execution_config .get ("work_dir" , "coding" ),
74- functions = functions ,
75- )
76- code_execution_config = {
77- "executor" : executor ,
78- "last_n_messages" : code_execution_config .get ("last_n_messages" , 1 ),
79- }
80- updated_user_proxy = UserProxyAgent (
81- name = agent .name ,
82- is_termination_msg = agent ._is_termination_msg ,
83- code_execution_config = code_execution_config ,
84- human_input_mode = "NEVER" ,
85- default_auto_reply = agent ._default_auto_reply ,
86- )
87- return updated_user_proxy
88-
89-
90- class LocalExecutorWithTools :
80+ if isinstance (tool_root , str ):
81+ # Find all the functions in the tool root
82+ functions = find_callables (tool_root )
83+
84+ code_execution_config = agent ._code_execution_config
85+ executor = LocalCommandLineCodeExecutor (
86+ timeout = code_execution_config .get ("timeout" , 180 ),
87+ work_dir = code_execution_config .get ("work_dir" , "coding" ),
88+ functions = functions ,
89+ )
90+ code_execution_config = {
91+ "executor" : executor ,
92+ "last_n_messages" : code_execution_config .get ("last_n_messages" , 1 ),
93+ }
94+ updated_user_proxy = UserProxyAgent (
95+ name = agent .name ,
96+ is_termination_msg = agent ._is_termination_msg ,
97+ code_execution_config = code_execution_config ,
98+ human_input_mode = "NEVER" ,
99+ default_auto_reply = agent ._default_auto_reply ,
100+ )
101+ return updated_user_proxy
102+ else :
103+ # second case: user defined tools
104+ code_execution_config = agent ._code_execution_config
105+ executor = LocalExecutorWithTools (
106+ tools = tool_root ,
107+ work_dir = code_execution_config .get ("work_dir" , "coding" ),
108+ )
109+ code_execution_config = {
110+ "executor" : executor ,
111+ "last_n_messages" : code_execution_config .get ("last_n_messages" , 1 ),
112+ }
113+ updated_user_proxy = UserProxyAgent (
114+ name = agent .name ,
115+ is_termination_msg = agent ._is_termination_msg ,
116+ code_execution_config = code_execution_config ,
117+ human_input_mode = "NEVER" ,
118+ default_auto_reply = agent ._default_auto_reply ,
119+ )
120+ return updated_user_proxy
121+
122+
123+ class LocalExecutorWithTools (CodeExecutor ):
91124 """
92125 An executor that executes code blocks with injected tools. In this executor, the func within the tools can be called directly without declaring in the code block.
93126
@@ -124,6 +157,11 @@ class LocalExecutorWithTools:
124157 work_dir: The working directory for the code execution. Default is the current directory.
125158 """
126159
160+ @property
161+ def code_extractor (self ) -> CodeExtractor :
162+ """(Experimental) Export a code extractor that can be used by an agent."""
163+ return MarkdownCodeExtractor ()
164+
127165 def __init__ (self , tools : Optional [List [Tool ]] = None , work_dir : Union [Path , str ] = Path ("." )):
128166 self .tools = tools if tools is not None else []
129167 self .work_dir = work_dir
@@ -189,6 +227,51 @@ def restart(self):
189227 pass
190228
191229
230+ def format_ag2_tool (tool : Tool ):
231+ # get the args first
232+ schema = get_function_schema (tool .func , description = tool .description )
233+
234+ arg_name = list (inspect .signature (tool .func ).parameters .keys ())[0 ]
235+ arg_info = schema ["function" ]["parameters" ]["properties" ][arg_name ]["properties" ]
236+
237+ content = f'def { tool .name } { arg_name } \n """\n '
238+ content += indent (tool .description , " " ) + "\n "
239+ content += (
240+ indent (
241+ f"You must format all the arguments into a dictionary and pass them as **kwargs to { arg_name } ,
242+ " " ,
243+ )
244+ + "\n "
245+ + indent (f"For example:\n \t result = { tool .name } { arg_name } , " " )
246+ + "\n "
247+ )
248+ content += indent (f"Arguments passed in { arg_name } \n " , " " )
249+ for arg , info in arg_info .items ():
250+ content += indent (f"{ arg } { info ['type' ]} { info ['description' ]} \n " , " " * 2 )
251+ content += ' """\n '
252+ return content
253+
254+
255+ def _wrap_function (func ):
256+ """Wrap the function to dump the return value to json.
257+
258+ Handles both sync and async functions.
259+
260+ Args:
261+ func: the function to be wrapped.
262+
263+ Returns:
264+ The wrapped function.
265+ """
266+
267+ @load_basemodels_if_needed
268+ @functools .wraps (func )
269+ def _wrapped_func (* args , ** kwargs ):
270+ return func (* args , ** kwargs )
271+
272+ return _wrapped_func
273+
274+
192275def get_full_tool_description (py_file ):
193276 """
194277 Retrieves the function signature for a given Python file.
0 commit comments