Skip to content

Commit 45244a7

Browse files
sincerity-beingsincerity-being
authored
feat(graph): update graph code node docker execution (alibaba#1056)
1 parent db36e2d commit 45244a7

File tree

6 files changed

+240
-28
lines changed

6 files changed

+240
-28
lines changed

‎spring-ai-alibaba-graph/spring-ai-alibaba-graph-core/pom.xml‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,6 +97,26 @@
9797
<version>${fastjson.version}</version>
9898
</dependency>
9999

100+
<dependency>
101+
<groupId>com.github.docker-java</groupId>
102+
<artifactId>docker-java</artifactId>
103+
<version>3.3.3</version>
104+
</dependency>
105+
106+
<dependency>
107+
<groupId>com.github.docker-java</groupId>
108+
<artifactId>docker-java-core</artifactId>
109+
<version>3.3.3</version>
110+
</dependency>
111+
112+
<dependency>
113+
<groupId>com.github.docker-java</groupId>
114+
<artifactId>docker-java-transport-httpclient5</artifactId>
115+
<version>3.3.3</version>
116+
</dependency>
117+
118+
119+
100120
<dependency>
101121
<groupId>org.apache.commons</groupId>
102122
<artifactId>commons-collections4</artifactId>

‎spring-ai-alibaba-graph/spring-ai-alibaba-graph-core/src/main/java/com/alibaba/cloud/ai/graph/node/code/DockerCodeExecutor.java‎

Lines changed: 151 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,151 @@
1+
/*
2+
* Copyright 2024-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.alibaba.cloud.ai.graph.node.code;
17+
18+
import com.alibaba.cloud.ai.graph.node.code.entity.CodeBlock;
19+
import com.alibaba.cloud.ai.graph.node.code.entity.CodeExecutionConfig;
20+
import com.alibaba.cloud.ai.graph.node.code.entity.CodeExecutionResult;
21+
import com.alibaba.cloud.ai.graph.utils.CodeUtils;
22+
import com.github.dockerjava.api.DockerClient;
23+
import com.github.dockerjava.api.async.ResultCallbackTemplate;
24+
import com.github.dockerjava.api.command.CreateContainerCmd;
25+
import com.github.dockerjava.api.command.CreateContainerResponse;
26+
import com.github.dockerjava.api.command.InspectContainerResponse;
27+
import com.github.dockerjava.api.model.Bind;
28+
import com.github.dockerjava.api.model.Frame;
29+
import com.github.dockerjava.api.model.Volume;
30+
import com.github.dockerjava.core.DockerClientBuilder;
31+
import org.slf4j.Logger;
32+
import org.slf4j.LoggerFactory;
33+
import org.apache.commons.codec.digest.DigestUtils;
34+
import com.alibaba.cloud.ai.graph.utils.FileUtils;
35+
36+
import java.util.List;
37+
import java.util.Objects;
38+
import java.util.concurrent.TimeUnit;
39+
40+
import static com.github.dockerjava.api.model.HostConfig.newHostConfig;
41+
42+
/**
43+
* @author HeYQ
44+
* @since 2025-06-01 20:15
45+
*/
46+
47+
public class DockerCodeExecutor implements CodeExecutor {
48+
49+
private static final Logger logger = LoggerFactory.getLogger(DockerCodeExecutor.class);
50+
51+
@Override
52+
public CodeExecutionResult executeCodeBlocks(List<CodeBlock> codeBlockList, CodeExecutionConfig codeExecutionConfig)
53+
throws Exception {
54+
StringBuilder allLogs = new StringBuilder();
55+
CodeExecutionResult result;
56+
57+
// Create Docker client
58+
try (DockerClient dockerClient = DockerClientBuilder.getInstance().build()) {
59+
60+
for (CodeBlock codeBlock : codeBlockList) {
61+
String language = codeBlock.language();
62+
String code = codeBlock.code();
63+
logger.info("\n>>>>>>>> EXECUTING CODE BLOCK (inferred language is {})...", language);
64+
65+
// Generate unique filename for each code block
66+
String codeHash = DigestUtils.md5Hex(code);
67+
String fileExt = CodeUtils.getFileExtForLanguage(language);
68+
String filename = String.format("tmp_code_%s.%s", codeHash, fileExt);
69+
70+
// Write code to working directory
71+
String hostWorkDir = codeExecutionConfig.getWorkDir();
72+
FileUtils.writeCodeToFile(hostWorkDir, filename, code);
73+
74+
// Create and configure container
75+
// Mount host directory to container's /workspace directory
76+
Volume containerVolume = new Volume("/workspace");
77+
Bind volumeBind = new Bind(hostWorkDir, containerVolume);
78+
79+
CreateContainerCmd createContainerCmd = dockerClient.createContainerCmd(codeExecutionConfig.getDocker())
80+
.withName(codeExecutionConfig.getContainerName() + "_" + codeBlockList.indexOf(codeBlock))
81+
.withCmd(CodeUtils.getExecutableForLanguage(language), filename)
82+
.withWorkingDir("/workspace")
83+
.withHostConfig(newHostConfig().withBinds(volumeBind));
84+
CreateContainerResponse container = createContainerCmd.exec();
85+
86+
try {
87+
// Start container
88+
dockerClient.startContainerCmd(container.getId()).exec();
89+
90+
// Wait for container execution to complete
91+
dockerClient.waitContainerCmd(container.getId())
92+
.start()
93+
.awaitCompletion(codeExecutionConfig.getTimeout(), TimeUnit.SECONDS);
94+
95+
// Get container logs
96+
String logs = dockerClient.logContainerCmd(container.getId())
97+
.withStdOut(true)
98+
.withStdErr(true)
99+
.exec(new LogContainerResultCallback())
100+
.toString();
101+
102+
// Get container exit code
103+
InspectContainerResponse containerInfo = dockerClient.inspectContainerCmd(container.getId()).exec();
104+
int exitCode = Objects.requireNonNull(containerInfo.getState().getExitCodeLong()).intValue();
105+
106+
// Append logs
107+
allLogs.append("\n").append(logs.trim());
108+
109+
// If execution failed, return result immediately
110+
if (exitCode != 0) {
111+
return new CodeExecutionResult(exitCode, allLogs.toString());
112+
}
113+
}
114+
finally {
115+
// Clean up container
116+
dockerClient.removeContainerCmd(container.getId()).withForce(true).exec();
117+
// Delete temporary file
118+
FileUtils.deleteFile(codeExecutionConfig.getWorkDir(), filename);
119+
}
120+
}
121+
122+
return new CodeExecutionResult(0, allLogs.toString());
123+
}
124+
catch (Exception e) {
125+
logger.error("Error executing code in Docker container", e);
126+
throw new RuntimeException("Error executing code in Docker container: " + e.getMessage(), e);
127+
}
128+
}
129+
130+
@Override
131+
public void restart() {
132+
133+
}
134+
135+
private static class LogContainerResultCallback extends ResultCallbackTemplate<LogContainerResultCallback, Frame> {
136+
137+
private final StringBuilder log = new StringBuilder();
138+
139+
@Override
140+
public void onNext(Frame frame) {
141+
log.append(new String(frame.getPayload()));
142+
}
143+
144+
@Override
145+
public String toString() {
146+
return log.toString();
147+
}
148+
149+
}
150+
151+
}

‎spring-ai-alibaba-graph/spring-ai-alibaba-graph-core/src/main/java/com/alibaba/cloud/ai/graph/node/code/LocalCommandlineCodeExecutor.java‎

Lines changed: 3 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
import com.alibaba.cloud.ai.graph.node.code.entity.CodeBlock;
1919
import com.alibaba.cloud.ai.graph.node.code.entity.CodeExecutionConfig;
2020
import com.alibaba.cloud.ai.graph.node.code.entity.CodeExecutionResult;
21+
import com.alibaba.cloud.ai.graph.utils.CodeUtils;
2122
import com.alibaba.cloud.ai.graph.utils.FileUtils;
2223
import org.apache.commons.codec.digest.DigestUtils;
2324
import org.apache.commons.exec.CommandLine;
@@ -76,7 +77,7 @@ public CodeExecutionResult executeCode(String language, String code, CodeExecuti
7677
}
7778
String workDir = config.getWorkDir();
7879
String codeHash = DigestUtils.md5Hex(code);
79-
String fileExt = getFileExtForLanguage(language);
80+
String fileExt = CodeUtils.getFileExtForLanguage(language);
8081
String filename = String.format("tmp_code_%s.%s", codeHash, fileExt);
8182

8283
// write the code string to a file specified by the filename.
@@ -91,7 +92,7 @@ public CodeExecutionResult executeCode(String language, String code, CodeExecuti
9192
private CodeExecutionResult executeCodeLocally(String language, String workDir, String filename, int timeout)
9293
throws Exception {
9394
// set up the command based on language
94-
String executable = getExecutableForLanguage(language);
95+
String executable = CodeUtils.getExecutableForLanguage(language);
9596
CommandLine commandLine = new CommandLine(executable);
9697
commandLine.addArgument(filename);
9798

@@ -129,22 +130,4 @@ private CodeExecutionResult executeCodeLocally(String language, String workDir,
129130
}
130131
}
131132

132-
private String getExecutableForLanguage(String language) throws Exception {
133-
return switch (language) {
134-
case "python3", "python" -> language;
135-
case "shell", "bash", "sh", "powershell" -> "sh";
136-
case "nodejs" -> "node";
137-
default -> throw new Exception("Language not recognized in code execution:" + language);
138-
};
139-
}
140-
141-
private String getFileExtForLanguage(String language) throws Exception {
142-
return switch (language) {
143-
case "python3", "python" -> "py";
144-
case "shell", "bash", "sh", "powershell" -> "sh";
145-
case "nodejs" -> "js";
146-
default -> throw new Exception("Language not recognized in code execution:" + language);
147-
};
148-
}
149-
150133
}

‎spring-ai-alibaba-graph/spring-ai-alibaba-graph-core/src/main/java/com/alibaba/cloud/ai/graph/node/code/entity/CodeExecutionConfig.java‎

Lines changed: 21 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,21 +16,26 @@
1616

1717
package com.alibaba.cloud.ai.graph.node.code.entity;
1818

19+
/**
20+
* @author HeYQ
21+
*/
1922
public class CodeExecutionConfig {
2023

21-
private String workDir = "extensions";
24+
private String workDir = "workspace";
2225

2326
/**
2427
* the docker image to use for code execution.
2528
*/
2629
private String docker;
2730

31+
private String containerName = "spring-ai-alibaba-container";
32+
33+
private String dockerHost = "unix:///var/run/docker.sock";
34+
2835
private int timeout = 600;
2936

3037
private int lastMessagesNumber = 1;
3138

32-
private int codeMaxDepth = 5;
33-
3439
public String getWorkDir() {
3540
return workDir;
3641
}
@@ -67,12 +72,21 @@ public CodeExecutionConfig setLastMessagesNumber(int lastMessagesNumber) {
6772
return this;
6873
}
6974

70-
public int getCodeMaxDepth() {
71-
return codeMaxDepth;
75+
public String getContainerName() {
76+
return containerName;
77+
}
78+
79+
public CodeExecutionConfig setCodeExecutionConfig(String containerName) {
80+
this.containerName = containerName;
81+
return this;
82+
}
83+
84+
public String getDockerHost() {
85+
return dockerHost;
7286
}
7387

74-
public CodeExecutionConfig setCodeMaxDepth(int codeMaxDepth) {
75-
this.codeMaxDepth = codeMaxDepth;
88+
public CodeExecutionConfig setDockerHost(String dockerHost) {
89+
this.dockerHost = dockerHost;
7690
return this;
7791
}
7892

‎spring-ai-alibaba-graph/spring-ai-alibaba-graph-core/src/main/java/com/alibaba/cloud/ai/graph/utils/CodeUtils.java‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
/*
2+
* Copyright 2024-2025 the original author or authors.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* https://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.alibaba.cloud.ai.graph.utils;
17+
18+
/**
19+
* @author HeYQ
20+
* @since 2025-06-01 20:35
21+
*/
22+
23+
public class CodeUtils {
24+
25+
public static String getExecutableForLanguage(String language) throws Exception {
26+
return switch (language) {
27+
case "python3", "python" -> language;
28+
case "shell", "bash", "sh", "powershell" -> "sh";
29+
case "nodejs" -> "node";
30+
default -> throw new Exception("Language not recognized in code execution:" + language);
31+
};
32+
}
33+
34+
public static String getFileExtForLanguage(String language) throws Exception {
35+
return switch (language) {
36+
case "python3", "python" -> "py";
37+
case "shell", "bash", "sh", "powershell" -> "sh";
38+
case "nodejs" -> "js";
39+
default -> throw new Exception("Language not recognized in code execution:" + language);
40+
};
41+
}
42+
43+
}

‎spring-ai-alibaba-graph/spring-ai-alibaba-graph-core/src/test/java/com/alibaba/cloud/ai/graph/CodeActionTest.java‎ renamed to ‎spring-ai-alibaba-graph/spring-ai-alibaba-graph-core/src/test/java/com/alibaba/cloud/ai/graph/node/CodeActionTest.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,8 +13,9 @@
1313
* See the License for the specific language governing permissions and
1414
* limitations under the License.
1515
*/
16-
package com.alibaba.cloud.ai.graph;
16+
package com.alibaba.cloud.ai.graph.node;
1717

18+
import com.alibaba.cloud.ai.graph.OverAllState;
1819
import com.alibaba.cloud.ai.graph.action.NodeAction;
1920
import com.alibaba.cloud.ai.graph.node.code.CodeExecutorNodeAction;
2021
import com.alibaba.cloud.ai.graph.node.code.LocalCommandlineCodeExecutor;

0 commit comments

Comments
 (0)