Move TransformContext to a dedicated module (LangStream#519)

vectorize-io · Oct 3, 2023 · e3953fe · e3953fe
1 parent 41fc767
commit e3953fe
Show file tree

Hide file tree

Showing 48 changed files with 226 additions and 122 deletions.
diff --git a/langstream-agents/langstream-agents-commons/pom.xml b/langstream-agents/langstream-agents-commons/pom.xml
@@ -0,0 +1,79 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+
+    Copyright DataStax, Inc.
+
+    Licensed under the Apache License, Version 2.0 (the "License");
+    you may not use this file except in compliance with the License.
+    You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+    Unless required by applicable law or agreed to in writing, software
+    distributed under the License is distributed on an "AS IS" BASIS,
+    WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+    See the License for the specific language governing permissions and
+    limitations under the License.
+
+-->
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/xsd/maven-4.0.0.xsd">
+  <parent>
+    <groupId>ai.langstream</groupId>
+    <artifactId>langstream-agents</artifactId>
+    <version>0.0.23-SNAPSHOT</version>
+  </parent>
+  <modelVersion>4.0.0</modelVersion>
+  <artifactId>langstream-agents-commons</artifactId>
+  <packaging>jar</packaging>
+  <name>LangStream - Common Agent utilities</name>
+  <properties>
+    <build.dir>${project.build.directory}</build.dir>
+  </properties>
+  <dependencies>
+    <dependency>
+      <groupId>${project.groupId}</groupId>
+      <artifactId>langstream-api</artifactId>
+      <version>${project.version}</version>
+      <scope>provided</scope>
+    </dependency>
+    <dependency>
+      <groupId>com.networknt</groupId>
+      <artifactId>json-schema-validator</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>com.fasterxml.jackson.dataformat</groupId>
+      <artifactId>jackson-dataformat-yaml</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.avro</groupId>
+      <artifactId>avro</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.tomcat.embed</groupId>
+      <artifactId>tomcat-embed-el</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>org.apache.commons</groupId>
+      <artifactId>commons-collections4</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>com.fasterxml.jackson.core</groupId>
+      <artifactId>jackson-databind</artifactId>
+    </dependency>
+    <dependency>
+      <!-- Transform Function -->
+      <groupId>org.apache.pulsar</groupId>
+      <artifactId>pulsar-client-original</artifactId>
+      <exclusions>
+        <exclusion>
+          <groupId>io.netty</groupId>
+          <artifactId>*</artifactId>
+        </exclusion>
+        <exclusion>
+          <groupId>io.netty.incubator</groupId>
+          <artifactId>*</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
+  </dependencies>
+</project>
diff --git a/...astax/oss/streaming/ai/util/AvroUtil.java → ...angstream/ai/agents/commons/AvroUtil.java b/...astax/oss/streaming/ai/util/AvroUtil.java → ...angstream/ai/agents/commons/AvroUtil.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.datastax.oss.streaming.ai.util;
+package ai.langstream.ai.agents.commons;
 
 import java.util.ArrayList;
 import java.util.Collection;

diff --git a/.../oss/streaming/ai/util/JsonConverter.java → ...ream/ai/agents/commons/JsonConverter.java b/.../oss/streaming/ai/util/JsonConverter.java → ...ream/ai/agents/commons/JsonConverter.java
@@ -13,9 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.datastax.oss.streaming.ai.util;
-
-import static com.datastax.oss.streaming.ai.util.TransformFunctionUtil.getBytes;
+package ai.langstream.ai.agents.commons;
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.node.ArrayNode;
@@ -247,4 +245,19 @@ JsonNode toJson(Schema schema, Object value) {
                     }
                 });
     }
+
+    public static byte[] getBytes(ByteBuffer byteBuffer) {
+        if (byteBuffer == null) {
+            return null;
+        }
+        if (byteBuffer.hasArray()
+                && byteBuffer.arrayOffset() == 0
+                && byteBuffer.array().length == byteBuffer.remaining()) {
+            return byteBuffer.array();
+        }
+        // Direct buffer is not backed by array and it needs to be read from direct memory
+        byte[] array = new byte[byteBuffer.remaining()];
+        byteBuffer.get(array);
+        return array;
+    }
 }
