Move custom tool template scripts from Azure/promptflow to microsoft/promptflow (#74)

Co-authored-by: yalu4 <[email protected]>

Author: 16oeahr

scripts/generate_tool_package_template.py

@@ -0,0 +1,149 @@
+import argparse
+import os
+import re
+from jinja2 import Environment, FileSystemLoader
+
+
+def make_pythonic_variable_name(input_string):
+    variable_name = input_string.strip()
+    variable_name = re.sub(r'\W|^(?=\d)', '_', variable_name)
+    if not variable_name[0].isalpha() and variable_name[0] != '_':
+        variable_name = f'_{variable_name}'
+
+    return variable_name
+
+
+def convert_tool_name_to_class_name(tool_name):
+    return ''.join(word.title() for word in tool_name.split('_'))
+
+
+def create_file(path):
+    with open(path, 'w'):
+        pass
+
+
+def create_folder(path):
+    os.makedirs(path, exist_ok=True)
+
+
+def create_tool_project_structure(destination: str, package_name: str, tool_name: str,
+                                  function_name: str, is_class_way=False):
+    if is_class_way:
+        class_name = convert_tool_name_to_class_name(tool_name)
+
+    # Load templates
+    file_loader = FileSystemLoader('scripts\\templates')
+    env = Environment(loader=file_loader)
+
+    # Create new directory
+    if os.path.exists(destination):
+        print("Destination already exists. Please choose another one.")
+        return
+
+    os.makedirs(destination, exist_ok=True)
+
+    # Generate setup.py
+    template = env.get_template('setup.py.j2')
+    output = template.render(package_name=package_name, tool_name=tool_name)
+    with open(os.path.join(destination, 'setup.py'), 'w') as f:
+        f.write(output)
+
+    # Generate MANIFEST.in
+    template = env.get_template('MANIFEST.in.j2')
+    output = template.render(package_name=package_name)
+    with open(os.path.join(destination, 'MANIFEST.in'), 'w') as f:
+        f.write(output)
+
+    # Create tools folder and __init__.py, tool.py inside it
+    tools_dir = os.path.join(destination, package_name, 'tools')
+    create_folder(tools_dir)
+    create_file(os.path.join(tools_dir, '__init__.py'))
+    with open(os.path.join(tools_dir, '__init__.py'), 'w') as f:
+        f.write('__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore\n')
+
+    # Generate tool.py
+    if is_class_way:
+        template = env.get_template('tool2.py.j2')
+        output = template.render(class_name=class_name, function_name=function_name)
+    else:
+        template = env.get_template('tool.py.j2')
+        output = template.render(function_name=function_name)
+    with open(os.path.join(tools_dir, f'{tool_name}.py'), 'w') as f:
+        f.write(output)
+
+    # Generate utils.py
+    template = env.get_template('utils.py.j2')
+    output = template.render()
+    with open(os.path.join(tools_dir, 'utils.py'), 'w') as f:
+        f.write(output)
+
+    create_file(os.path.join(destination, package_name, '__init__.py'))
+    with open(os.path.join(destination, package_name, '__init__.py'), 'w') as f:
+        f.write('__path__ = __import__("pkgutil").extend_path(__path__, __name__)  # type: ignore\n')
+
+    # Create yamls folder and __init__.py inside it
+    yamls_dir = os.path.join(destination, package_name, 'yamls')
+    create_folder(yamls_dir)
+
+    # Create tool yaml
+    if is_class_way:
+        template = env.get_template('tool2.yaml.j2')
+        output = template.render(package_name=package_name, tool_name=tool_name, class_name=class_name,
+                                 function_name=function_name)
+    else:
+        template = env.get_template('tool.yaml.j2')
+        output = template.render(package_name=package_name, tool_name=tool_name, function_name=function_name)
+    with open(os.path.join(yamls_dir, f'{tool_name}.yaml'), 'w') as f:
+        f.write(output)
+
+    # Create test folder and __init__.py inside it
+    tests_dir = os.path.join(destination, 'tests')
+    create_folder(tests_dir)
+    create_file(os.path.join(tests_dir, '__init__.py'))
+
+    # Create test_tool.py
+    if is_class_way:
+        template = env.get_template('test_tool2.py.j2')
+        output = template.render(package_name=package_name, tool_name=tool_name, class_name=class_name,
+                                 function_name=function_name)
+    else:
+        template = env.get_template('test_tool.py.j2')
+        output = template.render(package_name=package_name, tool_name=tool_name, function_name=function_name)
+    with open(os.path.join(tests_dir, f'test_{tool_name}.py'), 'w') as f:
+        f.write(output)
+
+    print(f'Generated tool package template for {package_name} at {destination}')
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="promptflow tool template generation arguments.")
+
+    parser.add_argument("--package-name", "-p", type=str, help="your tool package's name", required=True)
+    parser.add_argument("--destination", "-d", type=str,
+                        help="target folder you want to place the generated template", required=True)
+    parser.add_argument("--tool-name", "-t", type=str,
+                        help="your tool's name, by default is hello_world_tool", required=False)
+    parser.add_argument("--function-name", "-f", type=str,
+                        help="your tool's function name, by default is your tool's name", required=False)
+    parser.add_argument("--use-class", action='store_true', help="Specify whether to use a class implementation way.")
+
+    args = parser.parse_args()
+
+    destination = args.destination
+
+    package_name = make_pythonic_variable_name(args.package_name)
+    package_name = package_name.lower()
+
+    if args.tool_name:
+        tool_name = make_pythonic_variable_name(args.tool_name)
+    else:
+        tool_name = 'hello_world_tool'
+    tool_name = tool_name.lower()
+
+    if args.function_name:
+        function_name = make_pythonic_variable_name(args.function_name)
+    else:
+        function_name = tool_name
+    function_name = function_name.lower()
+
+    create_tool_project_structure(destination, package_name, tool_name, function_name, args.use_class)

