Zacharias
3c6a75a587
dont midn the version bump, it's easier to bump versions then to get Maven to bahave
558 lines
21 KiB
Java
558 lines
21 KiB
Java
package me.neurodock.display;
|
|
|
|
import com.github.dockerjava.api.DockerClient;
|
|
import com.github.dockerjava.api.model.HostConfig;
|
|
import com.github.dockerjava.core.DefaultDockerClientConfig;
|
|
import com.github.dockerjava.core.DockerClientImpl;
|
|
import com.github.dockerjava.httpclient5.ApacheDockerHttpClient;
|
|
import com.github.dockerjava.transport.DockerHttpClient;
|
|
import me.neurodock.core.Core;
|
|
import me.neurodock.core.Pair;
|
|
import me.neurodock.ollama.OllamaFunctionArgument;
|
|
import me.neurodock.ollama.OllamaFunctionTool;
|
|
import me.neurodock.ollama.OllamaPerameter;
|
|
import me.neurodock.ollama.OllamaToolResponse;
|
|
import me.neurodock.ollama.exceptions.OllamaToolErrorException;
|
|
import org.apache.commons.io.IOUtils;
|
|
import org.jetbrains.annotations.NotNull;
|
|
import org.json.JSONArray;
|
|
import org.json.JSONObject;
|
|
import org.jspecify.annotations.NonNull;
|
|
|
|
import java.io.*;
|
|
import java.net.ServerSocket;
|
|
import java.net.Socket;
|
|
import java.net.SocketException;
|
|
import java.nio.charset.Charset;
|
|
import java.nio.charset.StandardCharsets;
|
|
import java.security.MessageDigest;
|
|
import java.time.Duration;
|
|
import java.util.ArrayList;
|
|
import java.util.Arrays;
|
|
import java.util.List;
|
|
import java.util.concurrent.ConcurrentHashMap;
|
|
|
|
import static me.neurodock.core.Core.ensureDir;
|
|
import static me.neurodock.core.Core.writeLog;
|
|
|
|
/**
|
|
* A tool that runs python code.
|
|
* This is a wrapper around a docker container.
|
|
* This is partly meant as a proof of concept, but also as a way to run python code while keeping the executed code in a secure environment.
|
|
*/
|
|
public class PythonRunner extends OllamaFunctionTool {
|
|
/**
|
|
* The DockerClient instance.
|
|
*/
|
|
private DockerClient dockerClient;
|
|
/**
|
|
* The Core instance.
|
|
*/
|
|
private Core core;
|
|
|
|
/**
|
|
* The ServerSocket instance.
|
|
*/
|
|
private ServerSocket serverSocket;
|
|
|
|
/**
|
|
* A hash map of all python runners, this is to make it thread safe in a theoretical case of multiple python runners running at once
|
|
*/
|
|
private ConcurrentHashMap<String, StringBuilder> outputBuffers = new ConcurrentHashMap<>();
|
|
|
|
/**
|
|
* The docker client to be used throughout this tool
|
|
*/
|
|
DockerHttpClient dockerHttpClient;
|
|
|
|
/**
|
|
* Creates a new instance of PythonRunner.
|
|
* @param core The Core instance
|
|
*/
|
|
public PythonRunner(Core core) {
|
|
this.core = core;
|
|
|
|
ensureDir("./pythonFiles/");
|
|
|
|
try {
|
|
serverSocket = new ServerSocket(6050);
|
|
Thread thread = new Thread(() -> {
|
|
while (true) {
|
|
try {
|
|
Socket socket = serverSocket.accept();
|
|
|
|
BufferedReader in = new BufferedReader(new InputStreamReader(socket.getInputStream()));
|
|
|
|
String inputLine = in.readLine();
|
|
|
|
BufferedWriter out = new BufferedWriter(new OutputStreamWriter(socket.getOutputStream()));
|
|
|
|
try {
|
|
JSONObject data = new JSONObject(inputLine);
|
|
|
|
String toolName = data.optString("function", "");
|
|
|
|
if(toolName.equals("print_output")) {
|
|
// This is a "special" "tool" that is used to hijack the output of the python program so logs can be put else where since the LLM dose'nt need the entirety of the log from pacman and what not.
|
|
String containerId = data.optString("container_id", "");
|
|
if(containerId.isEmpty()) {
|
|
out.write(new JSONObject().put("error", "Missing container id").toString());
|
|
out.newLine();
|
|
out.flush();
|
|
out.close();
|
|
in.close();
|
|
socket.close();
|
|
continue;
|
|
}
|
|
|
|
outputBuffers.compute(containerId, (k, buf) -> {
|
|
if(buf == null) {
|
|
buf = new StringBuilder();
|
|
}
|
|
buf.append(data.optString("text", ""));
|
|
buf.append(System.lineSeparator());
|
|
return buf;
|
|
});
|
|
out.write(new JSONObject().put("container_id", containerId).toString());
|
|
out.newLine();
|
|
out.flush();
|
|
out.close();
|
|
in.close();
|
|
socket.close();
|
|
continue;
|
|
}
|
|
|
|
List<Pair<OllamaFunctionTool, String>> list = core.getFuntionTools().stream().filter(funtionTool -> funtionTool.getKey().name().equalsIgnoreCase(toolName)).toList();
|
|
|
|
if (list.isEmpty()) {
|
|
out.write(new JSONObject().put("error", "Function don't exist").toString());
|
|
out.newLine();
|
|
out.flush();
|
|
out.close();
|
|
in.close();
|
|
socket.close();
|
|
continue;
|
|
}
|
|
|
|
ArrayList<OllamaFunctionArgument> args = new ArrayList<>();
|
|
|
|
for (Object o : data.optJSONArray("arguments", new JSONArray())) {
|
|
if (o instanceof JSONObject obj) {
|
|
if(obj.has("value") && !obj.isNull("value")) {
|
|
OllamaFunctionArgument arg = new OllamaFunctionArgument(obj.getString("name"), obj.getString("value"));
|
|
args.add(arg);
|
|
}
|
|
}
|
|
}
|
|
|
|
out.write(list.getFirst().getKey().function(args.toArray(new OllamaFunctionArgument[0])).getResponse());
|
|
out.newLine();
|
|
out.flush();
|
|
out.close();
|
|
in.close();
|
|
socket.close();
|
|
} catch (Exception e) {
|
|
}
|
|
} catch (Exception e) {
|
|
|
|
}
|
|
}
|
|
});
|
|
|
|
thread.start();
|
|
}catch (Exception e) {
|
|
e.printStackTrace();
|
|
}
|
|
|
|
DefaultDockerClientConfig config
|
|
= DefaultDockerClientConfig.createDefaultConfigBuilder()
|
|
.build();
|
|
|
|
dockerHttpClient = new ApacheDockerHttpClient.Builder()
|
|
.dockerHost(config.getDockerHost())
|
|
.sslConfig(config.getSSLConfig())
|
|
.maxConnections(10)
|
|
.connectionTimeout(Duration.ofSeconds(100))
|
|
.responseTimeout(Duration.ofSeconds(100))
|
|
.build();
|
|
|
|
DockerHttpClient.Request ping = DockerHttpClient.Request.builder()
|
|
.method(DockerHttpClient.Request.Method.GET)
|
|
.path("/_ping")
|
|
.build();
|
|
|
|
try(DockerHttpClient.Response response = dockerHttpClient.execute(ping))
|
|
{
|
|
if(!(response.getStatusCode() == 200))
|
|
{
|
|
writeLog("Failed to ping docker: " + response.getStatusCode());
|
|
System.out.println("Failed to ping docker (" + response.getStatusCode() + "). Docker components is disabled.");
|
|
dockerClient = null;
|
|
return;
|
|
}
|
|
if(!(IOUtils.toString(response.getBody(), Charset.defaultCharset()).equals("OK")))
|
|
{
|
|
writeLog("Failed to ping docker: body not OK");
|
|
System.out.println("Failed to ping docker. Docker components is disabled.");
|
|
dockerClient = null;
|
|
return;
|
|
}
|
|
}
|
|
catch (RuntimeException re)
|
|
{
|
|
if(re.getCause() instanceof SocketException) {
|
|
writeLog("Failed to ping docker: " + re.getMessage());
|
|
System.out.println("Failed to ping docker. Docker component is disabled.");
|
|
dockerClient = null;
|
|
return;
|
|
}
|
|
else{
|
|
writeLog("Failed to ping docker: " + re.getMessage());
|
|
re.printStackTrace();
|
|
System.out.println("Failed to ping docker. Docker component is disabled.");
|
|
dockerClient = null;
|
|
return;
|
|
}
|
|
}
|
|
catch (Exception e)
|
|
{
|
|
writeLog("Failed to ping docker: " + e.getMessage());
|
|
e.printStackTrace();
|
|
System.out.println("Failed to ping docker. Docker component is disabled.");
|
|
dockerClient = null;
|
|
return;
|
|
}
|
|
|
|
dockerClient = DockerClientImpl.getInstance(config, dockerHttpClient);
|
|
|
|
initializingImage();
|
|
}
|
|
|
|
/**
|
|
* A small function to pull the archlinux:latest image from Docker Hub. See <a href="https://hub.docker.com/_/archlinux/">ArchLinux Docker Hub</a>
|
|
* <br>
|
|
* I want to note here that the reason for using an ArchLinux image and not somthing like python or Ubuntu. Is because I'm more
|
|
* comfortable with how to do this on an Arch host, and I could not manage to pull libraries for the python image. So those who despise Arch you're welcome to re-write the arch specific parts.
|
|
*/
|
|
private void initializingImage() {
|
|
// Check if archlinux:latest is pulled already
|
|
DockerHttpClient.Request checkArchlinux_latest = DockerHttpClient.Request.builder()
|
|
.method(DockerHttpClient.Request.Method.GET)
|
|
.path("/images/archlinux:latest/json")
|
|
.build();
|
|
|
|
boolean archlinux_latestAvailable = false;
|
|
|
|
try(DockerHttpClient.Response response = dockerHttpClient.execute(checkArchlinux_latest))
|
|
{
|
|
if(response.getStatusCode() == 200)
|
|
{
|
|
archlinux_latestAvailable = true;
|
|
}
|
|
}
|
|
catch (Exception e)
|
|
{
|
|
writeLog("Failed to fetch archlinux:latest: " + e.getMessage());
|
|
e.printStackTrace();
|
|
System.out.println("Failed to fetch archlinux:latest.");
|
|
return;
|
|
}
|
|
|
|
if(!archlinux_latestAvailable) {
|
|
|
|
// If it's not already downloaded, we will download it.
|
|
DockerHttpClient.Request downloadArchlinux_latest = DockerHttpClient.Request.builder()
|
|
.method(DockerHttpClient.Request.Method.POST)
|
|
.path("/images/create?fromImage=archlinux&tag=latest")
|
|
.build();
|
|
|
|
try (DockerHttpClient.Response response = dockerHttpClient.execute(downloadArchlinux_latest)) {
|
|
if (response.getStatusCode() == 200) {
|
|
writeLog("Successfully downloaded archlinux:latest");
|
|
}
|
|
else {
|
|
writeLog("Failed to download archlinux:latest: " + response.getStatusCode() + "\n" + response.getBody());
|
|
}
|
|
} catch (Exception e) {
|
|
writeLog("Failed to fetch archlinux:latest: " + e.getMessage());
|
|
e.printStackTrace();
|
|
System.out.println("Failed to fetch archlinux:latest.");
|
|
return;
|
|
}
|
|
}
|
|
else
|
|
{
|
|
writeLog("archlinux:latest already exists");
|
|
}
|
|
}
|
|
|
|
@Override
|
|
public @NonNull String name() {
|
|
return "python_runner";
|
|
}
|
|
|
|
@Override
|
|
public String description() {
|
|
return "Runs python code.\nuse python's print method for output";
|
|
}
|
|
|
|
@Override
|
|
public @NotNull OllamaPerameter parameters() {
|
|
return OllamaPerameter.builder()
|
|
.addProperty("code", OllamaPerameter.OllamaPerameterBuilder.Type.STRING, "The code to be executed, optional if a file with the same name exists.")
|
|
.addProperty("name", OllamaPerameter.OllamaPerameterBuilder.Type.STRING, "The name of the python code")
|
|
.addProperty("libs", OllamaPerameter.OllamaPerameterBuilder.Type.STRING, "A space separated list of pip packages needed")
|
|
.build();
|
|
}
|
|
|
|
@Override
|
|
public @NonNull OllamaToolResponse function(OllamaFunctionArgument... args) {
|
|
if(dockerClient == null)
|
|
{
|
|
return new OllamaToolResponse(name(), "Docker is disabled");
|
|
}
|
|
if(args.length == 0)
|
|
{
|
|
throw new OllamaToolErrorException(name(), "Missing code argument");
|
|
}
|
|
|
|
String name = null;
|
|
String code = null;
|
|
String libs = null;
|
|
|
|
for(OllamaFunctionArgument arg : args)
|
|
{
|
|
if(arg.argument().equals("name") && !arg.value().equals(""))
|
|
{
|
|
name = (String) arg.value();
|
|
if(!name.endsWith(".py"))
|
|
{
|
|
name += ".py";
|
|
}
|
|
} else if (arg.argument().equals("code") && !arg.value().equals("")) {
|
|
code = (String) arg.value();
|
|
}
|
|
else if(arg.argument().equals("libs") && !arg.value().equals(""))
|
|
{
|
|
libs = (String) arg.value();
|
|
}
|
|
}
|
|
|
|
if(name == null && code == null)
|
|
{
|
|
throw new OllamaToolErrorException(name(), "Missing code or name");
|
|
}
|
|
|
|
if(name == null)
|
|
{
|
|
try {
|
|
MessageDigest digest = MessageDigest.getInstance("SHA-256");
|
|
byte[] encodedhash = digest.digest(String.valueOf(args[0].value()).getBytes(StandardCharsets.UTF_8));
|
|
StringBuffer hexString = new StringBuffer();
|
|
for(byte b : encodedhash)
|
|
{
|
|
hexString.append(String.format("%02x", b));
|
|
}
|
|
name = hexString.toString()+".py";
|
|
}catch (Exception e) {}
|
|
}
|
|
|
|
name = name.replace(" ", "_");
|
|
|
|
writeLog("Running python code `" + name + "`");
|
|
|
|
File pythonFile = new File("./pythonFiles", name);
|
|
|
|
code = "from external_tools import *\n\n"+code;
|
|
|
|
if(!pythonFile.exists())
|
|
{
|
|
try {
|
|
BufferedWriter writer = new BufferedWriter(new FileWriter(pythonFile));
|
|
writer.write(code);
|
|
writer.close();
|
|
}catch(IOException e) {}
|
|
}
|
|
|
|
File f = new File("./pythonFiles", "external_tools.py");
|
|
|
|
String containerId = dockerClient.createContainerCmd("archlinux")
|
|
.withCmd("/bin/bash","/cmd.sh")
|
|
.withHostConfig(HostConfig.newHostConfig()
|
|
.withExtraHosts("host.docker.internal:host-gateway"))
|
|
.exec().getId();
|
|
|
|
outputBuffers.put(containerId, new StringBuilder());
|
|
|
|
try {
|
|
String external_tools = generateExternalTools(containerId);
|
|
BufferedWriter bw = new BufferedWriter(new FileWriter(f));
|
|
bw.write(external_tools/*.replace("\\n", "\n").replace("\\!n", "\\n")*/);
|
|
bw.flush();
|
|
bw.close();
|
|
}catch(IOException e) {}
|
|
|
|
try {
|
|
ArrayList<String> pythonArgs = new ArrayList<>();
|
|
|
|
pythonArgs.add("#!/bin/env bash\n");
|
|
|
|
StringBuilder cmd = new StringBuilder();
|
|
cmd.append("""
|
|
set -e
|
|
|
|
pacman --noconfirm -Sy > /dev/null
|
|
pacman --noconfirm -S python python-pip > /dev/null
|
|
mkdir pythonRun > /dev/null
|
|
python -m venv ./pythonRun > /dev/null
|
|
cp /external_tools.py ./pythonRun > /dev/null
|
|
cp\s/""")
|
|
.append(name).append("""
|
|
\s./pythonRun > /dev/null
|
|
cd pythonRun > /dev/null
|
|
source ./bin/activate > /dev/null
|
|
""");
|
|
|
|
if(libs != null && !libs.isEmpty())
|
|
{
|
|
cmd.append("pip install ");
|
|
for(String lib : libs.split(" "))
|
|
{
|
|
cmd.append(lib).append(" ");
|
|
}
|
|
cmd.append(" > /dev/null\n");
|
|
}
|
|
|
|
cmd.append("python ").append(name);//.append(" exit");
|
|
|
|
pythonArgs.add(cmd.toString());
|
|
|
|
StringBuilder fullCmd = new StringBuilder();
|
|
|
|
for(String arg : pythonArgs)
|
|
{
|
|
fullCmd.append(arg).append(" ");
|
|
}
|
|
|
|
File program = new File(Core.CACHE_DIRECTORY, "cmd.sh");
|
|
if(program.exists()){
|
|
program.delete();
|
|
}
|
|
program.createNewFile();
|
|
|
|
BufferedWriter bw = new BufferedWriter(new FileWriter(program));
|
|
bw.write(fullCmd.toString());
|
|
bw.flush();
|
|
bw.close();
|
|
|
|
dockerClient.copyArchiveToContainerCmd(containerId)
|
|
.withHostResource(pythonFile.getPath())
|
|
.exec();
|
|
|
|
dockerClient.copyArchiveToContainerCmd(containerId)
|
|
.withHostResource(f.getPath())
|
|
.exec();
|
|
|
|
dockerClient.copyArchiveToContainerCmd(containerId)
|
|
.withHostResource(program.getPath())
|
|
.exec();
|
|
|
|
dockerClient.startContainerCmd(containerId).exec();
|
|
|
|
dockerClient.waitContainerCmd(containerId).start().awaitCompletion();
|
|
|
|
String output = outputBuffers.remove(containerId).toString();
|
|
|
|
dockerClient.removeContainerCmd(containerId).exec();
|
|
|
|
return new OllamaToolResponse(name(), output.isEmpty() ? "Python script names "+name+" ran successfully with no output" : "Python script names "+name+" output: "+output);
|
|
}
|
|
catch (Exception e) {
|
|
throw new OllamaToolErrorException(name(), "Docker unavailable");
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Generates the external_tools.py file.<br>
|
|
* This is meant to provide the python code with all ExternalTools defined in the OllamaObject.
|
|
* @return The generated external_tools.py file
|
|
* @throws IllegalStateException This is thrown when the external_tool_base.py can not be read, this should be handled by disabling the python runner.
|
|
*/
|
|
private String generateExternalTools(String containerId) throws IllegalStateException{
|
|
StringBuilder code = new StringBuilder();
|
|
|
|
code.append("CONTAINER_ID = \"").append(containerId).append("\"\n\n");
|
|
|
|
try (InputStream in = getClass().getResourceAsStream("/external_tool_base.py");
|
|
BufferedReader reader = new BufferedReader(new InputStreamReader(in))) {
|
|
String tmp = null;
|
|
while((tmp = reader.readLine()) != null)
|
|
{
|
|
code.append(tmp).append('\n');
|
|
}
|
|
}catch (Exception ex)
|
|
{
|
|
throw new IllegalStateException("Can not read the base external tools file, this is a fetal error for the python runner!");
|
|
}
|
|
|
|
for(Pair<OllamaFunctionTool, String> funtionTool : core.getFuntionTools())
|
|
{
|
|
OllamaFunctionTool tool = funtionTool.getKey();
|
|
|
|
String name = tool.name();
|
|
|
|
code.append("def ").append(name).append("(");
|
|
|
|
ArrayList<String> args = new ArrayList<>();
|
|
|
|
boolean first = true;
|
|
try {
|
|
for(String argName : tool.parameters().getRequired())
|
|
{
|
|
args.add(argName);
|
|
if (!first) {
|
|
code.append(", ");
|
|
}
|
|
code.append(argName);
|
|
first = false;
|
|
}
|
|
|
|
for (String argName : tool.parameters().getProperties().keySet()) {
|
|
if(Arrays.asList(tool.parameters().getRequired()).contains(argName)) continue;
|
|
|
|
args.add(argName);
|
|
if (!first) {
|
|
code.append(", ");
|
|
}
|
|
code.append(argName).append("=None");
|
|
first = false;
|
|
}
|
|
}catch (Exception e) {}
|
|
code.append("):\n");
|
|
code.append(" data = {\"function\":\"").append(tool.name()).append("\"");
|
|
if(args.size() > 0)
|
|
{
|
|
code.append(",\"arguments\":[");
|
|
}
|
|
first = true;
|
|
for(String str : args)
|
|
{
|
|
if(!first)
|
|
{
|
|
code.append(", ");
|
|
}
|
|
code.append("{\"name\": \"").append(str).append("\", \"value\": ").append(str).append("}");
|
|
first = false;
|
|
}
|
|
if(args.size() > 0)
|
|
{
|
|
code.append("]");
|
|
}
|
|
code.append("}\n");
|
|
code.append(" return json.loads(connect(json.dumps(data)))\n\n");
|
|
}
|
|
|
|
return code.toString();
|
|
}
|
|
}
|