diff --git a/src/main/java/com/redislabs/redisai/Command.java b/src/main/java/com/redislabs/redisai/Command.java index f12e3e7..9180b02 100644 --- a/src/main/java/com/redislabs/redisai/Command.java +++ b/src/main/java/com/redislabs/redisai/Command.java @@ -13,9 +13,11 @@ public enum Command implements ProtocolCommand { MODEL_RUN("AI.MODELRUN"), MODEL_EXECUTE("AI.MODELEXECUTE"), SCRIPT_SET("AI.SCRIPTSET"), + SCRIPT_STORE("AI.SCRIPTSTORE"), SCRIPT_GET("AI.SCRIPTGET"), SCRIPT_DEL("AI.SCRIPTDEL"), SCRIPT_RUN("AI.SCRIPTRUN"), + SCRIPT_EXECUTE("AI.SCRIPTEXECUTE"), DAGRUN("AI.DAGRUN"), DAGRUN_RO("AI.DAGRUN_RO"), DAGEXECUTE("AI.DAGEXECUTE"), diff --git a/src/main/java/com/redislabs/redisai/Dag.java b/src/main/java/com/redislabs/redisai/Dag.java index 5cca32a..4c41587 100644 --- a/src/main/java/com/redislabs/redisai/Dag.java +++ b/src/main/java/com/redislabs/redisai/Dag.java @@ -68,6 +68,21 @@ public Dag runScript(String key, String function, String[] inputs, String[] outp return this; } + @Override + public Dag executeScript( + String key, + String function, + List keys, + List inputs, + List args, + List outputs) { + List binary = + Script.scriptExecuteFlatArgs(key, function, keys, inputs, args, outputs, -1L, true); + this.commands.add(binary); + this.tensorgetflag.add(false); + return this; + } + List dagRunFlatArgs(String[] loadKeys, String[] persistKeys) { List args = new ArrayList<>(); if (loadKeys != null && loadKeys.length > 0) { diff --git a/src/main/java/com/redislabs/redisai/DagRunCommands.java b/src/main/java/com/redislabs/redisai/DagRunCommands.java index 328eed8..a4bfb0a 100644 --- a/src/main/java/com/redislabs/redisai/DagRunCommands.java +++ b/src/main/java/com/redislabs/redisai/DagRunCommands.java @@ -1,5 +1,7 @@ package com.redislabs.redisai; +import java.util.List; + interface DagRunCommands { T setTensor(String key, Tensor tensor); @@ -10,4 +12,12 @@ interface DagRunCommands { T executeModel(String key, String[] inputs, String[] outputs); T runScript(String key, String function, String[] inputs, String[] outputs); + + T executeScript( + String key, + String function, + List keys, + List inputs, + List args, + List outputs); } diff --git a/src/main/java/com/redislabs/redisai/Device.java b/src/main/java/com/redislabs/redisai/Device.java index aebb77c..2d3b895 100644 --- a/src/main/java/com/redislabs/redisai/Device.java +++ b/src/main/java/com/redislabs/redisai/Device.java @@ -10,9 +10,10 @@ public enum Device implements ProtocolCommand { private final byte[] raw; Device() { - raw = SafeEncoder.encode(this.name()); + raw = SafeEncoder.encode(name()); } + @Override public byte[] getRaw() { return raw; } diff --git a/src/main/java/com/redislabs/redisai/Keyword.java b/src/main/java/com/redislabs/redisai/Keyword.java index 25d69d8..bcd7b3f 100644 --- a/src/main/java/com/redislabs/redisai/Keyword.java +++ b/src/main/java/com/redislabs/redisai/Keyword.java @@ -12,6 +12,7 @@ public enum Keyword implements ProtocolCommand { SOURCE, RESETSTAT, TAG, + ENTRY_POINTS, BATCHSIZE, MINBATCHSIZE, MINBATCHTIMEOUT, @@ -21,6 +22,7 @@ public enum Keyword implements ProtocolCommand { LOAD, PERSIST, KEYS, + ARGS, PIPE("|>"); private final byte[] raw; diff --git a/src/main/java/com/redislabs/redisai/RedisAI.java b/src/main/java/com/redislabs/redisai/RedisAI.java index 95a46c0..2701b4e 100644 --- a/src/main/java/com/redislabs/redisai/RedisAI.java +++ b/src/main/java/com/redislabs/redisai/RedisAI.java @@ -9,6 +9,7 @@ import java.util.Map; import org.apache.commons.pool2.impl.GenericObjectPoolConfig; import redis.clients.jedis.BinaryClient; +import redis.clients.jedis.Client; import redis.clients.jedis.HostAndPort; import redis.clients.jedis.Jedis; import redis.clients.jedis.JedisClientConfig; @@ -103,6 +104,22 @@ private static JedisPoolConfig initPoolConfig(int poolSize) { return conf; } + private Jedis getConnection() { + return pool.getResource(); + } + + private BinaryClient sendCommand(Jedis conn, Command command, byte[]... args) { + BinaryClient client = conn.getClient(); + client.sendCommand(command, args); + return client; + } + + private Client sendCommand(Jedis conn, Command command, String... args) { + Client client = conn.getClient(); + client.sendCommand(command, args); + return client; + } + /** * Direct mapping to AI.TENSORSET * @@ -320,6 +337,27 @@ public boolean setScript(String key, Script script) { } } + /** + * Direct mapping to AI.MODELSTORE command. + * + *

{@code AI.SCRIPTSTORE [TAG tag] ENTRY_POINTS + * [...] SOURCE "