[FLINK-38581][model] Support surfacing error message

Author: yunfengzhou-hub

docs/content.zh/docs/connectors/models/openai.md

@@ -168,7 +168,12 @@ FROM ML_PREDICT(
                 <ul>
                     <li><code>retry</code>: 重试发送请求。重试行为受 retry-num、retry-fallback-strategy、retry-backoff-strategy 和 retry-backoff-base-interval 限制。</li>
                     <li><code>failover</code>: 抛出异常并使 Flink 作业失败。</li>
-                    <li><code>ignore</code>: 忽略导致错误的输入并继续。错误本身将记录在日志中。</li>
+                    <li><code>ignore</code>: 忽略导致错误的输入并继续执行。错误本身将被记录在日志中。您还可以指定以下元数据列，以便在输出流中显示错误信息。
+                        <ul>
+                            <li><code>error-string</code>: 与错误相关的消息</li>
+                            <li><code>http-status-code</code>: HTTP状态码</li>
+                        </ul>
+                    </li>
                 </ul>
             </td>
         </tr>

docs/content/docs/connectors/models/openai.md

@@ -168,7 +168,12 @@ FROM ML_PREDICT(
                 <ul>
                     <li><code>retry</code>: Retry sending the request. The retrying behavior is limited by retry-num, retry-fallback-strategy, retry-backoff-strategy and retry-backoff-base-interval.</li>
                     <li><code>failover</code>: Throw exceptions and fail the Flink job.</li>
-                    <li><code>ignore</code>: Ignore the input that caused the error and continue. The error itself would be recorded in log.</li>
+                    <li><code>ignore</code>: Ignore the input that caused the error and continue. The error itself would be recorded in log. You can also specify the following metadata columns to surface the error message in the output stream.
+                        <ul>
+                            <li><code>error-string</code>: A message associated with the error</li>
+                            <li><code>http-status-code</code>: The HTTP status code</li>
+                        </ul>
+                    </li>
                 </ul>
             </td>
         </tr>

flink-models/flink-model-openai/src/main/java/org/apache/flink/model/openai/AbstractOpenAIModelFunction.java

@@ -19,13 +19,17 @@
 package org.apache.flink.model.openai;
 
 import org.apache.flink.configuration.ReadableConfig;
+import org.apache.flink.table.api.DataTypes;
 import org.apache.flink.table.api.config.ExecutionConfigOptions;
 import org.apache.flink.table.catalog.Column;
 import org.apache.flink.table.catalog.ResolvedSchema;
+import org.apache.flink.table.data.GenericRowData;
 import org.apache.flink.table.data.RowData;
+import org.apache.flink.table.data.binary.BinaryStringData;
 import org.apache.flink.table.factories.ModelProviderFactory;
 import org.apache.flink.table.functions.AsyncPredictFunction;
 import org.apache.flink.table.functions.FunctionContext;
+import org.apache.flink.table.types.DataType;
 import org.apache.flink.table.types.logical.LogicalType;
 import org.apache.flink.table.types.logical.VarCharType;
 import org.apache.flink.util.ExceptionUtils;
@@ -41,6 +45,7 @@
 
 import java.io.IOException;
 import java.time.Duration;
+import java.util.Arrays;
 import java.util.Collection;
 import java.util.Collections;
 import java.util.HashSet;
@@ -78,6 +83,7 @@ public abstract class AbstractOpenAIModelFunction extends AsyncPredictFunction {
     private final String model;
     @Nullable private final Integer maxContextSize;
     private final ContextOverflowAction contextOverflowAction;
+    protected final List<String> outputColumnNames;
 
     public AbstractOpenAIModelFunction(
             ModelProviderFactory.Context factoryContext, ReadableConfig config) {
@@ -140,6 +146,9 @@ public AbstractOpenAIModelFunction(
                 factoryContext.getCatalogModel().getResolvedInputSchema(),
                 new VarCharType(VarCharType.MAX_LENGTH),
                 "input");
+
+        this.outputColumnNames =
+                factoryContext.getCatalogModel().getResolvedOutputSchema().getColumnNames();
     }
 
     @Override
@@ -184,23 +193,19 @@ public void close() throws Exception {
     protected void validateSingleColumnSchema(
             ResolvedSchema schema, LogicalType expectedType, String inputOrOutput) {
         List<Column> columns = schema.getColumns();
-        if (columns.size() != 1) {
+        List<String> physicalColumnNames =
+                columns.stream()
+                        .filter(Column::isPhysical)
+                        .map(Column::getName)
+                        .collect(Collectors.toList());
+        if (physicalColumnNames.size() != 1) {
             throw new IllegalArgumentException(
                     String.format(
-                            "Model should have exactly one %s column, but actually has %s columns: %s",
-                            inputOrOutput,
-                            columns.size(),
-                            columns.stream().map(Column::getName).collect(Collectors.toList())));
-        }
-
-        Column column = columns.get(0);
-        if (!column.isPhysical()) {
-            throw new IllegalArgumentException(
-                    String.format(
-                            "%s column %s should be a physical column, but is a %s.",
-                            inputOrOutput, column.getName(), column.getClass()));
+                            "Model should have exactly one %s physical column, but actually has %s physical columns: %s",
+                            inputOrOutput, physicalColumnNames.size(), physicalColumnNames));
         }
 
+        Column column = schema.getColumn(physicalColumnNames.get(0)).get();
         if (!expectedType.equals(column.getDataType().getLogicalType())) {
             throw new IllegalArgumentException(
                     String.format(
@@ -210,6 +215,33 @@ protected void validateSingleColumnSchema(
                             expectedType,
                             column.getDataType().getLogicalType()));
         }
+
+        List<Column> metadataColumns =
+                columns.stream()
+                        .filter(x -> x instanceof Column.MetadataColumn)
+                        .collect(Collectors.toList());
+        if (!metadataColumns.isEmpty()) {
+            Preconditions.checkArgument(
+                    "output".equals(inputOrOutput), "Only output schema supports metadata column");
+
+            for (Column metadataColumn : metadataColumns) {
+                ErrorMessageMetadata errorMessageMetadata =
+                        ErrorMessageMetadata.get(metadataColumn.getName());
+                Preconditions.checkNotNull(
+                        errorMessageMetadata,
+                        String.format(
+                                "Unexpected metadata column %s. Supported metadata columns:\n%s",
+                                metadataColumn.getName(),
+                                ErrorMessageMetadata.getAllKeysAndDescriptions()));
+                Preconditions.checkArgument(
+                        errorMessageMetadata.dataType.equals(metadataColumn.getDataType()),
+                        String.format(
+                                "Expected metadata column %s to be of type %s, but is of type %s",
+                                metadataColumn.getName(),
+                                errorMessageMetadata.dataType,
+                                metadataColumn.getDataType()));
+            }
+        }
     }
 
     /**
@@ -223,30 +255,52 @@ protected void validateSingleColumnSchema(
      *     appropriate retry and error handling applied, or a null value if the request failed in
      *     the middle and the failure should be ignored.
      */
-    protected <T> CompletableFuture<T> sendAsyncOpenAIRequest(
-            Supplier<CompletableFuture<T>> requestSender) {
+    protected <T> CompletableFuture<Collection<RowData>> sendAsyncOpenAIRequest(
+            Supplier<CompletableFuture<T>> requestSender,
+            Function<T, Collection<RowData>> converter) {
         CompletableFuture<T> result =
                 retryAsync(
                         requestSender,
                         numRetry,
                         retryBackoffBaseIntervalMs,
                         retryBackoffStrategy,
                         null);
-        ErrorHandlingStrategy finalErrorHandlingStrategy =
-                this.errorHandlingStrategy == ErrorHandlingStrategy.RETRY
-                        ? this.retryFallbackStrategy
-                        : this.errorHandlingStrategy;
-        if (finalErrorHandlingStrategy == ErrorHandlingStrategy.IGNORE) {
-            result =
-                    result.exceptionally(
-                            (e) -> {
-                                LOG.warn(
-                                        "The input row data failed to acquire a valid response. Ignoring the input.",
-                                        e);
-                                return null;
-                            });
+        return result.handle((x, throwable) -> this.convertToRowData(x, throwable, converter));
+    }
+
+    private <T> Collection<RowData> convertToRowData(
+            @Nullable T t,
+            @Nullable Throwable throwable,
+            Function<T, Collection<RowData>> converter) {
+        if (throwable != null) {
+            ErrorHandlingStrategy finalErrorHandlingStrategy =
+                    this.errorHandlingStrategy == ErrorHandlingStrategy.RETRY
+                            ? this.retryFallbackStrategy
+                            : this.errorHandlingStrategy;
+            if (finalErrorHandlingStrategy == ErrorHandlingStrategy.FAILOVER) {
+                throw new RuntimeException(throwable);
+            } else {
+                LOG.warn(
+                        "The input row data failed to acquire a valid response. Ignoring the input.",
+                        throwable);
+                GenericRowData rowData = new GenericRowData(this.outputColumnNames.size());
+                boolean isMetadataSet = false;
+                for (int i = 0; i < this.outputColumnNames.size(); i++) {
+                    String columnName = this.outputColumnNames.get(i);
+                    ErrorMessageMetadata errorMessageMetadata =
+                            ErrorMessageMetadata.get(columnName);
+                    if (errorMessageMetadata != null) {
+                        rowData.setField(i, errorMessageMetadata.converter.apply(throwable));
+                        isMetadataSet = true;
+                    }
+                }
+                return isMetadataSet ? Collections.singletonList(rowData) : Collections.emptyList();
+            }
+        } else if (t == null) {
+            return Collections.emptyList();
+        } else {
+            return converter.apply(t);
         }
-        return result;
     }
 
     private <T> CompletableFuture<T> retryAsync(
@@ -348,4 +402,55 @@ public long getMinRetryTotalTime(long baseRetryInterval, int numRetry) {
 
         public abstract long getMinRetryTotalTime(long baseRetryInterval, int numRetry);
     }
+
+    /**
+     * Metadata that can be read from the output row about error messages. Referenced from Flink
+     * HTTP Connector's ReadableMetadata.
+     */
+    protected enum ErrorMessageMetadata {
+        ERROR_STRING(
+                "error-string",
+                DataTypes.STRING(),
+                x -> BinaryStringData.fromString(x.getMessage()),
+                "A message associated with the error"),
+        HTTP_STATUS_CODE(
+                "http-status-code",
+                DataTypes.INT(),
+                e ->
+                        ExceptionUtils.findThrowable(e, OpenAIServiceException.class)
+                                .map(OpenAIServiceException::statusCode)
+                                .orElse(null),
+                "The HTTP status code");
+
+        final String key;
+        final DataType dataType;
+        final Function<Throwable, Object> converter;
+        final String description;
+
+        ErrorMessageMetadata(
+                String key,
+                DataType dataType,
+                Function<Throwable, Object> converter,
+                String description) {
+            this.key = key;
+            this.dataType = dataType;
+            this.converter = converter;
+            this.description = description;
+        }
+
+        static @Nullable ErrorMessageMetadata get(String key) {
+            for (ErrorMessageMetadata value : values()) {
+                if (value.key.equals(key)) {
+                    return value;
+                }
+            }
+            return null;
+        }
+
+        static String getAllKeysAndDescriptions() {
+            return Arrays.stream(values())
+                    .map(value -> value.key + ":\t" + value.description)
+                    .collect(Collectors.joining("\n"));
+        }
+    }
 }

flink-models/flink-model-openai/src/main/java/org/apache/flink/model/openai/OpenAIChatModelFunction.java

@@ -33,11 +33,8 @@
 import com.openai.models.chat.completions.ChatCompletionCreateParams.ResponseFormat;
 import com.openai.services.async.chat.ChatCompletionServiceAsync;
 
-import javax.annotation.Nullable;
-
 import java.util.Arrays;
 import java.util.Collection;
-import java.util.Collections;
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.stream.Collectors;
@@ -53,6 +50,7 @@ public class OpenAIChatModelFunction extends AbstractOpenAIModelFunction {
     private final String model;
     private final String systemPrompt;
     private final Configuration config;
+    private final int outputColumnIndex;
 
     public OpenAIChatModelFunction(
             ModelProviderFactory.Context factoryContext, ReadableConfig config) {
@@ -64,6 +62,21 @@ public OpenAIChatModelFunction(
                 factoryContext.getCatalogModel().getResolvedOutputSchema(),
                 new VarCharType(VarCharType.MAX_LENGTH),
                 "output");
+        this.outputColumnIndex = getOutputColumnIndex();
+    }
+
+    private int getOutputColumnIndex() {
+        for (int i = 0; i < this.outputColumnNames.size(); i++) {
+            String columnName = this.outputColumnNames.get(i);
+            if (ErrorMessageMetadata.get(columnName) == null) {
+                // Prior checks have guaranteed that there is one and only one physical output
+                // column.
+                return i;
+            }
+        }
+        throw new IllegalArgumentException(
+                "There should be one and only one physical output column. Actual columns: "
+                        + this.outputColumnNames);
     }
 
     @Override
@@ -93,21 +106,21 @@ public CompletableFuture<Collection<RowData>> asyncPredictInternal(String input)
 
         ChatCompletionCreateParams params = builder.build();
         ChatCompletionServiceAsync serviceAsync = client.chat().completions();
-        return sendAsyncOpenAIRequest(() -> serviceAsync.create(params))
-                .thenApply(this::convertToRowData);
+        return sendAsyncOpenAIRequest(() -> serviceAsync.create(params), this::convertToRowData);
     }
 
-    private List<RowData> convertToRowData(@Nullable ChatCompletion chatCompletion) {
-        if (chatCompletion == null) {
-            return Collections.emptyList();
-        }
-
+    private List<RowData> convertToRowData(ChatCompletion chatCompletion) {
         return chatCompletion.choices().stream()
                 .map(
-                        choice ->
-                                GenericRowData.of(
-                                        BinaryStringData.fromString(
-                                                choice.message().content().orElse(""))))
+                        choice -> {
+                            GenericRowData rowData =
+                                    new GenericRowData(this.outputColumnNames.size());
+                            rowData.setField(
+                                    this.outputColumnIndex,
+                                    BinaryStringData.fromString(
+                                            choice.message().content().orElse("")));
+                            return rowData;
+                        })
                 .collect(Collectors.toList());
     }