From a1f1f3c8d29d0227b55e916736de7652a683e9e9 Mon Sep 17 00:00:00 2001
From: Enrico Olivelli <enrico.olivelli@datastax.com>
Date: Wed, 13 Dec 2023 14:11:10 +0100
Subject: [PATCH] Fix IT

---
 .../ai/agents/commons/jstl/JstlEvaluator.java       |  3 +++
 .../ai/agents/commons/jstl/JstlFunctions.java       |  4 ++++
 .../oss/streaming/ai/jstl/JstlEvaluatorTest.java    | 12 ++++++++++++
 .../java/ai/langstream/tests/PythonAgentsIT.java    |  6 +++---
 .../test/resources/apps/python-source/pipeline.yaml | 13 ++++++++++---
 5 files changed, 32 insertions(+), 6 deletions(-)

diff --git a/langstream-agents/langstream-agents-commons/src/main/java/ai/langstream/ai/agents/commons/jstl/JstlEvaluator.java b/langstream-agents/langstream-agents-commons/src/main/java/ai/langstream/ai/agents/commons/jstl/JstlEvaluator.java
index c8360288c..53dbf26e6 100644
--- a/langstream-agents/langstream-agents-commons/src/main/java/ai/langstream/ai/agents/commons/jstl/JstlEvaluator.java
+++ b/langstream-agents/langstream-agents-commons/src/main/java/ai/langstream/ai/agents/commons/jstl/JstlEvaluator.java
@@ -45,6 +45,9 @@ public JstlEvaluator(String expression, Class<? extends T> type) {
 
     @SneakyThrows
     private void registerFunctions() {
+        this.expressionContext
+                .getFunctionMapper()
+                .mapFunction("fn", "length", JstlFunctions.class.getMethod("length", Object.class));
         this.expressionContext
                 .getFunctionMapper()
                 .mapFunction("fn", "toJson", JstlFunctions.class.getMethod("toJson", Object.class));
diff --git a/langstream-agents/langstream-agents-commons/src/main/java/ai/langstream/ai/agents/commons/jstl/JstlFunctions.java b/langstream-agents/langstream-agents-commons/src/main/java/ai/langstream/ai/agents/commons/jstl/JstlFunctions.java
index 9d4055937..bddc4b780 100644
--- a/langstream-agents/langstream-agents-commons/src/main/java/ai/langstream/ai/agents/commons/jstl/JstlFunctions.java
+++ b/langstream-agents/langstream-agents-commons/src/main/java/ai/langstream/ai/agents/commons/jstl/JstlFunctions.java
@@ -187,6 +187,10 @@ public static Map<String, Object> emptyMap() {
         return Map.of();
     }
 
+    public static long length(Object o) {
+        return o == null ? 0 : toString(o).length();
+    }
+
     public static Map<String, Object> mapOf(Object... field) {
         Map<String, Object> result = new HashMap<>();
         for (int i = 0; i < field.length; i += 2) {
diff --git a/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/jstl/JstlEvaluatorTest.java b/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/jstl/JstlEvaluatorTest.java
index 15f0f8c51..810566f57 100644
--- a/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/jstl/JstlEvaluatorTest.java
+++ b/langstream-agents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/jstl/JstlEvaluatorTest.java
@@ -65,6 +65,18 @@ void testPrimitiveValue() {
         assertEquals("test-message", value);
     }
 
+    @Test
+    void testLength() {
+        MutableRecord primitiveStringContext =
+                Utils.createContextWithPrimitiveRecord(Schema.STRING, "test-message", "");
+
+        String value =
+                new JstlEvaluator<>("${fn:length(value)}", String.class)
+                        .evaluate(primitiveStringContext);
+
+        assertEquals("12", value);
+    }
+
     @Test
     void testNowFunction() {
         MutableRecord primitiveStringContext =
diff --git a/langstream-e2e-tests/src/test/java/ai/langstream/tests/PythonAgentsIT.java b/langstream-e2e-tests/src/test/java/ai/langstream/tests/PythonAgentsIT.java
index 3a3880b2b..5be8c4807 100644
--- a/langstream-e2e-tests/src/test/java/ai/langstream/tests/PythonAgentsIT.java
+++ b/langstream-e2e-tests/src/test/java/ai/langstream/tests/PythonAgentsIT.java
@@ -117,11 +117,11 @@ public void testSource() {
                                 .split(" "));
         log.info("Output: {}", output);
         String bigPayload = Strings.repeat("test", 10000);
+        String value = "the length is " + bigPayload.length();
         Assertions.assertTrue(
                 output.contains(
-                        "{\"record\":{\"key\":null,\"value\":\""
-                                + bigPayload
-                                + "\",\"headers\":{}}"));
+                        "{\"record\":{\"key\":null,\"value\":\"" + value + "\",\"headers\":{}}"),
+                "Output doesn't contain the expected payload: " + output);
 
         deleteAppAndAwaitCleanup(tenant, applicationId);
     }
diff --git a/langstream-e2e-tests/src/test/resources/apps/python-source/pipeline.yaml b/langstream-e2e-tests/src/test/resources/apps/python-source/pipeline.yaml
index 6e8b69411..357df2401 100644
--- a/langstream-e2e-tests/src/test/resources/apps/python-source/pipeline.yaml
+++ b/langstream-e2e-tests/src/test/resources/apps/python-source/pipeline.yaml
@@ -22,12 +22,19 @@ topics:
     creation-mode: create-if-not-exists
     schema:
       type: string
+resources:
+  size: 2
 pipeline:
   - name: "Source using Python"
-    resources:
-      size: 2
     id: "test-python-source"
     type: "python-source"
+    configuration:
+      className: example.TestSource
+  - name: "Compute length"
+    id: "test-python-source"
+    type: "compute"
     output: ls-test-output
     configuration:
-      className: example.TestSource
\ No newline at end of file
+      fields:
+        - name: "value"
+          expression: "fn:concat('the length is ', fn:len(value))"
\ No newline at end of file