scripts/templates/MANIFEST.in.j2

@@ -0,0 +1 @@
+include {{ package_name }}/yamls/*.yaml

scripts/templates/setup.py.j2

@@ -0,0 +1,14 @@
+from setuptools import find_packages, setup
+
+PACKAGE_NAME = "{{ package_name }}"
+
+setup(
+    name=PACKAGE_NAME,
+    version="0.0.1",
+    description="This is my tools package",
+    packages=find_packages(),
+    entry_points={
+        "package_tools": ["{{ tool_name }} = {{ package_name }}.tools.utils:list_package_tools"],
+    },
+    include_package_data=True,   # This line tells setuptools to include files from MANIFEST.in
+)

scripts/templates/test_tool.py.j2

@@ -0,0 +1,28 @@
+import pytest
+import unittest
+
+from promptflow.connections import CustomConnection
+from {{ package_name }}.tools.{{ tool_name }} import {{ function_name }}
+
+
+@pytest.fixture
+def my_custom_connection() -> CustomConnection:
+    my_custom_connection = CustomConnection(
+        {
+            "api-key" : "my-api-key",
+            "api-secret" : "my-api-secret",
+            "api-url" : "my-api-url"
+        }
+    )
+    return my_custom_connection
+
+
+class TestTool:
+    def test_{{ function_name }}(self, my_custom_connection):
+        result = {{ function_name }}(my_custom_connection, input_text="Microsoft")
+        assert result == "Hello Microsoft"
+
+
+# Run the unit tests
+if __name__ == "__main__":
+    unittest.main()

scripts/templates/test_tool2.py.j2

@@ -0,0 +1,27 @@
+import pytest
+import unittest
+
+from {{ package_name }}.tools.{{ tool_name }} import {{ class_name }}
+
+
+@pytest.fixture
+def my_url() -> str:
+    my_url = "https://www.bing.com"
+    return my_url
+
+
+@pytest.fixture
+def my_tool_provider(my_url) -> {{ class_name }}:
+    my_tool_provider = {{ class_name }}(my_url)
+    return my_tool_provider
+
+
+class TestTool:
+    def test_{{ tool_name }}(self, my_tool_provider):
+        result = my_tool_provider.{{ function_name }}(query="Microsoft")
+        assert result == "Hello Microsoft"
+
+
+# Run the unit tests
+if __name__ == "__main__":
+    unittest.main()

scripts/templates/tool.py.j2

@@ -0,0 +1,11 @@
+from promptflow import tool
+from promptflow.connections import CustomConnection
+
+
+@tool
+def {{ function_name }}(connection: CustomConnection, input_text: str) -> str:
+    # Replace with your tool code.
+    # Usually connection contains configs to connect to an API.
+    # Use CustomConnection is a dict. You can use it like: connection.api_key, connection.api_base
+    # Not all tools need a connection. You can remove it if you don't need it.
+    return "Hello " + input_text

scripts/templates/tool.yaml.j2

@@ -0,0 +1,13 @@
+{{ package_name }}.tools.{{ tool_name }}.{{ function_name }}:
+  function: {{ function_name }}
+  inputs:
+    connection:
+      type:
+      - CustomConnection
+    input_text:
+      type:
+      - string
+  module: {{ package_name }}.tools.{{ tool_name }}
+  name: Hello World Tool
+  description: This is hello world tool
+  type: python

scripts/templates/tool2.py.j2

@@ -0,0 +1,15 @@
+from promptflow import ToolProvider, tool
+import urllib.request
+
+
+class {{ class_name }}(ToolProvider):
+
+    def __init__(self, url: str):
+        super().__init__()
+        # Load content from url might be slow, so we do it in __init__ method to make sure it is loaded only once.
+        self.content = urllib.request.urlopen(url).read()
+
+    @tool
+    def {{ function_name }}(self, query: str) -> str:
+        # Replace with your tool code.
+        return "Hello " + query

scripts/templates/tool2.yaml.j2

@@ -0,0 +1,14 @@
+{{ package_name }}.tools.{{ tool_name }}.{{ class_name }}.{{ function_name }}:
+  class_name: {{ class_name }}
+  function: {{ function_name }}
+  inputs:
+    url:
+      type:
+      - string
+    query:
+      type:
+      - string
+  module: {{ package_name }}.tools.{{ tool_name }}
+  name: Hello World Tool
+  description: This is hello world tool
+  type: python

scripts/templates/utils.py.j2

@@ -0,0 +1,18 @@
+import yaml
+from pathlib import Path
+
+
+def collect_tools_from_directory(base_dir) -> dict:
+    tools = {}
+    for f in Path(base_dir).glob("**/*.yaml"):
+        with open(f, "r") as f:
+            tools_in_file = yaml.safe_load(f)
+            for identifier, tool in tools_in_file.items():
+                tools[identifier] = tool
+    return tools
+
+
+def list_package_tools():
+    """List package tools"""
+    yaml_dir = Path(__file__).parents[1] / "yamls"
+    return collect_tools_from_directory(yaml_dir)