Zacharias
96a6ed169e
build / build (push) Has been cancelled
Fixed URL building in APITool.java. Updated libraries for API, Launcher. Started applying NotNull annotation to methods and arguments in Core.java. Started handeling the big init block in Core.java. Added overloaded constructor for Core.java and made Core.ollamaIP a variabled that's definible Updated Tool.addTool return for 400 error.
488 lines
18 KiB
Java
488 lines
18 KiB
Java
package me.zacharias.chat.display;
|
|
|
|
import com.github.dockerjava.api.DockerClient;
|
|
import com.github.dockerjava.api.async.ResultCallback;
|
|
import com.github.dockerjava.api.async.ResultCallbackTemplate;
|
|
import com.github.dockerjava.api.command.*;
|
|
import com.github.dockerjava.api.model.BuildResponseItem;
|
|
import com.github.dockerjava.api.model.Frame;
|
|
import com.github.dockerjava.api.model.Statistics;
|
|
import com.github.dockerjava.core.DefaultDockerClientConfig;
|
|
import com.github.dockerjava.httpclient5.ApacheDockerHttpClient;
|
|
import com.github.dockerjava.transport.DockerHttpClient;
|
|
import me.zacharias.chat.core.Core;
|
|
import me.zacharias.chat.core.Pair;
|
|
import me.zacharias.chat.ollama.*;
|
|
import me.zacharias.chat.ollama.exceptions.OllamaToolErrorException;
|
|
import org.apache.commons.io.IOUtils;
|
|
import org.json.JSONArray;
|
|
import org.json.JSONObject;
|
|
|
|
import java.io.*;
|
|
import java.net.ServerSocket;
|
|
import java.net.Socket;
|
|
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.logging.Logger;
|
|
|
|
import static me.zacharias.chat.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;
|
|
|
|
/**
|
|
* Creates a new instance of PythonRunner.
|
|
* @param core The Core instance
|
|
*/
|
|
public PythonRunner(Core core) {
|
|
this.core = core;
|
|
|
|
generateExternalTools();
|
|
|
|
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);
|
|
List<Pair<OllamaFunctionTool, String>> list = core.getFuntionTools().stream().filter(funtionTool -> funtionTool.getKey().name().equalsIgnoreCase(data.optString("function", ""))).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()
|
|
.withDockerHost("tcp://localhost:2375")
|
|
.withDockerTlsVerify(false)
|
|
.withDockerCertPath("~/.docker")
|
|
.build();
|
|
DockerHttpClient dockerHttpClient = new ApacheDockerHttpClient.Builder()
|
|
.dockerHost(config.getDockerHost())
|
|
.maxConnections(10)
|
|
.connectionTimeout(Duration.ofSeconds(100))
|
|
.responseTimeout(Duration.ofSeconds(100))
|
|
.sslConfig(config.getSSLConfig())
|
|
.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");
|
|
System.out.println("Failed to ping docker. Docker components is disabled.");
|
|
dockerClient = null;
|
|
return;
|
|
}
|
|
if(!(IOUtils.toString(response.getBody(), Charset.defaultCharset()).equals("OK")))
|
|
{
|
|
writeLog("Failed to ping docker");
|
|
System.out.println("Failed to ping docker. Docker components is disabled.");
|
|
dockerClient = null;
|
|
return;
|
|
}
|
|
}
|
|
catch (Exception e)
|
|
{
|
|
writeLog("Failed to ping docker");
|
|
System.out.println("Failed to ping docker. Docker components is disabled.");
|
|
dockerClient = null;
|
|
}
|
|
}
|
|
|
|
@Override
|
|
public String name() {
|
|
return "python_runner";
|
|
}
|
|
|
|
@Override
|
|
public String description() {
|
|
return "Runs python code";
|
|
}
|
|
|
|
@Override
|
|
public OllamaPerameter parameters() {
|
|
return OllamaPerameter.builder()
|
|
.addProperty("code", OllamaPerameter.OllamaPerameterBuilder.Type.STRING, "The code to be executed", true)
|
|
.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 OllamaToolRespnce function(OllamaFunctionArgument... args) {
|
|
if(dockerClient == null)
|
|
{
|
|
return new OllamaToolRespnce(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)
|
|
{
|
|
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");
|
|
|
|
try {
|
|
String external_tools = generateExternalTools();
|
|
BufferedWriter bw = new BufferedWriter(new FileWriter(f));
|
|
bw.write(external_tools.replace("\\n", "\n"));
|
|
bw.flush();
|
|
bw.close();
|
|
}catch(IOException e) {}
|
|
|
|
try {
|
|
ArrayList<String> pythonArgs = new ArrayList<>();
|
|
|
|
pythonArgs.add("/bin/bash");
|
|
|
|
StringBuilder cmd = new StringBuilder();
|
|
cmd.append("set -e\n\npacman --noconfirm -Sy > /dev/null\npacman --noconfirm -S python python-pip > /dev/null\nmkdir pythonRun > /dev/null\npython -m venv ./pythonRun > /dev/null\ncp external_tools.py ./pythonRun > /dev/null\ncp ").append(name).append(" ./pythonRun > /dev/null\ncd pythonRun > /dev/null\nsource ./bin/activate > /dev/null\n");
|
|
|
|
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("-c");
|
|
pythonArgs.add(cmd.toString());
|
|
|
|
StringBuilder fullCmd = new StringBuilder();
|
|
|
|
for(String arg : pythonArgs)
|
|
{
|
|
fullCmd.append(arg).append(" ");
|
|
}
|
|
|
|
File program = new File("./cache", "cmd.sh");
|
|
if(program.exists()){
|
|
program.delete();
|
|
}
|
|
program.createNewFile();
|
|
|
|
BufferedWriter bw = new BufferedWriter(new FileWriter(program));
|
|
bw.write(cmd.toString());
|
|
bw.flush();
|
|
bw.close();
|
|
|
|
|
|
String containerId = dockerClient.createContainerCmd("archlinux").withCmd("/bin/bash","./cmd.sh").exec().getId();
|
|
|
|
dockerClient.copyArchiveToContainerCmd(containerId)
|
|
.withHostResource(pythonFile.getPath())
|
|
.exec();
|
|
|
|
dockerClient.copyArchiveToContainerCmd(containerId)
|
|
.withHostResource(f.getPath())
|
|
.exec();
|
|
|
|
dockerClient.copyArchiveToContainerCmd(containerId)
|
|
.withHostResource(program.getPath())
|
|
.exec();
|
|
|
|
//InputStream stdin = new ByteArrayInputStream(fullCmd.toString().getBytes(StandardCharsets.UTF_8));
|
|
|
|
dockerClient.startContainerCmd(containerId).exec();
|
|
|
|
//dockerClient.attachContainerCmd(containerId).withStdIn(stdin).exec(null);
|
|
|
|
//dockerClient.execCreateCmd(containerId).withCmd(fullCmd.toString()).exec();
|
|
|
|
GetContainerLog log = new GetContainerLog(dockerClient, containerId);
|
|
|
|
List<String> logs = new ArrayList<>();
|
|
|
|
do {
|
|
try {
|
|
Thread.sleep(2000);
|
|
} catch (InterruptedException e) {
|
|
throw new RuntimeException(e);
|
|
}
|
|
logs.addAll(log.getDockerLogs());
|
|
}
|
|
while (isRunning(containerId));
|
|
|
|
StringBuilder output = new StringBuilder();
|
|
|
|
for (String s : logs) {
|
|
output.append(s).append("\n");
|
|
}
|
|
|
|
//writeLog("Result from python: " + output.toString());
|
|
|
|
return new OllamaToolRespnce(name(), output.toString());
|
|
}
|
|
catch (Exception e) {
|
|
throw new OllamaToolErrorException(name(), "Docker unavalible");
|
|
}
|
|
}
|
|
|
|
/**
|
|
* 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() throws IllegalStateException{
|
|
StringBuilder code = new StringBuilder();
|
|
|
|
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().getProperties().keySet()) {
|
|
args.add(argName);
|
|
if (!first) {
|
|
code.append(", ");
|
|
}
|
|
code.append(argName);
|
|
if (Arrays.stream(tool.parameters().getRequired()).noneMatch(required -> required.equals(argName))) {
|
|
code.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();
|
|
}
|
|
|
|
/**
|
|
* Checks if a Docker Container is running
|
|
* @param containerId The ID of the Container
|
|
* @return a boolean weather it's running or not
|
|
*/
|
|
public boolean isRunning(String containerId)
|
|
{
|
|
InspectContainerResponse cmd = dockerClient.inspectContainerCmd(containerId).exec();
|
|
return Boolean.TRUE.equals(cmd.getState().getRunning());
|
|
}
|
|
|
|
/**
|
|
* A Helper class to get the logs from a docker container.
|
|
*/
|
|
public class GetContainerLog {
|
|
private DockerClient dockerClient;
|
|
private String containerId;
|
|
private int lastLogTime;
|
|
|
|
private static String nameOfLogger = "dockertest.PrintContainerLog";
|
|
private static Logger myLogger = Logger.getLogger(nameOfLogger);
|
|
|
|
/**
|
|
* Creates a new instance of {@link GetContainerLog}
|
|
* @param dockerClient The DockerClient instance
|
|
* @param containerId The container id
|
|
*/
|
|
public GetContainerLog(DockerClient dockerClient, String containerId) {
|
|
this.dockerClient = dockerClient;
|
|
this.containerId = containerId;
|
|
this.lastLogTime = (int) (System.currentTimeMillis() / 1000);
|
|
}
|
|
|
|
/**
|
|
* Gets the logs of the container.
|
|
* @return The logs of the container
|
|
*/
|
|
public List<String> getDockerLogs() {
|
|
|
|
final List<String> logs = new ArrayList<>();
|
|
|
|
LogContainerCmd logContainerCmd = dockerClient.logContainerCmd(containerId);
|
|
logContainerCmd.withStdOut(true).withStdErr(true);
|
|
logContainerCmd.withSince(lastLogTime); // UNIX timestamp (integer) to filter logs. Specifying a timestamp will only output log-entries since that timestamp.
|
|
// logContainerCmd.withTail(4); // get only the last 4 log entries
|
|
|
|
logContainerCmd.withTimestamps(true);
|
|
|
|
try {
|
|
logContainerCmd.exec(new ResultCallback.Adapter<Frame>() {
|
|
@Override
|
|
public void onNext(Frame item) {
|
|
logs.add(new String(item.getPayload()).trim());
|
|
}
|
|
}).awaitCompletion();
|
|
} catch (InterruptedException e) {
|
|
myLogger.severe("Interrupted Exception!" + e.getMessage());
|
|
}
|
|
|
|
lastLogTime = (int) (System.currentTimeMillis() / 1000) + 5; // assumes at least a 5 second wait between calls to getDockerLogs
|
|
|
|
return logs;
|
|
}
|
|
}
|
|
}
|