feat: supports functions in bash/js/python/ruby (#6)

3 files changed, 115 insertions, 0 deletions

diff --git a/cmd/cmd.js b/cmd/cmd.js
new file mode 100755
index 0000000..81eedc6
--- /dev/null
+++ b/cmd/cmd.js
@@ -0,0 +1,36 @@
+#!/usr/bin/env node
+
+function loadModule() {
+  const path = require("path");
+  let func_name = process.argv[1];
+  if (func_name.endsWith("cmd.js")) {
+    func_name = process.argv[2]
+  } else {
+    func_name = path.basename(func_name)
+  }
+  if (!func_name.endsWith(".js")) {
+    func_name += '.js'
+  }
+  const func_path = path.resolve(__dirname, `../js/${func_name}`)
+  try {
+    return require(func_path);
+  } catch {
+    console.log(`Invalid js function: ${func_name}`)
+    process.exit(1)
+  }
+}
+
+if (process.env["LLM_FUNCTION_DECLARATE"]) {
+  const { declarate } = loadModule();
+  console.log(JSON.stringify(declarate(), null, 2))
+} else {
+  let data = null;
+  try {
+    data = JSON.parse(process.env["LLM_FUNCTION_DATA"])
+  } catch {
+    console.log("Invalid LLM_FUNCTION_DATA")
+    process.exit(1)
+  }
+  const { execute } = loadModule();
+  execute(data)
+}
\ No newline at end of file
diff --git a/cmd/cmd.py b/cmd/cmd.py
new file mode 100755
index 0000000..be0fc93
--- /dev/null
+++ b/cmd/cmd.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python
+
+import os
+import json
+import sys
+import importlib.util
+
+def load_module(func_name):
+    base_dir = os.path.dirname(os.path.abspath(__file__))
+    func_path = os.path.join(base_dir, f"../py/{func_name}")
+    if os.path.exists(func_path):
+        spec = importlib.util.spec_from_file_location(func_name, func_path)
+        module = importlib.util.module_from_spec(spec)
+        spec.loader.exec_module(module)
+        return module
+    else:
+        print(f"Invalid py function: {func_name}")
+        sys.exit(1)
+
+func_name = sys.argv[0]
+if func_name.endswith("cmd.py"):
+    func_name = sys.argv[1]
+else:
+    func_name = os.path.basename(func_name)
+
+if not func_name.endswith(".py"):
+    func_name += ".py"
+
+if os.getenv("LLM_FUNCTION_DECLARATE"):
+    module = load_module(func_name)
+    declarate = getattr(module, 'declarate')
+    print(json.dumps(declarate(), indent=2))
+else:
+    data = None
+    try:
+        data = json.loads(os.getenv("LLM_FUNCTION_DATA"))
+    except (json.JSONDecodeError, TypeError):
+        print("Invalid LLM_FUNCTION_DATA")
+        sys.exit(1)
+
+    module = load_module(func_name)
+    execute = getattr(module, 'execute')
+    execute(data)
\ No newline at end of file
diff --git a/cmd/cmd.rb b/cmd/cmd.rb
new file mode 100755
index 0000000..cb67f4d
--- /dev/null
+++ b/cmd/cmd.rb
@@ -0,0 +1,36 @@
+#!/usr/bin/env ruby
+
+require 'json'
+require 'pathname'
+
+def load_module
+  if __FILE__.end_with?("cmd.rb")
+    func_name = ARGV[0]
+  else
+    func_name = Pathname.new(__FILE__).basename.to_s
+  end
+
+  func_name += '.rb' unless func_name.end_with?('.rb')
+  func_path = File.expand_path("../rb/#{func_name}", __dir__)
+
+  begin
+    return require_relative func_path
+  rescue LoadError
+    puts "Invalid ruby function: #{func_name}"
+    exit 1
+  end
+end
+
+if ENV["LLM_FUNCTION_DECLARATE"]
+  declarate = load_module.method(:declarate)
+  puts JSON.pretty_generate(declarate.call)
+else
+  begin
+    data = JSON.parse(ENV["LLM_FUNCTION_DATA"])
+  rescue JSON::ParserError
+    puts "Invalid LLM_FUNCTION_DATA"
+    exit 1
+  end
+  execute = load_module.method(:execute)
+  execute.call(data)
+end
\ No newline at end of file