ag2 0.9.1a1__py3-none-any.whl → 0.9.1.post0__py3-none-any.whl

autogen/agentchat/contrib/captainagent/tool_retriever.py

@@ -0,0 +1,335 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 Optional, Union
+
+from .... import AssistantAgent, UserProxyAgent
+from ....coding import CodeExecutor, CodeExtractor, LocalCommandLineCodeExecutor, MarkdownCodeExtractor
+from ....coding.base import CodeBlock, CodeResult
+from ....doc_utils import export_module
+from ....import_utils import optional_import_block, require_optional_import
+from ....tools import Tool, get_function_schema, load_basemodels_if_needed
+
+with optional_import_block():
+    import pandas as pd
+    from sentence_transformers import SentenceTransformer, util
+
+
+@require_optional_import(["pandas", "sentence_transformers"], "retrievechat")
+@export_module("autogen.agentchat.contrib.captainagent")
+class ToolBuilder:
+    TOOL_PROMPT_DEFAULT = """\n## Functions
+You have access to the following functions. They can be accessed from the module called 'functions' by their function names.
+For example, if there is a function called `foo` you could import it by writing `from functions import foo`
+{functions}
+"""
+    TOOL_PROMPT_USER_DEFINED = """\n## Functions
+You have access to the following functions. You can write python code to call these functions directly without importing them.
+{functions}
+"""
+
+    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
+
+    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
+
+    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
+
+    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
+
+
+class LocalExecutorWithTools(CodeExecutor):
+    """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.
+    """
+
+    @property
+    def code_extractor(self) -> CodeExtractor:
+        """(Experimental) Export a code extractor that can be used by an agent."""
+        return MarkdownCodeExtractor()
+
+    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
+
+
+@export_module("autogen.agentchat.contrib.captainagent")
+def format_ag2_tool(tool: Tool):
+    # get the args first
+    schema = get_function_schema(tool.func, description=tool.description)
+
+    arg_name = list(inspect.signature(tool.func).parameters.keys())[0]
+    arg_info = schema["function"]["parameters"]["properties"][arg_name]["properties"]
+
+    content = f'def {tool.name}({arg_name}):\n    """\n'
+    content += indent(tool.description, "    ") + "\n"
+    content += (
+        indent(
+            f"You must format all the arguments into a dictionary and pass them as **kwargs to {arg_name}. You should use print function to get the results.",
+            "    ",
+        )
+        + "\n"
+        + indent(f"For example:\n\tresult = {tool.name}({arg_name}={{'arg1': 'value1' }})", "    ")
+        + "\n"
+    )
+    content += indent(f"Arguments passed in {arg_name}:\n", "    ")
+    for arg, info in arg_info.items():
+        content += indent(f"{arg} ({info['type']}): {info['description']}\n", "    " * 2)
+    content += '    """\n'
+    return content
+
+
+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
+
+
+@export_module("autogen.agentchat.contrib.captainagent")
+def get_full_tool_description(py_file):
+    """Retrieves the function signature for a given Python file."""
+    with open(py_file) as f:
+        code = f.read()
+        exec(code)
+        function_name = os.path.splitext(os.path.basename(py_file))[0]
+        if function_name in locals():
+            func = locals()[function_name]
+            content = f"def {func.__name__}{inspect.signature(func)}:\n"
+            docstring = func.__doc__
+
+            if docstring:
+                docstring = dedent(docstring)
+                docstring = '"""' + docstring + '"""'
+                docstring = indent(docstring, "    ")
+                content += docstring + "\n"
+            return content
+        else:
+            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."""
+    callables = []
+    for root, dirs, files in os.walk(directory):
+        for file in files:
+            if file.endswith(".py"):
+                module_name = os.path.splitext(file)[0]
+                module_path = os.path.join(root, file)
+                spec = importlib.util.spec_from_file_location(module_name, module_path)
+                module = importlib.util.module_from_spec(spec)
+                spec.loader.exec_module(module)
+                for name, value in module.__dict__.items():
+                    if callable(value) and name == module_name:
+                        callables.append(value)
+                        break
+    return callables

autogen/agentchat/contrib/captainagent/tools/README.md

@@ -0,0 +1,44 @@
+# Introduction
+
+This directory contains a library of manually created python tools. These tools have three categories: math, data_analysis and information_retrieval.
+
+# Directory Layout
+```
+tools
+├── README.md
+├── data_analysis
+│   ├── calculate_correlation.py
+│   └── ...
+├── information_retrieval
+│   ├── arxiv_download.py
+│   ├── arxiv_search.py
+│   └── ...
+├── math
+│   ├── calculate_circle_area_from_diameter.py
+│   └── ...
+└── tool_description.tsv
+```
+
+Tools can be imported from `tools/{category}/{tool_name}.py` with exactly the same function name.
+
+`tool_description.tsv` contains descriptions of tools for retrieval.
+
+# How to use
+Some tools require Bing Search API key and RapidAPI key. For Bing API, you can read more about how to get an API on the [Bing Web Search API](https://www.microsoft.com/en-us/bing/apis/bing-web-search-api) page. For RapidAPI, you can [sign up](https://rapidapi.com/auth/sign-up) and subscribe to these two links([link1](https://rapidapi.com/solid-api-solid-api-default/api/youtube-transcript3), [link2](https://rapidapi.com/420vijay47/api/youtube-mp3-downloader2)). These apis have free billing options and there is no need to worry about extra costs.
+
+To install the requirements for running tools, use pip install.
+```bash
+pip install -r autogen/agentchat/contrib/captainagent/tools/requirements.txt
+```
+
+Whenever you run the tool-related code, remember to export the api keys to system variables.
+```bash
+export BING_API_KEY=""
+export RAPID_API_KEY=""
+```
+or
+```python
+import os
+os.environ["BING_API_KEY"] = ""
+os.environ["RAPID_API_KEY"] = ""
+```

autogen/agentchat/contrib/captainagent/tools/__init__.py

@@ -0,0 +1,5 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+__all__: list[str] = []

autogen/agentchat/contrib/captainagent/tools/data_analysis/calculate_correlation.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+def calculate_correlation(csv_path: str, column1: str, column2: str, method: str = "pearson") -> float:
+    """Calculate the correlation between two columns in a CSV file.
+
+    Args:
+    csv_path (str): The path to the CSV file.
+    column1 (str): The name of the first column.
+    column2 (str): The name of the second column.
+    method (str or callable, optional): The method used to calculate the correlation.
+        - 'pearson' (default): Pearson correlation coefficient.
+        - 'kendall': Kendall Tau correlation coefficient.
+        - 'spearman': Spearman rank correlation coefficient.
+        - callable: A custom correlation function that takes two arrays and returns a scalar.
+
+    Returns:
+    float: The correlation coefficient between the two columns.
+    """
+    import pandas as pd
+
+    # Read the CSV file into a pandas DataFrame
+    df = pd.read_csv(csv_path)
+
+    # Select the specified columns
+    selected_columns = df[[column1, column2]]
+
+    # Calculate the correlation based on the specified method
+    if method == "pearson":
+        correlation = selected_columns.corr().iloc[0, 1]
+    elif method == "kendall":
+        correlation = selected_columns.corr(method="kendall").iloc[0, 1]
+    elif method == "spearman":
+        correlation = selected_columns.corr(method="spearman").iloc[0, 1]
+    elif callable(method):
+        correlation = selected_columns.corr(method=method).iloc[0, 1]
+    else:
+        raise ValueError("Invalid correlation method. Please choose 'pearson', 'kendall', 'spearman', or a callable.")
+
+    return correlation

autogen/agentchat/contrib/captainagent/tools/data_analysis/calculate_skewness_and_kurtosis.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+def calculate_skewness_and_kurtosis(csv_file: str, column_name: str) -> tuple:
+    """Calculate the skewness and kurtosis of a specified column in a CSV file. The kurtosis is calculated using the Fisher definition.
+    The two metrics are computed using scipy.stats functions.
+
+    Args:
+    csv_file (str): The path to the CSV file.
+    column_name (str): The name of the column to calculate skewness and kurtosis for.
+
+    Returns:
+    tuple: (skewness, kurtosis)
+    """
+    import pandas as pd
+    from scipy.stats import kurtosis, skew
+
+    # Read the CSV file into a pandas DataFrame
+    df = pd.read_csv(csv_file)
+
+    # Extract the specified column
+    column = df[column_name]
+
+    # Calculate the skewness and kurtosis
+    skewness = skew(column)
+    kurt = kurtosis(column)
+
+    return skewness, kurt

autogen/agentchat/contrib/captainagent/tools/data_analysis/detect_outlier_iqr.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+def detect_outlier_iqr(csv_file: str, column_name: str):
+    """Detect outliers in a specified column of a CSV file using the IQR method.
+
+    Args:
+    csv_file (str): The path to the CSV file.
+    column_name (str): The name of the column to detect outliers in.
+
+    Returns:
+    list: A list of row indices that correspond to the outliers.
+    """
+    import pandas as pd
+
+    # Read the CSV file into a pandas DataFrame
+    df = pd.read_csv(csv_file)
+
+    # Calculate the quartiles and IQR for the specified column
+    q1 = df[column_name].quantile(0.25)
+    q3 = df[column_name].quantile(0.75)
+    iqr = q3 - q1
+
+    # Find the outliers based on the defined criteria
+    outliers = df[(df[column_name] < q1 - 1.5 * iqr) | (df[column_name] > q3 + 1.5 * iqr)]
+
+    # Return the row indices of the outliers
+    return outliers.index.tolist()

autogen/agentchat/contrib/captainagent/tools/data_analysis/detect_outlier_zscore.py

@@ -0,0 +1,28 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+def detect_outlier_zscore(csv_file, column_name, threshold=3):
+    """Detect outliers in a CSV file based on a specified column. The outliers are determined by calculating the z-score of the data points in the column.
+
+    Args:
+    csv_file (str): The path to the CSV file.
+    column_name (str): The name of the column to calculate z-scores for.
+    threshold (float, optional): The threshold value for determining outliers. By default set to 3.
+
+    Returns:
+    list: A list of row indices where the z-score is above the threshold.
+    """
+    import numpy as np
+    import pandas as pd
+
+    # Read the CSV file into a pandas DataFrame
+    df = pd.read_csv(csv_file)
+
+    # Calculate the z-score for the specified column
+    z_scores = np.abs((df[column_name] - df[column_name].mean()) / df[column_name].std())
+
+    # Find the row indices where the z-score is above the threshold
+    outlier_indices = np.where(z_scores > threshold)[0]
+
+    # Return the row indices of the outliers
+    return outlier_indices

autogen/agentchat/contrib/captainagent/tools/data_analysis/explore_csv.py

@@ -0,0 +1,21 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+def explore_csv(file_path, num_lines=5):
+    """Reads a CSV file and prints the column names, shape, data types, and the first few lines of data.
+
+    Args:
+        file_path (str): The path to the CSV file.
+        num_lines (int, optional): The number of lines to print. Defaults to 5.
+    """
+    import pandas as pd
+
+    df = pd.read_csv(file_path)
+    header = df.columns
+    print("Columns:")
+    print(", ".join(header))
+    print("Shape:", df.shape)
+    print("Data Types:")
+    print(df.dtypes)
+    print("First", num_lines, "lines:")
+    print(df.head(num_lines))

autogen/agentchat/contrib/captainagent/tools/data_analysis/shapiro_wilk_test.py

@@ -0,0 +1,30 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+from autogen.coding.func_with_reqs import with_requirements
+
+
+@with_requirements(["pandas", "scipy"])
+def shapiro_wilk_test(csv_file, column_name):
+    """Perform the Shapiro-Wilk test on a specified column of a CSV file.
+
+    Args:
+    csv_file (str): The path to the CSV file.
+    column_name (str): The name of the column to perform the test on.
+
+    Returns:
+    float: The p-value resulting from the Shapiro-Wilk test.
+    """
+    import pandas as pd
+    from scipy.stats import shapiro
+
+    # Read the CSV file into a pandas DataFrame
+    df = pd.read_csv(csv_file)
+
+    # Extract the specified column as a numpy array
+    column_data = df[column_name].values
+
+    # Perform the Shapiro-Wilk test
+    _, p_value = shapiro(column_data)
+
+    return p_value

autogen/agentchat/contrib/captainagent/tools/information_retrieval/arxiv_download.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+from autogen.coding.func_with_reqs import with_requirements
+
+
+@with_requirements(["arxiv"])
+def arxiv_download(id_list: list[str], download_dir="./") -> list[str]:
+    """Downloads PDF files from ArXiv based on a list of arxiv paper IDs.
+
+    Args:
+        id_list (list): A list of paper IDs to download. e.g. [2302.00006v1]
+        download_dir (str, optional): The directory to save the downloaded PDF files. Defaults to './'.
+
+    Returns:
+        list: A list of paths to the downloaded PDF files.
+    """
+    import arxiv
+
+    paths = []
+    for paper in arxiv.Client().results(arxiv.Search(id_list=id_list)):
+        path = paper.download_pdf(download_dir, filename=paper.get_short_id() + ".pdf")
+        paths.append(path)
+        print("Paper id:", paper.get_short_id(), "Downloaded to:", path)
+    return paths

autogen/agentchat/contrib/captainagent/tools/information_retrieval/arxiv_search.py

@@ -0,0 +1,53 @@
+# Copyright (c) 2023 - 2025, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+from autogen.coding.func_with_reqs import with_requirements
+
+
+@with_requirements(["arxiv"])
+def arxiv_search(query, max_results=10, sortby="relevance"):
+    """Search for articles on arXiv based on the given query.
+
+    Args:
+        query (str): The search query.
+        max_results (int, optional): The maximum number of results to retrieve. Defaults to 10.
+        sortby (str, optional): The sorting criterion for the search results. Can be 'relevance' or 'submittedDate'. Defaults to 'relevance'.
+
+    Returns:
+        list: A list of dictionaries containing information about the search results. Each dictionary contains the following keys:
+            - 'title': The title of the article.
+            - 'authors': The authors of the article.
+            - 'summary': The summary of the article.
+            - 'entry_id': The entry ID of the article.
+            - 'doi': The DOI of the article (If applicable).
+            - 'published': The publication date of the article in the format 'Y-M'.
+    """
+    import arxiv
+
+    def get_author(r):
+        return ", ".join(a.name for a in r.authors)
+
+    criterion = {"relevance": arxiv.SortCriterion.Relevance, "submittedDate": arxiv.SortCriterion.SubmittedDate}[sortby]
+
+    client = arxiv.Client()
+    search = arxiv.Search(query=query, max_results=max_results, sort_by=criterion)
+    res = []
+    results = client.results(search)
+    for r in results:
+        print("Entry id:", r.entry_id)
+        print("Title:", r.title)
+        print("Authors:", get_author(r))
+        print("DOI:", r.doi)
+        print("Published:", r.published.strftime("%Y-%m"))
+        # print("Summary:", r.summary)
+        res.append({
+            "title": r.title,
+            "authors": get_author(r),
+            "summary": r.summary,
+            "entry_id": r.entry_id,
+            "doi": r.doi,
+            "published": r.published.strftime("%Y-%m"),
+        })
+    return res