[query agent] Add test case with output-field that contains a dash ch…

…aracter (LangStream#616)
vectorize-io · Oct 19, 2023 · 095ffcf · 095ffcf
1 parent dec327a
commit 095ffcf
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 4 deletions.
diff --git a/...am-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/QueryStep.java b/...am-agents/langstream-ai-agents/src/main/java/com/datastax/oss/streaming/ai/QueryStep.java
@@ -70,10 +70,12 @@ public class QueryStep implements TransformStep {
         this.dataSource = dataSource;
         this.loopOver = loopOver;
         this.loopOverAccessor = loopOverAccessor;
-        this.fields.forEach(
-                field -> {
-                    fieldsEvaluators.add(new JstlEvaluator<>("${" + field + "}", Object.class));
-                });
+        if (this.fields != null) {
+            this.fields.forEach(
+                    field -> {
+                        fieldsEvaluators.add(new JstlEvaluator<>("${" + field + "}", Object.class));
+                    });
+        }
         if (loopOver != null && !loopOver.isEmpty()) {
             this.loopOverAccessor = new JstlEvaluator<>("${" + loopOver + "}", List.class);
         } else {

diff --git a/...gents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/QueryStepTest.java b/...gents/langstream-ai-agents/src/test/java/com/datastax/oss/streaming/ai/QueryStepTest.java
@@ -552,4 +552,40 @@ public List<Map<String, Object>> fetchData(String query, List<Object> params) {
                 (List<Map<String, Object>>) result.get("command_results");
         assertEquals(List.of(Map.of("foo", "bar"), Map.of("foo", "bar2")), command_results);
     }
+
+    @Test
+    void testSetFieldWithDash() throws Exception {
+
+        String value = "{}";
+
+        QueryStepDataSource dataSource =
+                new QueryStepDataSource() {
+
+                    @Override
+                    public List<Map<String, Object>> fetchData(String query, List<Object> params) {
+                        return List.of(Map.of("foo", "bar"));
+                    }
+                };
+
+        QueryStep queryStep =
+                QueryStep.builder()
+                        .dataSource(dataSource)
+                        .outputFieldName("value.command-results")
+                        .query("select 1")
+                        .build();
+        queryStep.start();
+
+        MutableRecord context =
+                MutableRecord.recordToMutableRecord(SimpleRecord.of(null, value), true);
+
+        queryStep.process(context);
+        ai.langstream.api.runner.code.Record record =
+                MutableRecord.mutableRecordToRecord(context).orElseThrow();
+        Map<String, Object> result = (Map<String, Object>) record.value();
+        log.info("Result: {}", result);
+
+        List<Map<String, Object>> command_results =
+                (List<Map<String, Object>>) result.get("command-results");
+        assertEquals(List.of(Map.of("foo", "bar")), command_results);
+    }
 }