feat: js/py generate declarations from comments (#30)

scripts/build-declarations.js

@@ -0,0 +1,203 @@
+#!/usr/bin/env node
+
+const fs = require("fs");
+
+function main() {
+  const scriptfile = process.argv[2];
+  const contents = fs.readFileSync(process.argv[2], "utf8");
+  const functions = extractFunctions(contents);
+  let declarations = functions.map(({ funcName, jsdoc }) => {
+    const { description, params } = parseJsDoc(jsdoc, funcName);
+    const declaration = buildDeclaration(funcName, description, params);
+    return declaration;
+  });
+  const name = getBasename(scriptfile);
+  if (declarations.length > 0) {
+    declarations = declarations.slice(0, 1);
+    declarations[0].name = name;
+  }
+  console.log(JSON.stringify(declarations, null, 2));
+}
+
+/**
+ * @param {string} contents
+ * @param {bool} isTool
+ */
+function extractFunctions(contents, isTool = true) {
+  const output = [];
+  const lines = contents.split("\n");
+  let isInComment = false;
+  let jsdoc = "";
+  let incompleteComment = "";
+  for (let line of lines) {
+    if (/^\s*\/\*/.test(line)) {
+      isInComment = true;
+      incompleteComment += `\n${line}`;
+    } else if (/^\s*\*\//.test(line)) {
+      isInComment = false;
+      incompleteComment += `\n${line}`;
+      jsdoc = incompleteComment;
+      incompleteComment = "";
+    } else if (isInComment) {
+      incompleteComment += `\n${line}`;
+    } else {
+      if (!jsdoc || line.trim() === "") {
+        continue;
+      }
+      if (isTool) {
+        if (/function main/.test(line)) {
+          output.push({
+            funcName: "main",
+            jsdoc,
+          });
+        }
+      } else {
+        const match = /function *([_A-Za-z]+)/.exec(line);
+        if (match) {
+          const funcName = match[1];
+          if (!funcName.startsWith("_")) {
+            output.push({ funcName, jsdoc });
+          }
+        }
+      }
+      jsdoc = "";
+    }
+  }
+  return output;
+}
+
+/**
+ * @param {string} jsdoc
+ * @param {string} funcName,
+ */
+function parseJsDoc(jsdoc, funcName) {
+  const lines = jsdoc.split("\n");
+  let description = "";
+  const rawParams = [];
+  let tag = "";
+  for (let line of lines) {
+    line = line.replace(/^\s*(\/\*\*|\*\/|\*)/, "").trim();
+    let match = /^@(\w+)/.exec(line);
+    if (match) {
+      tag = match[1];
+    }
+    if (!tag) {
+      description += `\n${line}`;
+    } else if (tag == "property") {
+      if (match) {
+        rawParams.push(line.slice(tag.length + 1).trim());
+      } else {
+        rawParams[rawParams.length - 1] += `\n${line}`;
+      }
+    }
+  }
+  const params = [];
+  for (const rawParam of rawParams) {
+    try {
+      params.push(parseParam(rawParam));
+    } catch (err) {
+      throw new Error(
+        `Unable to parse function '${funcName}' of jsdoc '@property ${rawParam}'`,
+      );
+    }
+  }
+  return {
+    description: description.trim(),
+    params,
+  };
+}
+
+/**
+ * @typedef {ReturnType<parseParam>} Param
+ */
+
+/**
+ * @param {string} rawParam
+ */
+function parseParam(rawParam) {
+  const regex = /^{([^}]+)} +(\S+)( *- +| +)?/;
+  const match = regex.exec(rawParam);
+  if (!match) {
+    throw new Error(`Invalid jsdoc comment`);
+  }
+  const type = match[1];
+  let name = match[2];
+  const description = rawParam.replace(regex, "");
+
+  let required = true;
+  if (/^\[.*\]$/.test(name)) {
+    name = name.slice(1, -1);
+    required = false;
+  }
+  let property = buildProperty(type, description);
+  return { name, property, required };
+}
+
+/**
+ * @param {string} type
+ * @param {string} description
+ */
+function buildProperty(type, description) {
+  type = type.toLowerCase();
+  const property = {};
+  if (type.includes("|")) {
+    property.type = "string";
+    property.enum = type.replace(/'/g, "").split("|");
+  } else if (type === "boolean") {
+    property.type = "boolean";
+  } else if (type === "string") {
+    property.type = "string";
+  } else if (type === "integer") {
+    property.type = "integer";
+  } else if (type === "number") {
+    property.type = "number";
+  } else if (type === "string[]") {
+    property.type = "array";
+    property.items = { type: "string" };
+  } else {
+    throw new Error(`Unsupported type '${type}'`);
+  }
+  property.description = description;
+  return property;
+}
+
+/**
+ * @param {string} filePath
+ */
+function getBasename(filePath) {
+  const filenameWithExt = filePath.split(/[/\\]/).pop();
+
+  const lastDotIndex = filenameWithExt.lastIndexOf(".");
+
+  if (lastDotIndex === -1) {
+    return filenameWithExt;
+  }
+
+  return filenameWithExt.substring(0, lastDotIndex);
+}
+
+/**
+ * @param {string} name
+ * @param {string} description
+ * @param {Param[]} params
+ */
+function buildDeclaration(name, description, params) {
+  const schema = {
+    name,
+    description,
+    properties: {},
+  };
+  const requiredParams = [];
+  for (const { name, property, required } of params) {
+    schema.properties[name] = property;
+    if (required) {
+      requiredParams.push(name);
+    }
+  }
+  if (requiredParams.length > 0) {
+    schema.required = requiredParams;
+  }
+  return schema;
+}
+
+main();

scripts/build-declarations.py

@@ -0,0 +1,177 @@
+#!/usr/bin/env python
+
+import ast
+import os
+import json
+import re
+import sys
+from collections import OrderedDict
+
+
+def main():
+    scriptfile = sys.argv[1]
+    with open(scriptfile, "r", encoding="utf-8") as f:
+        contents = f.read()
+
+    functions = extract_functions(contents)
+    declarations = []
+    for function in functions:
+        func_name, docstring, func_args = function
+        description, params = parse_docstring(docstring)
+        declarations.append(
+            build_declaration(func_name, description, params, func_args)
+        )
+
+    name = os.path.splitext(os.path.basename(scriptfile))[0]
+    if declarations:
+        declarations = declarations[0:1]
+        declarations[0]["name"] = name
+
+    print(json.dumps(declarations, indent=2))
+
+
+def extract_functions(contents: str):
+    tree = ast.parse(contents)
+    output = []
+    for node in ast.walk(tree):
+        if not isinstance(node, ast.FunctionDef):
+            continue
+        func_name = node.name
+        if func_name.startswith("_"):
+            continue
+        docstring = ast.get_docstring(node) or ""
+        func_args = OrderedDict()
+        for arg in node.args.args:
+            arg_name = arg.arg
+            arg_type = get_arg_type(arg.annotation)
+            func_args[arg_name] = arg_type
+        output.append((func_name, docstring, func_args))
+    return output
+
+
+def get_arg_type(annotation) -> str:
+    if annotation is None:
+        return ""
+    elif isinstance(annotation, ast.Name):
+        return annotation.id
+    elif isinstance(annotation, ast.Subscript):
+        if isinstance(annotation.value, ast.Name):
+            type_name = annotation.value.id
+            if type_name == "List":
+                child = get_arg_type(annotation.slice)
+                return f"list[{child}]"
+            if type_name == "Literal":
+                literals = [ast.unparse(el) for el in annotation.slice.elts]
+                return f"{'|'.join(literals)}"
+            if type_name == "Optional":
+                child = get_arg_type(annotation.slice)
+                return f"{child}?"
+    return "any"
+
+
+def parse_docstring(docstring: str):
+    lines = docstring.splitlines()
+    description = ""
+    rawParams = []
+    is_in_args = False
+    for line in lines:
+        if not is_in_args:
+            if line.startswith("Args:"):
+                is_in_args = True
+            else:
+                description += f"\n{line}"
+            continue
+        else:
+            if re.search(r"^\s+", line):
+                rawParams.append(line.strip())
+            else:
+                break
+    params = {}
+    for rawParam in rawParams:
+        name, type_, description = parse_param(rawParam)
+        params[name] = (type_, description)
+    return (description.strip(), params)
+
+
+def parse_param(raw_param: str):
+    name = ""
+    description = ""
+    type_from_comment = ""
+    if ":" in raw_param:
+        name, description = raw_param.split(":", 1)
+        name = name.strip()
+        description = description.strip()
+    else:
+        name = raw_param
+    if " " in name:
+        name, type_from_comment = name.split(" ", 1)
+        type_from_comment = type_from_comment.strip()
+
+    if type_from_comment.startswith("(") and type_from_comment.endswith(")"):
+        type_from_comment = type_from_comment[1:-1]
+    type_parts = [value.strip() for value in type_from_comment.split(",")]
+    type_ = type_parts[0]
+    if "optional" in type_parts[1:]:
+        type_ = f"{type_}?"
+
+    return (name, type_, description)
+
+
+def build_declaration(
+    name: str, description: str, params: dict, args: OrderedDict[str, str]
+) -> dict[str, dict]:
+    schema = {
+        "name": name,
+        "description": description,
+        "properties": {},
+    }
+    required_params = []
+    for arg_name, arg_type in args.items():
+        type_ = arg_type
+        description = ""
+        required = True
+        if params.get(arg_name):
+            param_type, description = params[arg_name]
+            if not type_:
+                type_ = param_type
+        if type_.endswith("?"):
+            type_ = type_[:-1]
+            required = False
+        try:
+            property = build_property(type_, description)
+        except:
+            raise ValueError(f"Unable to parse arg '{arg_name}' of function '{name}'")
+        schema["properties"][arg_name] = property
+        if required:
+            required_params.append(arg_name)
+    if required_params:
+        schema["required"] = required_params
+    return schema
+
+
+def build_property(type_: str, description: str):
+    property = {}
+    if "|" in type_:
+        property["type"] = "string"
+        property["enum"] = type_.replace("'", "").split("|")
+    elif type_ == "bool":
+        property["type"] = "boolean"
+    elif type_ == "str":
+        property["type"] = "string"
+    elif type_ == "int":
+        property["type"] = "integer"
+    elif type_ == "float":
+        property["type"] = "number"
+    elif type_ == "list[str]":
+        property["type"] = "array"
+        property["items"] = {"type": "string"}
+    elif type_ == "":
+        property["type"] = "string"
+    else:
+        raise ValueError(f"Unsupported type `{type_}`")
+    property["description"] = description
+    return property
+
+
+if __name__ == "__main__":
+    main()

scripts/build-declarations.sh

@@ -0,0 +1,43 @@
+#!/usr/bin/env bash
+
+argc --argc-export "$1" | \
+jq -r '
+    def parse_description(flag_option):
+        if flag_option.describe == "" then
+            {}
+        else
+            { "description": flag_option.describe }
+        end;
+
+    def parse_enum(flag_option):
+        if flag_option.choice.type == "Values" then
+            { "enum": flag_option.choice.data }
+        else
+            {}
+        end;
+
+    def parse_property(flag_option):
+        [
+            { condition: (flag_option.flag == true), result: { type: "boolean" } },
+            { condition: (flag_option.multiple_occurs == true), result: { type: "array", items: { type: "string" } } },
+            { condition: (flag_option.notations[0] == "INT"), result: { type: "integer" } },
+            { condition: (flag_option.notations[0] == "NUM"), result: { type: "number" } },
+            { condition: true, result: { type: "string" } } ]
+        | map(select(.condition) | .result) | first 
+        | (. + parse_description(flag_option))
+        | (. + parse_enum(flag_option))
+        ;
+
+
+    def parse_parameter(flag_options):
+        {
+            type: "object",
+            properties: (reduce flag_options[] as $item ({}; . + { ($item.id | sub("-"; "_"; "g")): parse_property($item) })),
+            required: [flag_options[] | select(.required == true) | .id | sub("-"; "_"; "g")],
+        };
+
+    [{
+        name: (.name | sub("-"; "_"; "g")),
+        description: .describe,
+        parameters: parse_parameter([.flag_options[] | select(.id != "help" and .id != "version")])
+    }]'

scripts/create-tool.sh

@@ -1,17 +1,20 @@
 #!/usr/bin/env bash
 set -e
 
-# @describe Create a boilplate tool script file.
+# @describe Create a boilplate tool scriptfile.
+#
 # It automatically generate declaration json for `*.py` and `*.js` and generate `@option` tags for `.sh`.
 # Examples:
 #   argc create abc.sh foo bar! baz+ qux*
 #   ./scripts/create-tool.sh test.py foo bar! baz+ qux*
+#
+# @flag --force Override the exist tool file
 # @arg name! The script file name.
 # @arg params* The script parameters
 
 main() {
     output="tools/$argc_name"
-    if [[ -f "$output" ]]; then
+    if [[ -f "$output" ]] && [[ -z "$argc_force" ]]; then
         _die "$output already exists"
     fi
     ext="${argc_name##*.}"
@@ -25,6 +28,7 @@ main() {
     py) create_py ;;
     *) _die "Invalid extension name: $ext, must be one of ${support_exts[*]}" ;; 
     esac
+    _die "$output generated"
 }
 
 create_sh() {
@@ -49,25 +53,91 @@ EOF
 }
 
 create_js() {
+    properties=''
+    for param in "${argc_params[@]}"; do
+        if [[ "$param" == *'!' ]]; then
+            param="${param:0:$((${#param}-1))}"
+            property=" * @property {string} $param - "
+        elif [[ "$param" == *'+' ]]; then
+            param="${param:0:$((${#param}-1))}"
+            property=" * @property {string[]} $param - "
+        elif [[ "$param" == *'*' ]]; then
+            param="${param:0:$((${#param}-1))}"
+            property=" * @property {string[]} [$param] - "
+        else
+            property=" * @property {string} [$param] - "
+        fi
+        properties+=$'\n'"$property"
+    done
     cat <<EOF > "$output"
-exports.declarate = function declarate() {
-  return $(build_schema)
-}
-
-exports.execute = function execute(data) {
-  console.log(data)
+/**
+ *
+ * @typedef {Object} Args${properties}
+ * @param {Args} args
+ */
+exports.main = function main(args) {
+  console.log(args);
 }
 EOF
 }
 
 create_py() {
+    has_array_param=false
+    has_optional_pram=false
+    required_properties=''
+    optional_properties=''
+    required_arguments=()
+    optional_arguments=()
+    indent="    "
+    for param in "${argc_params[@]}"; do
+        optional=false
+        if [[ "$param" == *'!' ]]; then
+            param="${param:0:$((${#param}-1))}"
+            type="str"
+        elif [[ "$param" == *'+' ]]; then
+            param="${param:0:$((${#param}-1))}"
+            type="List[str]"
+            has_array_param=true
+        elif [[ "$param" == *'*' ]]; then
+            param="${param:0:$((${#param}-1))}"
+            type="Optional[List[str]] = None"
+            optional=true
+            has_array_param=true
+        else
+            optional=true
+            type="Optional[str] = None"
+        fi
+        if [[ "$optional" == "true" ]]; then
+            has_optional_pram=true
+            optional_arguments+="$param: $type, "
+            optional_properties+=$'\n'"$indent$indent$param: -"
+        else
+            required_arguments+="$param: $type, "
+            required_properties+=$'\n'"$indent$indent$param: -"
+        fi
+    done
+    import_typing_members=()
+    if [[ "$has_array_param" == "true" ]]; then
+        import_typing_members+=("List")
+    fi
+    if [[ "$has_optional_pram" == "true" ]]; then
+        import_typing_members+=("Optional")
+    fi
+    imports=""
+    if [[ -n "$import_typing_members" ]]; then
+        members="$(echo "${import_typing_members[*]}" | sed 's/ /, /')"
+        imports="from typing import $members"$'\n'
+    fi
+    if [[ -n "$imports" ]]; then
+        imports="$imports"$'\n'
+    fi
     cat <<EOF > "$output"
-def declarate():
-  return $(build_schema) 
-
-
-def execute(data):
-  print(data)
+${imports}
+def main(${required_arguments}${optional_arguments}):
+    """
+    Args:${required_properties}${optional_properties}
+    """
+    pass
 EOF
 }
 

scripts/run-tool.js

@@ -58,23 +58,18 @@ const [funcName, funcData] = parseArgv();
 
 process.env["LLM_FUNCTION_NAME"] = funcName;
 
-if (process.env["LLM_FUNCTION_ACTION"] == "declarate") {
-  const { declarate } = loadFunc(funcName);
-  console.log(JSON.stringify(declarate(), null, 2));
-} else {
-  if (!funcData) {
-    console.log("No json data");
-    process.exit(1);
-  }
-
-  let args;
-  try {
-    args = JSON.parse(funcData);
-  } catch {
-    console.log("Invalid json data");
-    process.exit(1);
-  }
-
-  const { execute } = loadFunc(funcName);
-  execute(args);
+if (!funcData) {
+  console.log("No json data");
+  process.exit(1);
 }
+
+let args;
+try {
+  args = JSON.parse(funcData);
+} catch {
+  console.log("Invalid json data");
+  process.exit(1);
+}
+
+const { main } = loadFunc(funcName);
+main(args);

scripts/run-tool.py

@@ -5,6 +5,7 @@ import json
 import sys
 import importlib.util
 
+
 def parse_argv():
     func_name = sys.argv[0]
     func_data = None
@@ -21,6 +22,7 @@ def parse_argv():
 
     return func_name, func_data
 
+
 def load_func(func_name):
     func_file_name = f"{func_name}.py"
     func_path = os.path.join(os.environ["LLM_FUNCTIONS_DIR"], f"tools/{func_file_name}")
@@ -33,20 +35,24 @@ def load_func(func_name):
         print(f"Invalid function: {func_file_name}")
         sys.exit(1)
 
+
 def load_env(file_path):
     try:
-        with open(file_path, 'r') as f:
+        with open(file_path, "r") as f:
             for line in f:
                 line = line.strip()
-                if line.startswith('#') or line == '':
+                if line.startswith("#") or line == "":
                     continue
 
-                key, *value = line.split('=')
-                os.environ[key.strip()] = '='.join(value).strip()
+                key, *value = line.split("=")
+                os.environ[key.strip()] = "=".join(value).strip()
     except FileNotFoundError:
         pass
 
-os.environ["LLM_FUNCTIONS_DIR"] = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+
+os.environ["LLM_FUNCTIONS_DIR"] = os.path.abspath(
+    os.path.join(os.path.dirname(__file__), "..")
+)
 
 load_env(os.path.join(os.environ["LLM_FUNCTIONS_DIR"], ".env"))
 
@@ -54,20 +60,16 @@ func_name, func_data = parse_argv()
 
 os.environ["LLM_FUNCTION_NAME"] = func_name
 
-if os.getenv("LLM_FUNCTION_ACTION") == "declarate":
-    module = load_func(func_name)
-    print(json.dumps(module.declarate(), indent=2))
-else:
-    if not func_data:
-        print("No json data")
-        sys.exit(1)
+if not func_data:
+    print("No json data")
+    sys.exit(1)
 
-    args = None
-    try:
-        args = json.loads(func_data)
-    except (json.JSONDecodeError, TypeError):
-        print("Invalid json data")
-        sys.exit(1)
+args = None
+try:
+    args = json.loads(func_data)
+except (json.JSONDecodeError, TypeError):
+    print("Invalid json data")
+    sys.exit(1)
 
-    module = load_func(func_name)
-    module.execute(args)
+module = load_func(func_name)
+module.main(**args)

scripts/run-tool.sh

@@ -27,76 +27,32 @@ if [[ "$OS" == "Windows_NT" ]]; then
     func_file="$(cygpath -w "$func_file")"
 fi
 
-if [[ "$LLM_FUNCTION_ACTION" == "declarate" ]]; then
-    argc --argc-export "$func_file" | \
+if [[ -z "$func_data" ]]; then
+    echo "No json data"
+    exit 1
+fi
+
+data="$(
+    echo "$func_data" | \
     $JQ -r '
-    def parse_description(flag_option):
-        if flag_option.describe == "" then
-            {}
-        else
-            { "description": flag_option.describe }
-        end;
-
-    def parse_enum(flag_option):
-        if flag_option.choice.type == "Values" then
-            { "enum": flag_option.choice.data }
-        else
-            {}
-        end;
-
-    def parse_property(flag_option):
-        [
-            { condition: (flag_option.flag == true), result: { type: "boolean" } },
-            { condition: (flag_option.multiple_occurs == true), result: { type: "array", items: { type: "string" } } },
-            { condition: (flag_option.notations[0] == "INT"), result: { type: "integer" } },
-            { condition: (flag_option.notations[0] == "NUM"), result: { type: "number" } },
-            { condition: true, result: { type: "string" } } ]
-        | map(select(.condition) | .result) | first 
-        | (. + parse_description(flag_option))
-        | (. + parse_enum(flag_option))
-        ;
-
-
-    def parse_parameter(flag_options):
-        {
-            type: "object",
-            properties: (reduce flag_options[] as $item ({}; . + { ($item.id | sub("-"; "_"; "g")): parse_property($item) })),
-            required: [flag_options[] | select(.required == true) | .id],
-        };
-
-    {
-        name: (.name | sub("-"; "_"; "g")),
-        description: .describe,
-        parameters: parse_parameter([.flag_options[] | select(.id != "help" and .id != "version")])
-    }'
-else
-    if [[ -z "$func_data" ]]; then
-        echo "No json data"
-        exit 1
+    to_entries | .[] | 
+    (.key | split("_") | join("-")) as $key |
+    if .value | type == "array" then
+        .value | .[] | "--\($key)\n\(. | @json)"
+    elif .value | type == "boolean" then
+        if .value then "--\($key)" else "" end
+    else
+        "--\($key)\n\(.value | @json)"
+    end'
+)" || {
+    echo "Invalid json data"
+    exit 1
+}
+while IFS= read -r line; do
+    if [[ "$line" == '--'* ]]; then
+        args+=("$line")
+    else
+        args+=("$(echo "$line" | $JQ -r '.')")
     fi
-
-    data="$(
-        echo "$func_data" | \
-        $JQ -r '
-        to_entries | .[] | 
-        (.key | split("_") | join("-")) as $key |
-        if .value | type == "array" then
-            .value | .[] | "--\($key)\n\(. | @json)"
-        elif .value | type == "boolean" then
-            if .value then "--\($key)" else "" end
-        else
-            "--\($key)\n\(.value | @json)"
-        end'
-    )" || {
-        echo "Invalid json data"
-        exit 1
-    }
-    while IFS= read -r line; do
-        if [[ "$line" == '--'* ]]; then
-            args+=("$line")
-        else
-            args+=("$(echo "$line" | $JQ -r '.')")
-        fi
-    done <<< "$data"
-    "$func_file" "${args[@]}"
-fi
+done <<< "$data"
+"$func_file" "${args[@]}"