diff --git a/...ax/oss/streaming/ai/model/JsonRecord.java → ...gstream/ai/agents/commons/JsonRecord.java b/...ax/oss/streaming/ai/model/JsonRecord.java → ...gstream/ai/agents/commons/JsonRecord.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.datastax.oss.streaming.ai.model;
+package ai.langstream.ai.agents.commons;
 
 import java.util.Map;
 import lombok.Data;

diff --git a/...ax/oss/streaming/ai/TransformContext.java → ...m/ai/agents/commons/TransformContext.java b/...ax/oss/streaming/ai/TransformContext.java → ...m/ai/agents/commons/TransformContext.java
@@ -13,28 +13,28 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.datastax.oss.streaming.ai;
+package ai.langstream.ai.agents.commons;
 
-import static com.datastax.oss.streaming.ai.util.TransformFunctionUtil.safeClone;
-
-import com.datastax.oss.streaming.ai.model.JsonRecord;
-import com.datastax.oss.streaming.ai.model.TransformSchemaType;
-import com.datastax.oss.streaming.ai.util.AvroUtil;
-import com.datastax.oss.streaming.ai.util.JsonConverter;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.core.type.TypeReference;
+import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import com.fasterxml.jackson.databind.node.ObjectNode;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
 import java.util.Map;
+import java.util.Set;
 import lombok.Data;
 import lombok.SneakyThrows;
 import lombok.extern.slf4j.Slf4j;
 import org.apache.avro.Conversions;
 import org.apache.avro.Schema;
+import org.apache.avro.generic.GenericData;
 import org.apache.avro.generic.GenericDatumWriter;
 import org.apache.avro.generic.GenericRecord;
 import org.apache.avro.io.BinaryEncoder;
@@ -324,4 +324,38 @@ public void setResultField(
             }
         }
     }
+
+    public static Object safeClone(Object object) {
+        if (object == null) {
+            return null;
+        }
+        if (object.getClass().isPrimitive()
+                || object instanceof String
+                || object instanceof Number
+                || object instanceof Boolean) {
+            return object;
+        }
+        if (object instanceof Map map) {
+            HashMap<Object, Object> res = new HashMap<>();
+            map.forEach((k, v) -> res.put(safeClone(k), safeClone(v)));
+            return res;
+        }
+        if (object instanceof List list) {
+            List<Object> res = new ArrayList<>();
+            list.forEach(v -> res.add(safeClone(v)));
+            return res;
+        }
+        if (object instanceof Set set) {
+            Set<Object> res = new HashSet<>();
+            set.forEach(v -> res.add(safeClone(v)));
+            return res;
+        }
+        if (object instanceof GenericRecord genericRecord) {
+            return GenericData.get().deepCopy(genericRecord.getSchema(), genericRecord);
+        }
+        if (object instanceof JsonNode jsonNode) {
+            return jsonNode.deepCopy();
+        }
+        throw new UnsupportedOperationException("Cannot copy a value of " + object.getClass());
+    }
 }
