[integrations][plan] Handle finish_reason in OpenAI-family connections

integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnection.java

@@ -253,6 +253,11 @@ private boolean apiVersionSupportsStructuredOutput() {
                 >= 0;
     }
 
+    /**
+     * Returns the model response. When the provider reports a finish reason it is carried verbatim
+     * in {@code extraArgs} under {@code finish_reason}, including values outside the documented
+     * set, and the entry is absent when the provider reports none.
+     */
     @Override
     public ChatMessage chat(
             List<ChatMessage> messages, List<Tool> tools, Map<String, Object> modelParams) {
@@ -271,6 +276,10 @@ public ChatMessage chat(
      * carries no model information. Leaving that parameter unset therefore keeps even a capable
      * deployment on the fallback.
      *
+     * <p>When the provider reports a finish reason it is carried verbatim in {@code extraArgs}
+     * under {@code finish_reason}, including values outside the documented set, and the entry is
+     * absent when the provider reports none.
+     *
      * @throws IllegalArgumentException if the schema is applied natively while {@code
      *     additional_kwargs} also carries a {@code response_format}, since the two would otherwise
      *     compete on the same request
@@ -305,9 +314,16 @@ ChatMessage toResponse(ChatCompletion completion, Map<String, Object> modelParam
         String modelOfAzureDeployment =
                 modelParams != null ? (String) modelParams.get("model_of_azure_deployment") : null;
 
+        ChatCompletion.Choice choice = completion.choices().get(0);
         ChatMessage response =
-                OpenAIChatCompletionsUtils.convertFromOpenAIMessage(
-                        completion.choices().get(0).message());
+                OpenAIChatCompletionsUtils.convertFromOpenAIMessage(choice.message());
+
+        // ChatCompletion.Choice#finishReason throws OpenAIInvalidDataException when the member is
+        // absent or null, so the value is read through the raw field.
+        choice._finishReason()
+                .asKnown()
+                .ifPresent(
+                        reason -> response.getExtraArgs().put("finish_reason", reason.asString()));
 
         if (modelOfAzureDeployment != null
                 && !modelOfAzureDeployment.isBlank()

integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnection.java

@@ -171,12 +171,22 @@ protected boolean supportsNativeStructuredOutput(String effectiveModel) {
                 || NATIVE_STRUCTURED_OUTPUT_MODELS.contains(effectiveModel);
     }
 
+    /**
+     * Returns the model response. When the provider reports a finish reason it is carried verbatim
+     * in {@code extraArgs} under {@code finish_reason}, including values outside the documented
+     * set, and the entry is absent when the provider reports none.
+     */
     @Override
     public ChatMessage chat(
             List<ChatMessage> messages, List<Tool> tools, Map<String, Object> modelParams) {
         return doChat(messages, tools, modelParams, null);
     }
 
+    /**
+     * Returns the model response. When the provider reports a finish reason it is carried verbatim
+     * in {@code extraArgs} under {@code finish_reason}, including values outside the documented
+     * set, and the entry is absent when the provider reports none.
+     */
     @Override
     public ChatMessage chat(
             List<ChatMessage> messages,
@@ -194,9 +204,16 @@ private ChatMessage doChat(
         ChatCompletionCreateParams params =
                 buildRequest(messages, tools, modelParams, outputSchema);
         ChatCompletion completion = client.chat().completions().create(params);
+        ChatCompletion.Choice choice = completion.choices().get(0);
         ChatMessage response =
-                OpenAIChatCompletionsUtils.convertFromOpenAIMessage(
-                        completion.choices().get(0).message());
+                OpenAIChatCompletionsUtils.convertFromOpenAIMessage(choice.message());
+
+        // ChatCompletion.Choice#finishReason throws OpenAIInvalidDataException when the member is
+        // absent or null, so the value is read through the raw field.
+        choice._finishReason()
+                .asKnown()
+                .ifPresent(
+                        reason -> response.getExtraArgs().put("finish_reason", reason.asString()));
 
         // Stash token usage
         if (completion.usage().isPresent()) {

integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnectionTest.java

@@ -19,6 +19,9 @@
 package org.apache.flink.agents.integrations.chatmodels.openai;
 
 import com.fasterxml.jackson.core.type.TypeReference;
+import com.openai.core.JsonField;
+import com.openai.core.JsonMissing;
+import com.openai.core.JsonValue;
 import com.openai.errors.BadRequestException;
 import com.openai.models.ChatModel;
 import com.openai.models.ResponseFormatJsonSchema;
@@ -50,9 +53,11 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.stream.Stream;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /**
@@ -560,6 +565,106 @@ void testNoTokenMetricsWhenMetricsInputsAreIncomplete(
                 .doesNotContainKeys("model_name", "promptTokens", "completionTokens");
     }
 
+    @Test
+    @DisplayName("The finish reason reported by the provider reaches the response extra args")
+    void testResponseCarriesFinishReason() {
+        ChatMessage response =
+                connection()
+                        .toResponse(
+                                completionWithFinishReasonAndUsage(
+                                        JsonField.of(ChatCompletion.Choice.FinishReason.LENGTH)),
+                                params("gpt-4o-mini"));
+
+        assertThat(response.getExtraArgs()).containsEntry("finish_reason", "length");
+    }
+
+    @Test
+    @DisplayName("A finish reason outside the documented set is stored as received")
+    void testResponseCarriesUnknownFinishReasonVerbatim() {
+        ChatMessage response =
+                connection()
+                        .toResponse(
+                                completionWithFinishReasonAndUsage(
+                                        JsonField.of(
+                                                ChatCompletion.Choice.FinishReason.of(
+                                                        "some_vendor_reason"))),
+                                params("gpt-4o-mini"));
+
+        assertThat(response.getExtraArgs()).containsEntry("finish_reason", "some_vendor_reason");
+    }
+
+    @Test
+    @DisplayName("An empty finish reason is recorded rather than discarded")
+    void testResponseCarriesEmptyFinishReason() {
+        // The choice carries a value, so it is recorded; emptiness is not treated as absence.
+        ChatMessage response =
+                connection()
+                        .toResponse(
+                                completionWithFinishReasonAndUsage(
+                                        JsonField.of(ChatCompletion.Choice.FinishReason.of(""))),
+                                params("gpt-4o-mini"));
+
+        assertThat(response.getExtraArgs()).containsEntry("finish_reason", "");
+    }
+
+    @Test
+    @DisplayName("The finish reason is captured independently of the token metrics")
+    void testResponseCarriesFinishReasonWithoutTokenMetrics() {
+        // The metrics need both a backing model and a usage report, and neither is supplied here,
+        // so the absent promptTokens proves that branch did not run and could not have written it.
+        ChatMessage response =
+                connection()
+                        .toResponse(
+                                completionWithFinishReasonWithoutUsage(
+                                        JsonField.of(
+                                                ChatCompletion.Choice.FinishReason.TOOL_CALLS)),
+                                params(null));
+
+        assertThat(response.getExtraArgs())
+                .containsEntry("finish_reason", "tool_calls")
+                .doesNotContainKey("promptTokens");
+    }
+
+    private static Stream<Arguments> finishReasonsWithoutAValue() {
+        return Stream.of(
+                Arguments.of("field_absent", JsonMissing.of()),
+                Arguments.of("json_null", JsonValue.from(null)));
+    }
+
+    @ParameterizedTest(name = "{0}")
+    @MethodSource("finishReasonsWithoutAValue")
+    @DisplayName("A choice carrying no finish reason value yields no key and no error")
+    void testNoFinishReasonKeyWhenTheChoiceCarriesNoValue(
+            String label, JsonField<ChatCompletion.Choice.FinishReason> finishReason) {
+        // ChatCompletion.Choice#finishReason throws OpenAIInvalidDataException for both of these
+        // inputs, so reading the value has to go through the raw field.
+        AzureOpenAIChatModelConnection connection = connection();
+        ChatCompletion completion = completionWithFinishReasonAndUsage(finishReason);
+        AtomicReference<ChatMessage> response = new AtomicReference<>();
+
+        assertThatCode(() -> response.set(connection.toResponse(completion, params("gpt-4o-mini"))))
+                .doesNotThrowAnyException();
+
+        assertThat(response.get().getExtraArgs()).doesNotContainKey("finish_reason");
+    }
+
+    private static ChatCompletion completionWithFinishReasonAndUsage(
+            JsonField<ChatCompletion.Choice.FinishReason> finishReason) {
+        return completionBuilder(finishReason)
+                .usage(
+                        CompletionUsage.builder()
+                                .promptTokens(11L)
+                                .completionTokens(7L)
+                                .totalTokens(18L)
+                                .build())
+                .build();
+    }
+
+    private static ChatCompletion completionWithFinishReasonWithoutUsage(
+            JsonField<ChatCompletion.Choice.FinishReason> finishReason) {
+        return completionBuilder(finishReason).build();
+    }
+
     private static ChatCompletion completionWithUsage(long promptTokens, long completionTokens) {
         return completionBuilder()
                 .usage(
@@ -576,6 +681,11 @@ private static ChatCompletion completionWithoutUsage() {
     }
 
     private static ChatCompletion.Builder completionBuilder() {
+        return completionBuilder(JsonField.of(ChatCompletion.Choice.FinishReason.STOP));
+    }
+
+    private static ChatCompletion.Builder completionBuilder(
+            JsonField<ChatCompletion.Choice.FinishReason> finishReason) {
         ChatCompletionMessage message =
                 ChatCompletionMessage.builder().content("hi").refusal(Optional.empty()).build();
         return ChatCompletion.builder()
@@ -584,7 +694,7 @@ private static ChatCompletion.Builder completionBuilder() {
                 .model(DEPLOYMENT)
                 .addChoice(
                         ChatCompletion.Choice.builder()
-                                .finishReason(ChatCompletion.Choice.FinishReason.STOP)
+                                .finishReason(finishReason)
                                 .index(0L)
                                 .logprobs(Optional.empty())
                                 .message(message)

integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/FakeOpenAICompletionsEndpoint.java

@@ -0,0 +1,106 @@
+/*
+ * 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.
+ */
+
+package org.apache.flink.agents.integrations.chatmodels.openai;
+
+import com.sun.net.httpserver.HttpServer;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.nio.charset.StandardCharsets;
+
+/**
+ * A loopback HTTP endpoint that answers every request with a successful chat completion, so a
+ * connection's response-handling path can be exercised without a live API call. The choice's {@code
+ * finish_reason} and the top-level {@code usage} block are chosen per instance, including the
+ * shapes that carry no finish reason at all: the member set to JSON null, and the member left out
+ * of the choice entirely.
+ */
+final class FakeOpenAICompletionsEndpoint implements AutoCloseable {
+
+    private static final String CONTENT = "hi";
+
+    private static final String NO_FINISH_REASON_MEMBER = "";
+
+    private static final String USAGE_MEMBER =
+            ",\"usage\":{\"prompt_tokens\":11,\"completion_tokens\":7,\"total_tokens\":18}";
+
+    private final HttpServer server;
+
+    private FakeOpenAICompletionsEndpoint(HttpServer server) {
+        this.server = server;
+    }
+
+    static FakeOpenAICompletionsEndpoint servingFinishReason(String finishReason)
+            throws IOException {
+        return serving(finishReasonMember("\"" + finishReason + "\""), USAGE_MEMBER);
+    }
+
+    static FakeOpenAICompletionsEndpoint servingFinishReasonWithoutUsage(String finishReason)
+            throws IOException {
+        return serving(finishReasonMember("\"" + finishReason + "\""), "");
+    }
+
+    static FakeOpenAICompletionsEndpoint servingNullFinishReason() throws IOException {
+        return serving(finishReasonMember("null"), USAGE_MEMBER);
+    }
+
+    static FakeOpenAICompletionsEndpoint servingNoFinishReasonMember() throws IOException {
+        return serving(NO_FINISH_REASON_MEMBER, USAGE_MEMBER);
+    }
+
+    private static String finishReasonMember(String value) {
+        return "\"finish_reason\":" + value + ",";
+    }
+
+    private static FakeOpenAICompletionsEndpoint serving(
+            String finishReasonMember, String usageMember) throws IOException {
+        String completion =
+                "{\"id\":\"completion-1\",\"object\":\"chat.completion\",\"created\":0,"
+                        + "\"model\":\"gpt-4o\",\"choices\":[{"
+                        + finishReasonMember
+                        + "\"index\":0,\"message\":{\"role\":\"assistant\",\"content\":\""
+                        + CONTENT
+                        + "\"}}]"
+                        + usageMember
+                        + "}";
+
+        HttpServer server = HttpServer.create(new InetSocketAddress("127.0.0.1", 0), 0);
+        byte[] body = completion.getBytes(StandardCharsets.UTF_8);
+        server.createContext(
+                "/",
+                exchange -> {
+                    exchange.getResponseHeaders().add("Content-Type", "application/json");
+                    exchange.sendResponseHeaders(200, body.length);
+                    exchange.getResponseBody().write(body);
+                    exchange.close();
+                });
+        server.setExecutor(null);
+        server.start();
+        return new FakeOpenAICompletionsEndpoint(server);
+    }
+
+    String baseUrl() {
+        return "http://127.0.0.1:" + server.getAddress().getPort();
+    }
+
+    @Override
+    public void close() {
+        server.stop(0);
+    }
+}