From f95151d3dbec09416548aa732681b474dd556ef6 Mon Sep 17 00:00:00 2001 From: LeoLjl Date: Sat, 28 Dec 2024 00:44:53 +0000 Subject: [PATCH 1/2] Add executor injected with tools. --- autogen/agentchat/contrib/tool_retriever.py | 132 ++++++++++++++++++++ test/agentchat/test_executor_with_tools.py | 37 ++++++ 2 files changed, 169 insertions(+) create mode 100644 test/agentchat/test_executor_with_tools.py diff --git a/autogen/agentchat/contrib/tool_retriever.py b/autogen/agentchat/contrib/tool_retriever.py index 4ca2ddcda..d875ae851 100644 --- a/autogen/agentchat/contrib/tool_retriever.py +++ b/autogen/agentchat/contrib/tool_retriever.py @@ -4,16 +4,26 @@ # # Portions derived from https://github.com/microsoft/autogen are under the MIT License. # SPDX-License-Identifier: MIT +import contextlib +import functools import importlib.util import inspect +import io import os +import traceback +from hashlib import md5 +from pathlib import Path from textwrap import dedent, indent +from typing import List, Optional, Union import pandas as pd from sentence_transformers import SentenceTransformer, util from autogen import AssistantAgent, UserProxyAgent from autogen.coding import LocalCommandLineCodeExecutor +from autogen.coding.base import CodeBlock, CodeResult +from autogen.function_utils import load_basemodels_if_needed +from autogen.tools import Tool class ToolBuilder: @@ -77,6 +87,108 @@ def bind_user_proxy(self, agent: UserProxyAgent, tool_root: str): return updated_user_proxy +class LocalExecutorWithTools: + """ + 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. + + For example, for a tool converted from langchain, the relevant functions can be called directly. + ```python + from langchain_community.tools import WikipediaQueryRun + from langchain_community.utilities import WikipediaAPIWrapper + from autogen.interop import Interoperability + + api_wrapper = WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=3000) + langchain_tool = WikipediaQueryRun(api_wrapper=api_wrapper) + interop = Interoperability() + ag2_tool = interop.convert_tool(tool=langchain_tool, type="langchain") + + # `ag2_tool.name` is wikipedia + local_executor = LocalExecutorWithTools(tools=[ag2_tool], work_dir='./') + + code = ''' + result = wikipedia(tool_input={"query":"Christmas"}) + print(result) + ''' + print( + local_executor.execute_code_blocks( + code_blocks=[ + CodeBlock(language="python", code=code), + ] + ) + ) + ``` + In this case, the `wikipedia` function can be called directly in the code block. This hides the complexity of the tool. + + Args: + tools: The tools to inject into the code execution environment. Default is an empty list. + work_dir: The working directory for the code execution. Default is the current directory. + """ + + def __init__(self, tools: Optional[List[Tool]] = None, work_dir: Union[Path, str] = Path(".")): + self.tools = tools if tools is not None else [] + self.work_dir = work_dir + if not os.path.exists(work_dir): + os.makedirs(work_dir, exist_ok=True) + + def execute_code_blocks(self, code_blocks: List[CodeBlock]) -> CodeResult: + """Execute code blocks and return the result. + + Args: + code_blocks (List[CodeBlock]): The code blocks to execute. + + Returns: + CodeResult: The result of the code execution. + """ + logs_all = "" + exit_code = 0 # Success code + code_file = None # Path to the first saved codeblock content + + for idx, code_block in enumerate(code_blocks): + code = code_block.code + code_hash = md5(code.encode()).hexdigest() + filename = f"tmp_code_{code_hash}.py" + filepath = os.path.join(self.work_dir, filename) + # Save code content to file + with open(filepath, "w", encoding="utf-8") as f: + f.write(code) + + if idx == 0: + code_file = filepath + + # Create a new execution environment + execution_env = {} + # Inject the tools + for tool in self.tools: + execution_env[tool.name] = _wrap_function(tool.func) + + # Prepare to capture stdout and stderr + stdout = io.StringIO() + stderr = io.StringIO() + + # Execute the code block + try: + # Redirect stdout and stderr + with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr): + # Exec the code in the execution environment + exec(code, execution_env) + except Exception: + # Capture exception traceback + tb = traceback.format_exc() + stderr.write(tb) + exit_code = 1 # Non-zero exit code indicates failure + + # Collect outputs + stdout_content = stdout.getvalue() + stderr_content = stderr.getvalue() + logs_all += stdout_content + stderr_content + + return CodeResult(exit_code=exit_code, output=logs_all, code_file=code_file) + + def restart(self): + """Restart the code executor. Since this executor is stateless, no action is needed.""" + pass + + def get_full_tool_description(py_file): """ Retrieves the function signature for a given Python file. @@ -100,6 +212,26 @@ def get_full_tool_description(py_file): raise ValueError(f"Function {function_name} not found in {py_file}") +def _wrap_function(func): + """Wrap the function to dump the return value to json. + + Handles both sync and async functions. + + Args: + func: the function to be wrapped. + + Returns: + The wrapped function. + """ + + @load_basemodels_if_needed + @functools.wraps(func) + def _wrapped_func(*args, **kwargs): + return func(*args, **kwargs) + + return _wrapped_func + + def find_callables(directory): """ Find all callable objects defined in Python files within the specified directory. diff --git a/test/agentchat/test_executor_with_tools.py b/test/agentchat/test_executor_with_tools.py new file mode 100644 index 000000000..a32ed8ea6 --- /dev/null +++ b/test/agentchat/test_executor_with_tools.py @@ -0,0 +1,37 @@ +# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai +# +# SPDX-License-Identifier: Apache-2.0 +import sys + +import pytest +from langchain_community.tools import WikipediaQueryRun +from langchain_community.utilities import WikipediaAPIWrapper + +from autogen.agentchat.contrib.tool_retriever import LocalExecutorWithTools +from autogen.coding.base import CodeBlock +from autogen.interop import Interoperability + + +# skip if python version is not >= 3.9 +@pytest.mark.skipif( + sys.version_info < (3, 9), reason="Only Python 3.9 and above are supported for LangchainInteroperability" +) +def test_execution(): + api_wrapper = WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=3000) + langchain_tool = WikipediaQueryRun(api_wrapper=api_wrapper) + interop = Interoperability() + ag2_tool = interop.convert_tool(tool=langchain_tool, type="langchain") + + # `ag2_tool.name` is wikipedia + local_executor = LocalExecutorWithTools(tools=[ag2_tool], work_dir="./") + + code = """ + result = wikipedia(tool_input={"query":"Christmas"}) + print(result) + """ + result = local_executor.execute_code_blocks( + code_blocks=[ + CodeBlock(language="python", code=code), + ] + ) + print(result) From a1a1a22cb34c8ca3d0c46a9c977c49132f8c2c46 Mon Sep 17 00:00:00 2001 From: LeoLjl Date: Sat, 28 Dec 2024 01:09:09 +0000 Subject: [PATCH 2/2] fix --- test/agentchat/test_executor_with_tools.py | 37 ---------------------- 1 file changed, 37 deletions(-) delete mode 100644 test/agentchat/test_executor_with_tools.py diff --git a/test/agentchat/test_executor_with_tools.py b/test/agentchat/test_executor_with_tools.py deleted file mode 100644 index a32ed8ea6..000000000 --- a/test/agentchat/test_executor_with_tools.py +++ /dev/null @@ -1,37 +0,0 @@ -# Copyright (c) 2023 - 2024, Owners of https://github.com/ag2ai -# -# SPDX-License-Identifier: Apache-2.0 -import sys - -import pytest -from langchain_community.tools import WikipediaQueryRun -from langchain_community.utilities import WikipediaAPIWrapper - -from autogen.agentchat.contrib.tool_retriever import LocalExecutorWithTools -from autogen.coding.base import CodeBlock -from autogen.interop import Interoperability - - -# skip if python version is not >= 3.9 -@pytest.mark.skipif( - sys.version_info < (3, 9), reason="Only Python 3.9 and above are supported for LangchainInteroperability" -) -def test_execution(): - api_wrapper = WikipediaAPIWrapper(top_k_results=1, doc_content_chars_max=3000) - langchain_tool = WikipediaQueryRun(api_wrapper=api_wrapper) - interop = Interoperability() - ag2_tool = interop.convert_tool(tool=langchain_tool, type="langchain") - - # `ag2_tool.name` is wikipedia - local_executor = LocalExecutorWithTools(tools=[ag2_tool], work_dir="./") - - code = """ - result = wikipedia(tool_input={"query":"Christmas"}) - print(result) - """ - result = local_executor.execute_code_blocks( - code_blocks=[ - CodeBlock(language="python", code=code), - ] - ) - print(result)