From 6d639be2e5ba1270df88d255671842bca73601fe Mon Sep 17 00:00:00 2001 From: Duy Huynh Date: Mon, 8 Jul 2024 18:19:23 +0700 Subject: [PATCH] feat: add golang support --- README.md | 2 +- examples/code_runner.py | 27 +++++++++++++++++++-------- llm_sandbox/session.py | 34 ++++++++++++++++++++++++---------- llm_sandbox/utils.py | 10 ++++------ 4 files changed, 48 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index 0484ea1..992260f 100644 --- a/README.md +++ b/README.md @@ -129,7 +129,7 @@ Here is a list of things you can do to contribute: - [ ] Add Java maven support. - [x] Add support for JavaScript. - [x] Add support for C++. -- [ ] Add support for Go. +- [x] Add support for Go. - [ ] Add support for Ruby. - [ ] Add remote Docker host support. - [ ] Add remote Kubernetes cluster support. diff --git a/examples/code_runner.py b/examples/code_runner.py index 3c716c4..6dd7451 100644 --- a/examples/code_runner.py +++ b/examples/code_runner.py @@ -99,7 +99,7 @@ def run_cpp_code(): def run_go_code(): - with SandboxSession(lang="go", keep_template=False, verbose=True) as session: + with SandboxSession(lang="go", keep_template=True, verbose=True) as session: output = session.run( """ package main @@ -117,20 +117,31 @@ def run_go_code(): package main import ( "fmt" - "github.com/gorilla/mux" + "github.com/spyzhov/ajson" ) func main() { fmt.Println("Hello, World!") + json := []byte(`...`) + + root, _ := ajson.Unmarshal(json) + nodes, _ := root.JSONPath("$..price") + for _, node := range nodes { + node.SetNumeric(node.MustNumeric() * 1.25) + node.Parent().AppendObject("currency", ajson.StringNode("", "EUR")) + } + result, _ := ajson.Marshal(root) + + fmt.Printf("%s", result) } """, - libraries=["github.com/gorilla/mux"], + libraries=["github.com/spyzhov/ajson"], ) print(output) if __name__ == "__main__": - run_python_code() - run_java_code() - run_javascript_code() - run_cpp_code() - # run_go_code() \ No newline at end of file + # run_python_code() + # run_java_code() + # run_javascript_code() + # run_cpp_code() + run_go_code() diff --git a/llm_sandbox/session.py b/llm_sandbox/session.py index 287ebf6..9932e34 100644 --- a/llm_sandbox/session.py +++ b/llm_sandbox/session.py @@ -139,27 +139,36 @@ def run(self, code: str, libraries: Optional[List] = None): ) if self.lang == SupportedLanguage.GO: - init_cmd = "sh -c 'cd /sandbox && go mod init example.com/sandbox && go mod tidy'" - self.execute_command(init_cmd) + self.execute_command("mkdir -p /example") + self.execute_command("go mod init example", worKdir="/example") + self.execute_command("go mod tidy", worKdir="/example") for lib in libraries: - install_cmd = f"sh -c 'cd /sandbox && go get {lib}'" - self.execute_command(install_cmd) + command = get_libraries_installation_command(self.lang, lib) + self.execute_command(command, worKdir="/example") else: for lib in libraries: command = get_libraries_installation_command(self.lang, lib) self.execute_command(command) code_file = f"/tmp/code.{get_code_file_extension(self.lang)}" + if self.lang == SupportedLanguage.GO: + code_dest_file = "/example/code.go" + else: + code_dest_file = code_file + with open(code_file, "w") as f: f.write(code) - self.copy_to_runtime(code_file, code_file) + self.copy_to_runtime(code_file, code_dest_file) output = "" - commands = get_code_execution_command(self.lang, code_file) + commands = get_code_execution_command(self.lang, code_dest_file) for command in commands: - output = self.execute_command(command) + if self.lang == SupportedLanguage.GO: + output = self.execute_command(command, worKdir="/example") + else: + output = self.execute_command(command) return output @@ -204,7 +213,7 @@ def copy_to_runtime(self, src: str, dest: str): tarstream.seek(0) self.container.put_archive(os.path.dirname(dest), tarstream) - def execute_command(self, command: Optional[str]): + def execute_command(self, command: Optional[str], worKdir: Optional[str] = None): if not command: raise ValueError("Command cannot be empty") @@ -216,9 +225,14 @@ def execute_command(self, command: Optional[str]): if self.verbose: print(f"Executing command: {command}") - _, exec_log = self.container.exec_run(command, stream=True, tty=True) - output = "" + if worKdir: + _, exec_log = self.container.exec_run( + command, stream=True, tty=True, workdir=worKdir + ) + else: + _, exec_log = self.container.exec_run(command, stream=True, tty=True) + output = "" if self.verbose: print("Output:", end=" ") diff --git a/llm_sandbox/utils.py b/llm_sandbox/utils.py index 0a27ce6..62452f4 100644 --- a/llm_sandbox/utils.py +++ b/llm_sandbox/utils.py @@ -1,6 +1,6 @@ import docker import docker.errors -from typing import List, Optional +from typing import Optional from docker import DockerClient from llm_sandbox.const import SupportedLanguage @@ -22,13 +22,11 @@ def image_exists(client: DockerClient, image: str) -> bool: raise e -def get_libraries_installation_command( - lang: str, library: str -) -> Optional[str]: +def get_libraries_installation_command(lang: str, library: str) -> Optional[str]: """ Get the command to install libraries for the given language :param lang: Programming language - :param libraries: List of libraries + :param library: List of libraries :return: Installation command """ if lang == SupportedLanguage.PYTHON: @@ -40,7 +38,7 @@ def get_libraries_installation_command( elif lang == SupportedLanguage.CPP: return f"apt-get install {library}" elif lang == SupportedLanguage.GO: - return f"go get {library}" + return f"go get -u {library}" elif lang == SupportedLanguage.RUBY: return f"gem install {library}" else: