[Feature][Transforms-V2] llm trans support field projection

docs/en/transform-v2/llm.md

@@ -11,17 +11,18 @@ more.
 ## Options
 
 | name                   | type   | required | default value |
-|------------------------|--------|----------|---------------|
+| ---------------------- | ------ | -------- | ------------- |
 | model_provider         | enum   | yes      |               |
 | output_data_type       | enum   | no       | String        |
 | prompt                 | string | yes      |               |
+| inference_columns   | list   | no       |               |
 | model                  | string | yes      |               |
 | api_key                | string | yes      |               |
 | api_path               | string | no       |               |
-| custom_config          | map    | no       |               | 
-| custom_response_parse  | string | no       |               | 
+| custom_config          | map    | no       |               |
+| custom_response_parse  | string | no       |               |
 | custom_request_headers | map    | no       |               |
-| custom_request_body    | map    | no       |               | 
+| custom_request_body    | map    | no       |               |
 
 ### model_provider
 
@@ -62,6 +63,23 @@ The result will be:
 | Eric          | 20  | American   |
 | Guangdong Liu | 20  | Chinese    |
 
+### inference_columns
+
+The `inference_columns` option allows you to specify which columns from the input data should be used as inputs for the LLM. By default, all columns will be used as inputs.
+
+For example:
+```hocon
+transform {
+  LLM {
+    model_provider = OPENAI
+    model = gpt-4o-mini
+    api_key = sk-xxx
+    inference_columns = ["name", "age"]
+    prompt = "Determine whether someone is Chinese or American by their name"
+  }
+}
+```
+
 ### model
 
 The model to use. Different model providers have different models. For example, the OpenAI model can be `gpt-4o-mini`.

docs/zh/transform-v2/llm.md

@@ -8,18 +8,19 @@
 
 ## 属性
 
-| 名称                     | 类型     | 是否必须 | 默认值    |
-|------------------------|--------|------|--------|
-| model_provider         | enum   | yes  |        |
-| output_data_type       | enum   | no   | String |
-| prompt                 | string | yes  |        |
-| model                  | string | yes  |        |
-| api_key                | string | yes  |        |
-| api_path               | string | no   |        |
-| custom_config          | map    | no   |        | 
-| custom_response_parse  | string | no   |        | 
-| custom_request_headers | map    | no   |        |
-| custom_request_body    | map    | no   |        | 
+| 名称                   | 类型   | 是否必须 | 默认值 |
+| ---------------------- | ------ | -------- | ------ |
+| model_provider         | enum   | yes      |        |
+| output_data_type       | enum   | no       | String |
+| prompt                 | string | yes      |        |
+| inference_columns   | list   | no       |        |
+| model                  | string | yes      |        |
+| api_key                | string | yes      |        |
+| api_path               | string | no       |        |
+| custom_config          | map    | no       |        |
+| custom_response_parse  | string | no       |        |
+| custom_request_headers | map    | no       |        |
+| custom_request_body    | map    | no       |        |
 
 ### model_provider
 
@@ -60,6 +61,23 @@ Determine whether someone is Chinese or American by their name
 | Eric          | 20  | American   |
 | Guangdong Liu | 20  | Chinese    |
 
+### inference_columns
+
+`inference_columns`选项允许您指定应该将输入数据中的哪些列用作LLM的输入。默认情况下，所有列都将用作输入。
+
+For example:
+```hocon
+transform {
+  LLM {
+    model_provider = OPENAI
+    model = gpt-4o-mini
+    api_key = sk-xxx
+    inference_columns = ["name", "age"]
+    prompt = "Determine whether someone is Chinese or American by their name"
+  }
+}
+```
+
 ### model
 
 要使用的模型。不同的模型提供者有不同的模型。例如，OpenAI 模型可以是 `gpt-4o-mini`。
@@ -253,4 +271,3 @@ sink {
 }
 ```
 
-

seatunnel-e2e/seatunnel-transforms-v2-e2e/seatunnel-transforms-v2-e2e-part-1/src/test/java/org/apache/seatunnel/e2e/transform/TestLLMIT.java

@@ -93,6 +93,13 @@ public void testLLMWithOpenAIBoolean(TestContainer container)
             throws IOException, InterruptedException {
         Container.ExecResult execResult =
                 container.executeJob("/llm_openai_transform_boolean.conf");
+    }
+
+    @TestTemplate
+    public void testLLMWithOpenAIColumns(TestContainer container)
+            throws IOException, InterruptedException {
+        Container.ExecResult execResult =
+                container.executeJob("/llm_openai_transform_columns.conf");
         Assertions.assertEquals(0, execResult.getExitCode());
     }
 

seatunnel-e2e/seatunnel-transforms-v2-e2e/seatunnel-transforms-v2-e2e-part-1/src/test/resources/llm_openai_transform_columns.conf

@@ -0,0 +1,76 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements.  See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You 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.
+#
+######
+###### This config file is a demonstration of streaming processing in seatunnel config
+######
+
+env {
+  job.mode = "BATCH"
+}
+
+source {
+  FakeSource {
+    row.num = 5
+    schema = {
+      fields {
+        id = "int"
+        name = "string"
+      }
+    }
+    rows = [
+      {fields = [1, "Jia Fan"], kind = INSERT}
+      {fields = [2, "Hailin Wang"], kind = INSERT}
+      {fields = [3, "Tomas"], kind = INSERT}
+      {fields = [4, "Eric"], kind = INSERT}
+      {fields = [5, "Guangdong Liu"], kind = INSERT}
+    ]
+    result_table_name = "fake"
+  }
+}
+
+transform {
+  LLM {
+    source_table_name = "fake"
+    model_provider = OPENAI
+    model = gpt-4o-mini
+    api_key = sk-xxx
+    inference_columns = ["name"]
+    prompt = "Determine whether someone is Chinese or American by their name"
+    openai.api_path = "http://mockserver:1080/v1/chat/completions"
+    result_table_name = "llm_output"
+  }
+}
+
+sink {
+  Assert {
+    source_table_name = "llm_output"
+    rules =
+      {
+        field_rules = [
+          {
+            field_name = llm_output
+            field_type = string
+            field_value = [
+              {
+                rule_type = NOT_NULL
+              }
+            ]
+          }
+        ]
+      }
+  }
+}

seatunnel-transforms-v2/src/main/java/org/apache/seatunnel/transform/nlpmodel/llm/LLMTransform.java

@@ -79,6 +79,7 @@ public void open() {
                         new CustomModel(
                                 inputCatalogTable.getSeaTunnelRowType(),
                                 outputDataType.getSqlType(),
+                                config.get(LLMTransformConfig.INFERENCE_COLUMNS),
                                 config.get(LLMTransformConfig.PROMPT),
                                 config.get(LLMTransformConfig.MODEL),
                                 provider.usedLLMPath(config.get(LLMTransformConfig.API_PATH)),
@@ -97,6 +98,7 @@ public void open() {
                         new OpenAIModel(
                                 inputCatalogTable.getSeaTunnelRowType(),
                                 outputDataType.getSqlType(),
+                                config.get(LLMTransformConfig.INFERENCE_COLUMNS),
                                 config.get(LLMTransformConfig.PROMPT),
                                 config.get(LLMTransformConfig.MODEL),
                                 config.get(LLMTransformConfig.API_KEY),

seatunnel-transforms-v2/src/main/java/org/apache/seatunnel/transform/nlpmodel/llm/LLMTransformConfig.java

@@ -21,6 +21,8 @@
 import org.apache.seatunnel.api.configuration.Options;
 import org.apache.seatunnel.transform.nlpmodel.ModelTransformConfig;
 
+import java.util.List;
+
 public class LLMTransformConfig extends ModelTransformConfig {
 
     public static final Option<String> PROMPT =
@@ -29,6 +31,12 @@ public class LLMTransformConfig extends ModelTransformConfig {
                     .noDefaultValue()
                     .withDescription("The prompt of LLM");
 
+    public static final Option<List<String>> INFERENCE_COLUMNS =
+            Options.key("inference_columns")
+                    .listType()
+                    .noDefaultValue()
+                    .withDescription("The row projection field of each inference");
+
     public static final Option<Integer> INFERENCE_BATCH_SIZE =
             Options.key("inference_batch_size")
                     .intType()

seatunnel-transforms-v2/src/main/java/org/apache/seatunnel/transform/nlpmodel/llm/remote/AbstractModel.java

@@ -21,25 +21,59 @@
 import org.apache.seatunnel.shade.com.fasterxml.jackson.databind.node.ArrayNode;
 import org.apache.seatunnel.shade.com.fasterxml.jackson.databind.node.ObjectNode;
 
+import org.apache.seatunnel.api.table.type.SeaTunnelDataType;
 import org.apache.seatunnel.api.table.type.SeaTunnelRow;
 import org.apache.seatunnel.api.table.type.SeaTunnelRowType;
 import org.apache.seatunnel.api.table.type.SqlType;
 import org.apache.seatunnel.format.json.RowToJsonConverters;
 
+import com.google.common.annotations.VisibleForTesting;
+
 import java.io.IOException;
+import java.util.ArrayList;
 import java.util.List;
 
 public abstract class AbstractModel implements Model {
 
     protected static final ObjectMapper OBJECT_MAPPER = new ObjectMapper();
-    private final RowToJsonConverters.RowToJsonConverter rowToJsonConverters;
+    private final RowToJsonConverters.RowToJsonConverter rowToJsonConverter;
+    private final SeaTunnelRowType rowType;
     private final String prompt;
     private final SqlType outputType;
+    private final List<String> projectionColumns;
 
-    public AbstractModel(SeaTunnelRowType rowType, SqlType outputType, String prompt) {
+    public AbstractModel(
+            SeaTunnelRowType rowType,
+            SqlType outputType,
+            List<String> projectionColumns,
+            String prompt) {
+        this.rowType = rowType;
         this.prompt = prompt;
         this.outputType = outputType;
-        this.rowToJsonConverters = new RowToJsonConverters().createConverter(rowType, null);
+        this.projectionColumns = projectionColumns;
+        this.rowToJsonConverter = getRowToJsonConverter();
+    }
+
+    public RowToJsonConverters.RowToJsonConverter getRowToJsonConverter() {
+        RowToJsonConverters converters = new RowToJsonConverters();
+        if (projectionColumns != null && !projectionColumns.isEmpty()) {
+            List<SeaTunnelDataType> fieldTypes = new ArrayList<>();
+            for (String fieldName : projectionColumns) {
+                int fieldIndex = rowType.indexOf(fieldName);
+                if (fieldIndex != -1) {
+                    fieldTypes.add(rowType.getFieldType(fieldIndex));
+                } else {
+                    throw new IllegalArgumentException(
+                            "Field name " + fieldName + " does not exist in the row type.");
+                }
+            }
+            SeaTunnelRowType projectionRowType =
+                    new SeaTunnelRowType(
+                            projectionColumns.toArray(new String[0]),
+                            fieldTypes.toArray(new SeaTunnelDataType[0]));
+            return converters.createConverter(projectionRowType, null);
+        }
+        return converters.createConverter(rowType, null);
     }
 
     private String getPromptWithLimit() {
@@ -58,12 +92,31 @@ public List<String> inference(List<SeaTunnelRow> rows) throws IOException {
         ArrayNode rowsNode = OBJECT_MAPPER.createArrayNode();
         for (SeaTunnelRow row : rows) {
             ObjectNode rowNode = OBJECT_MAPPER.createObjectNode();
-            rowToJsonConverters.convert(OBJECT_MAPPER, rowNode, row);
+            rowToJsonConverter.convert(OBJECT_MAPPER, rowNode, createProjectionSeaTunnelRow(row));
             rowsNode.add(rowNode);
         }
         return chatWithModel(getPromptWithLimit(), OBJECT_MAPPER.writeValueAsString(rowsNode));
     }
 
+    @VisibleForTesting
+    public SeaTunnelRow createProjectionSeaTunnelRow(SeaTunnelRow row) {
+        if (row == null || projectionColumns == null || projectionColumns.isEmpty()) {
+            return row;
+        }
+        SeaTunnelRow projectionRow = new SeaTunnelRow(projectionColumns.size());
+        for (int i = 0; i < projectionColumns.size(); i++) {
+            String fieldName = projectionColumns.get(i);
+            int fieldIndex = rowType.indexOf(fieldName);
+            if (fieldIndex != -1) {
+                projectionRow.setField(i, row.getField(fieldIndex));
+            } else {
+                throw new IllegalArgumentException(
+                        "Field name " + fieldName + " does not exist in the row type.");
+            }
+        }
+        return projectionRow;
+    }
+
     protected abstract List<String> chatWithModel(String promptWithLimit, String rowsJson)
             throws IOException;
 

seatunnel-transforms-v2/src/main/java/org/apache/seatunnel/transform/nlpmodel/llm/remote/custom/CustomModel.java

@@ -56,13 +56,14 @@ public class CustomModel extends AbstractModel {
     public CustomModel(
             SeaTunnelRowType rowType,
             SqlType outputType,
+            List<String> projectionColumns,
             String prompt,
             String model,
             String apiPath,
             Map<String, String> header,
             Map<String, Object> body,
             String parse) {
-        super(rowType, outputType, prompt);
+        super(rowType, outputType, projectionColumns, prompt);
         this.apiPath = apiPath;
         this.model = model;
         this.header = header;

seatunnel-transforms-v2/src/main/java/org/apache/seatunnel/transform/nlpmodel/llm/remote/openai/OpenAIModel.java

@@ -54,11 +54,12 @@ public class OpenAIModel extends AbstractModel {
     public OpenAIModel(
             SeaTunnelRowType rowType,
             SqlType outputType,
+            List<String> projectionColumns,
             String prompt,
             String model,
             String apiKey,
             String apiPath) {
-        super(rowType, outputType, prompt);
+        super(rowType, outputType, projectionColumns, prompt);
         this.apiKey = apiKey;
         this.apiPath = apiPath;
         this.model = model;