diff --git a/...reaming/ai/model/TransformSchemaType.java → ...i/agents/commons/TransformSchemaType.java b/...reaming/ai/model/TransformSchemaType.java → ...i/agents/commons/TransformSchemaType.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.datastax.oss.streaming.ai.model;
+package ai.langstream.ai.agents.commons;
 
 public enum TransformSchemaType {
     STRING,

diff --git a/.../jstl/DisabledInvocationBeanResolver.java → .../jstl/DisabledInvocationBeanResolver.java b/.../jstl/DisabledInvocationBeanResolver.java → .../jstl/DisabledInvocationBeanResolver.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.datastax.oss.streaming.ai.jstl;
+package ai.langstream.ai.agents.commons.jstl;
 
 import jakarta.el.BeanELResolver;
 import jakarta.el.ELContext;

diff --git a/.../oss/streaming/ai/jstl/JstlEvaluator.java → ...ai/agents/commons/jstl/JstlEvaluator.java b/.../oss/streaming/ai/jstl/JstlEvaluator.java → ...ai/agents/commons/jstl/JstlEvaluator.java
@@ -13,9 +13,9 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.datastax.oss.streaming.ai.jstl;
+package ai.langstream.ai.agents.commons.jstl;
 
-import com.datastax.oss.streaming.ai.TransformContext;
+import ai.langstream.ai.agents.commons.TransformContext;
 import jakarta.el.ELContext;
 import jakarta.el.ExpressionFactory;
 import jakarta.el.ValueExpression;

diff --git a/.../oss/streaming/ai/jstl/JstlFunctions.java → ...ai/agents/commons/jstl/JstlFunctions.java b/.../oss/streaming/ai/jstl/JstlFunctions.java → ...ai/agents/commons/jstl/JstlFunctions.java
@@ -13,10 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.datastax.oss.streaming.ai.jstl;
+package ai.langstream.ai.agents.commons.jstl;
 
-import com.datastax.oss.streaming.ai.TransformContext;
-import com.datastax.oss.streaming.ai.jstl.predicate.JstlPredicate;
+import ai.langstream.ai.agents.commons.TransformContext;
+import ai.langstream.ai.agents.commons.jstl.predicate.JstlPredicate;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import jakarta.el.ELException;
 import java.lang.reflect.Array;

diff --git a/.../ai/jstl/JstlTransformContextAdapter.java → ...ons/jstl/JstlTransformContextAdapter.java b/.../ai/jstl/JstlTransformContextAdapter.java → ...ons/jstl/JstlTransformContextAdapter.java
@@ -13,10 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.datastax.oss.streaming.ai.jstl;
+package ai.langstream.ai.agents.commons.jstl;
 
-import com.datastax.oss.streaming.ai.TransformContext;
-import com.datastax.oss.streaming.ai.util.AvroUtil;
+import ai.langstream.ai.agents.commons.AvroUtil;
+import ai.langstream.ai.agents.commons.TransformContext;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;

diff --git a/.../streaming/ai/jstl/JstlTypeConverter.java → ...gents/commons/jstl/JstlTypeConverter.java b/.../streaming/ai/jstl/JstlTypeConverter.java → ...gents/commons/jstl/JstlTypeConverter.java
@@ -13,9 +13,8 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.datastax.oss.streaming.ai.jstl;
+package ai.langstream.ai.agents.commons.jstl;
 
-import edu.umd.cs.findbugs.annotations.SuppressFBWarnings;
 import jakarta.el.ELContext;
 import jakarta.el.ELException;
 import jakarta.el.TypeConverter;
@@ -46,7 +45,6 @@
 public class JstlTypeConverter extends TypeConverter {
     public static final JstlTypeConverter INSTANCE = new JstlTypeConverter();
 
-    @SuppressFBWarnings("NP_BOOLEAN_RETURN_NULL")
     protected Boolean coerceToBoolean(Object value) {
         if (value instanceof byte[]) {
             return Schema.BOOL.decode((byte[]) value);

diff --git a/...ss/streaming/ai/jstl/StandardContext.java → .../agents/commons/jstl/StandardContext.java b/...ss/streaming/ai/jstl/StandardContext.java → .../agents/commons/jstl/StandardContext.java
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.datastax.oss.streaming.ai.jstl;
+package ai.langstream.ai.agents.commons.jstl;
 
 import jakarta.el.ArrayELResolver;
 import jakarta.el.BeanNameELResolver;

diff --git a/...ming/ai/jstl/predicate/JstlPredicate.java → ...commons/jstl/predicate/JstlPredicate.java b/...ming/ai/jstl/predicate/JstlPredicate.java → ...commons/jstl/predicate/JstlPredicate.java
@@ -13,10 +13,10 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.datastax.oss.streaming.ai.jstl.predicate;
+package ai.langstream.ai.agents.commons.jstl.predicate;
 
-import com.datastax.oss.streaming.ai.TransformContext;
-import com.datastax.oss.streaming.ai.jstl.JstlEvaluator;
+import ai.langstream.ai.agents.commons.TransformContext;
+import ai.langstream.ai.agents.commons.jstl.JstlEvaluator;
 import jakarta.el.ELException;
 import jakarta.el.PropertyNotFoundException;
 import lombok.extern.slf4j.Slf4j;

diff --git a/...ai/jstl/predicate/TransformPredicate.java → ...ns/jstl/predicate/TransformPredicate.java b/...ai/jstl/predicate/TransformPredicate.java → ...ns/jstl/predicate/TransformPredicate.java
@@ -13,9 +13,9 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.datastax.oss.streaming.ai.jstl.predicate;
+package ai.langstream.ai.agents.commons.jstl.predicate;
 
-import com.datastax.oss.streaming.ai.TransformContext;
+import ai.langstream.ai.agents.commons.TransformContext;
 import java.util.function.Predicate;
 
 /**

diff --git a/langstream-agents/langstream-ai-agents/pom.xml b/langstream-agents/langstream-ai-agents/pom.xml
@@ -36,6 +36,12 @@
       <version>${project.version}</version>
       <scope>provided</scope>
     </dependency>
+    <dependency>
+      <groupId>${project.groupId}</groupId>
+      <artifactId>langstream-agents-commons</artifactId>
+      <version>${project.version}</version>
+    </dependency>
+
     <dependency>
       <groupId>com.google.api-client</groupId>
       <artifactId>google-api-client</artifactId>
@@ -206,18 +212,7 @@
       <groupId>com.samskivert</groupId>
       <artifactId>jmustache</artifactId>
     </dependency>
-    <dependency>
-      <groupId>org.apache.tomcat.embed</groupId>
-      <artifactId>tomcat-embed-el</artifactId>
-    </dependency>
-    <dependency>
-      <groupId>org.apache.commons</groupId>
-      <artifactId>commons-collections4</artifactId>
-    </dependency>
-    <dependency>
-      <groupId>com.fasterxml.jackson.core</groupId>
-      <artifactId>jackson-databind</artifactId>
-    </dependency>
+
     <dependency>
       <groupId>org.projectlombok</groupId>
       <artifactId>lombok</artifactId>

diff --git a/...-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/GenAIToolKitAgent.java b/...-agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/GenAIToolKitAgent.java
@@ -15,6 +15,8 @@
  */
 package ai.langstream.ai.agents;
 
+import ai.langstream.ai.agents.commons.TransformContext;
+import ai.langstream.ai.agents.commons.TransformSchemaType;
 import ai.langstream.ai.agents.datasource.DataSourceProviderRegistry;
 import ai.langstream.ai.agents.services.ServiceProviderRegistry;
 import ai.langstream.api.runner.code.AbstractAgentCode;
@@ -26,11 +28,9 @@
 import ai.langstream.api.runner.code.SimpleRecord;
 import ai.langstream.api.runner.topics.TopicProducer;
 import ai.langstream.api.runtime.ComponentType;
-import com.datastax.oss.streaming.ai.TransformContext;
+import com.datastax.oss.streaming.ai.StepPredicatePair;
 import com.datastax.oss.streaming.ai.TransformStep;
 import com.datastax.oss.streaming.ai.datasource.QueryStepDataSource;
-import com.datastax.oss.streaming.ai.jstl.predicate.StepPredicatePair;
-import com.datastax.oss.streaming.ai.model.TransformSchemaType;
 import com.datastax.oss.streaming.ai.model.config.StepConfig;
 import com.datastax.oss.streaming.ai.model.config.TransformStepConfig;
 import com.datastax.oss.streaming.ai.services.ServiceProvider;

diff --git a/...agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/rerank/ReRankAgent.java b/...agents/langstream-ai-agents/src/main/java/ai/langstream/ai/agents/rerank/ReRankAgent.java
@@ -18,11 +18,11 @@
 import static ai.langstream.ai.agents.GenAIToolKitAgent.transformContextToRecord;
 
 import ai.langstream.ai.agents.GenAIToolKitAgent;
+import ai.langstream.ai.agents.commons.TransformContext;
+import ai.langstream.ai.agents.commons.jstl.JstlEvaluator;
 import ai.langstream.api.runner.code.Record;
 import ai.langstream.api.runner.code.SingleRecordAgentProcessor;
 import ai.langstream.api.util.ConfigurationUtils;
-import com.datastax.oss.streaming.ai.TransformContext;
-import com.datastax.oss.streaming.ai.jstl.JstlEvaluator;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.HashMap;

diff --git a/...eam-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/CastStep.java b/...eam-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/CastStep.java
@@ -17,8 +17,9 @@
 
 import static com.datastax.oss.streaming.ai.util.TransformFunctionUtil.attemptJsonConversion;
 
-import com.datastax.oss.streaming.ai.jstl.JstlTypeConverter;
-import com.datastax.oss.streaming.ai.model.TransformSchemaType;
+import ai.langstream.ai.agents.commons.TransformContext;
+import ai.langstream.ai.agents.commons.TransformSchemaType;
+import ai.langstream.ai.agents.commons.jstl.JstlTypeConverter;
 import java.sql.Time;
 import java.sql.Timestamp;
 import java.time.Instant;