This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new c4a29ac6f [integrations][plan] Handle finish_reason in OpenAI-family
connections (#1040)
c4a29ac6f is described below
commit c4a29ac6f8fee703f0cb3a57a514dfed626cc155
Author: Weiqing Yang <[email protected]>
AuthorDate: Tue Sep 1 22:02:48 2026 -0700
[integrations][plan] Handle finish_reason in OpenAI-family connections
(#1040)
Generated-by: Claude Code 2.1.234 (Claude Opus 5)
---------
Co-authored-by: weiqingy <[email protected]>
---
.../openai/AzureOpenAIChatModelConnection.java | 20 +-
.../openai/OpenAICompletionsConnection.java | 21 ++-
.../openai/AzureOpenAIChatModelConnectionTest.java | 112 ++++++++++-
.../openai/FakeOpenAICompletionsEndpoint.java | 106 +++++++++++
.../openai/OpenAICompletionsConnectionTest.java | 94 ++++++++++
.../flink/agents/plan/actions/ChatModelAction.java | 33 +++-
.../agents/plan/actions/ChatModelInvoker.java | 3 +
.../plan/actions/ChatModelActionRetryTest.java | 208 ++++++++++++++++++++-
.../agents/plan/actions/ChatModelActionTest.java | 40 ++++
.../chat_models/azure/azure_openai_chat_model.py | 9 +-
.../tests/test_azure_openai_response_parsing.py | 140 ++++++++++++++
.../chat_models/openai/openai_chat_model.py | 9 +-
.../chat_models/openai/openai_utils.py | 11 +-
.../openai/tests/test_openai_response_parsing.py | 203 +++++++++++++++++++-
.../flink_agents/plan/actions/chat_model_action.py | 39 ++++
.../plan/tests/actions/test_chat_model_action.py | 30 +++
.../tests/actions/test_chat_model_action_retry.py | 203 +++++++++++++++++++-
17 files changed, 1262 insertions(+), 19 deletions(-)
diff --git
a/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnection.java
b/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnection.java
index db3270db0..4ce122ad0 100644
---
a/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnection.java
+++
b/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnection.java
@@ -253,6 +253,11 @@ public class AzureOpenAIChatModelConnection extends
BaseChatModelConnection {
>= 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 class AzureOpenAIChatModelConnection extends
BaseChatModelConnection {
* 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 @@ public class AzureOpenAIChatModelConnection extends
BaseChatModelConnection {
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()
diff --git
a/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnection.java
b/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnection.java
index 3af78b72c..4864769e5 100644
---
a/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnection.java
+++
b/integrations/chat-models/openai/src/main/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnection.java
@@ -171,12 +171,22 @@ public class OpenAICompletionsConnection extends
BaseChatModelConnection {
|| 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 @@ public class OpenAICompletionsConnection extends
BaseChatModelConnection {
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()) {
diff --git
a/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnectionTest.java
b/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnectionTest.java
index b47cbbdce..24d448a74 100644
---
a/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnectionTest.java
+++
b/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/AzureOpenAIChatModelConnectionTest.java
@@ -21,6 +21,9 @@ package
org.apache.flink.agents.integrations.chatmodels.openai;
import com.fasterxml.jackson.annotation.JsonSubTypes;
import com.fasterxml.jackson.annotation.JsonTypeInfo;
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;
@@ -52,9 +55,11 @@ import java.util.HashMap;
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;
/**
@@ -617,6 +622,106 @@ class AzureOpenAIChatModelConnectionTest {
.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(
@@ -633,6 +738,11 @@ class AzureOpenAIChatModelConnectionTest {
}
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()
@@ -641,7 +751,7 @@ class AzureOpenAIChatModelConnectionTest {
.model(DEPLOYMENT)
.addChoice(
ChatCompletion.Choice.builder()
-
.finishReason(ChatCompletion.Choice.FinishReason.STOP)
+ .finishReason(finishReason)
.index(0L)
.logprobs(Optional.empty())
.message(message)
diff --git
a/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/FakeOpenAICompletionsEndpoint.java
b/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/FakeOpenAICompletionsEndpoint.java
new file mode 100644
index 000000000..6e1e633b9
--- /dev/null
+++
b/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);
+ }
+}
diff --git
a/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnectionTest.java
b/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnectionTest.java
index 5204de2cf..851e90043 100644
---
a/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnectionTest.java
+++
b/integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/OpenAICompletionsConnectionTest.java
@@ -41,8 +41,10 @@ import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatCode;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
/**
@@ -183,6 +185,98 @@ class OpenAICompletionsConnectionTest {
}
}
+ @Test
+ @DisplayName("The finish reason reported by the provider reaches the
response extra args")
+ void testResponseCarriesFinishReason() throws IOException {
+ try (FakeOpenAICompletionsEndpoint endpoint =
+ FakeOpenAICompletionsEndpoint.servingFinishReason("length")) {
+ ChatMessage response =
+ connection(endpoint.baseUrl())
+ .chat(userMessage(), List.of(), params("gpt-4o"),
null);
+
+ assertThat(response.getExtraArgs()).containsEntry("finish_reason",
"length");
+ }
+ }
+
+ @Test
+ @DisplayName("A finish reason outside the documented set is stored as
received")
+ void testResponseCarriesUnknownFinishReasonVerbatim() throws IOException {
+ try (FakeOpenAICompletionsEndpoint endpoint =
+
FakeOpenAICompletionsEndpoint.servingFinishReason("some_vendor_reason")) {
+ ChatMessage response =
+ connection(endpoint.baseUrl())
+ .chat(userMessage(), List.of(), params("gpt-4o"),
null);
+
+ assertThat(response.getExtraArgs())
+ .containsEntry("finish_reason", "some_vendor_reason");
+ }
+ }
+
+ @Test
+ @DisplayName("An empty finish reason is recorded rather than discarded")
+ void testResponseCarriesEmptyFinishReason() throws IOException {
+ // The choice carries a value, so it is recorded; emptiness is not
treated as absence.
+ try (FakeOpenAICompletionsEndpoint endpoint =
+ FakeOpenAICompletionsEndpoint.servingFinishReason("")) {
+ ChatMessage response =
+ connection(endpoint.baseUrl())
+ .chat(userMessage(), List.of(), params("gpt-4o"),
null);
+
+ assertThat(response.getExtraArgs()).containsEntry("finish_reason",
"");
+ }
+ }
+
+ @Test
+ @DisplayName("The finish reason is captured independently of the token
metrics")
+ void testResponseCarriesFinishReasonWithoutUsage() throws IOException {
+ // The metrics come from the usage report the response here omits, so
the absent
+ // promptTokens proves that branch did not run and could not have
written the reason.
+ try (FakeOpenAICompletionsEndpoint endpoint =
+
FakeOpenAICompletionsEndpoint.servingFinishReasonWithoutUsage("tool_calls")) {
+ ChatMessage response =
+ connection(endpoint.baseUrl())
+ .chat(userMessage(), List.of(), params("gpt-4o"),
null);
+
+ assertThat(response.getExtraArgs())
+ .containsEntry("finish_reason", "tool_calls")
+ .doesNotContainKey("promptTokens");
+ }
+ }
+
+ @Test
+ @DisplayName("A choice with no finish_reason member yields no key and no
error")
+ void testNoFinishReasonKeyWhenMemberAbsent() throws IOException {
+ try (FakeOpenAICompletionsEndpoint endpoint =
+ FakeOpenAICompletionsEndpoint.servingNoFinishReasonMember()) {
+ assertNoFinishReasonKey(endpoint);
+ }
+ }
+
+ @Test
+ @DisplayName("A choice whose finish_reason is JSON null yields no key and
no error")
+ void testNoFinishReasonKeyWhenJsonNull() throws IOException {
+ try (FakeOpenAICompletionsEndpoint endpoint =
+ FakeOpenAICompletionsEndpoint.servingNullFinishReason()) {
+ assertNoFinishReasonKey(endpoint);
+ }
+ }
+
+ private static void assertNoFinishReasonKey(FakeOpenAICompletionsEndpoint
endpoint) {
+ // ChatCompletion.Choice#finishReason throws
OpenAIInvalidDataException for both of these
+ // response shapes, so reading the value has to go through the raw
field.
+ OpenAICompletionsConnection connection =
connection(endpoint.baseUrl());
+ AtomicReference<ChatMessage> response = new AtomicReference<>();
+
+ assertThatCode(
+ () ->
+ response.set(
+ connection.chat(
+ userMessage(), List.of(),
params("gpt-4o"), null)))
+ .doesNotThrowAnyException();
+
+
assertThat(response.get().getExtraArgs()).doesNotContainKey("finish_reason");
+ }
+
@Test
@DisplayName("Native response_format json_schema strict applied for a POJO
on a capable model")
void testNativeAppliedForPojoCapableModel() {
diff --git
a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java
b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java
index 57c191c75..2e8c2476d 100644
---
a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java
+++
b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelAction.java
@@ -97,6 +97,9 @@ public class ChatModelAction {
private static final String RETRY_STATS_CONTEXT = "_RETRY_STATS_CONTEXT";
private static final String TOTAL_RETRY_COUNT = "totalRetryCount";
private static final String TOTAL_RETRY_WAIT_SEC = "totalRetryWaitSec";
+ private static final String FINISH_REASON = "finish_reason";
+ private static final String TRUNCATED_FINISH_REASON = "length";
+ private static final String CONTENT_FILTERED_FINISH_REASON =
"content_filter";
private static final ObjectMapper mapper = new ObjectMapper();
@@ -337,7 +340,7 @@ public class ChatModelAction {
throw new RuntimeException(
String.format("Unsupported output schema %s.",
outputSchema));
}
- Map<String, Object> extraArgs = new HashMap<>();
+ Map<String, Object> extraArgs = new HashMap<>(response.getExtraArgs());
extraArgs.put(STRUCTURED_OUTPUT, structuredOutput);
return new ChatMessage(response.getRole(), output, extraArgs);
}
@@ -576,6 +579,34 @@ public class ChatModelAction {
return routing;
}
+ /**
+ * Rejects a response the provider did not finish emitting. Evaluated once
per chat response,
+ * before it is dispatched as text, structured output, or tool calls. A
finish reason reporting
+ * the content as cut off by the token budget or withheld by content
filtering raises {@link
+ * IllegalStateException}; any other reason, and an absent one, are
accepted.
+ */
+ static void rejectIncompleteResponse(ChatMessage response) {
+ Object finishReason = response.getExtraArgs().get(FINISH_REASON);
+ if (TRUNCATED_FINISH_REASON.equals(finishReason)) {
+ throw new IllegalStateException(
+ String.format(
+ "ChatModel response is truncated
(finish_reason='%s'): it"
+ + " exhausted the completion token budget
before the model"
+ + " finished, so the content is
incomplete. Raise the"
+ + " model's max output tokens, or ask for
a smaller output.",
+ finishReason));
+ }
+ if (CONTENT_FILTERED_FINISH_REASON.equals(finishReason)) {
+ throw new IllegalStateException(
+ String.format(
+ "ChatModel response was withheld by the provider's
content"
+ + " filter (finish_reason='%s'), so the
content is"
+ + " incomplete. Adjust the prompt or the
provider's content"
+ + " filtering configuration.",
+ finishReason));
+ }
+ }
+
static ChatMessage generateStructuredOutputWithReport(
RunnerContext ctx, ChatMessage response, Object outputSchema)
throws Exception {
ExecutionReporters.started(ctx, ExecutionReporter.EntityTypes.PARSER,
STRUCTURED_OUTPUT);
diff --git
a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java
b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java
index 8dfacee4e..4f49feac7 100644
---
a/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java
+++
b/plan/src/main/java/org/apache/flink/agents/plan/actions/ChatModelInvoker.java
@@ -174,6 +174,9 @@ final class ChatModelInvoker {
ExecutionReporters.succeeded(
ctx, ExecutionReporter.EntityTypes.LLM, model,
llmMetadata);
ChatModelAction.recordChatTokenMetrics(chatModel, response,
requestMetricGroup);
+ // A truncated response consumed its full token budget, so the
token metrics
+ // above are recorded before this rejects and abandons the
response.
+ ChatModelAction.rejectIncompleteResponse(response);
// only generate structured output for final response.
if (outputSchema != null && response.getToolCalls().isEmpty())
{
response =
diff --git
a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRetryTest.java
b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRetryTest.java
index a1d7b0cf3..450ef22dc 100644
---
a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRetryTest.java
+++
b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionRetryTest.java
@@ -37,6 +37,9 @@ import org.apache.flink.metrics.Counter;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
+import org.junit.jupiter.params.provider.ValueSource;
import org.mockito.ArgumentCaptor;
import org.mockito.Mock;
import org.mockito.MockitoAnnotations;
@@ -47,6 +50,7 @@ import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;
+import java.util.stream.Stream;
import static org.assertj.core.api.Assertions.assertThat;
import static org.assertj.core.api.Assertions.assertThatThrownBy;
@@ -54,7 +58,10 @@ import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.*;
-/** Tests for retry behavior in {@link ChatModelAction}. */
+/**
+ * Tests for {@link ChatModelAction#chat} driven end to end: retry behavior,
execution reporting,
+ * the finish-reason gate, and tool-response handling.
+ */
class ChatModelActionRetryTest {
private static final Map<String, Object> LLM_METADATA =
@@ -425,6 +432,205 @@ class ChatModelActionRetryTest {
assertThat(promptArgsCaptor.getValue()).isEqualTo(savedPromptArgs);
}
+ @Test
+ void chatRejectsTruncatedTextResponse() throws Exception {
+ RunnerContext reportingCtx = reportingRunnerContext();
+ BaseChatModelSetup chatModel =
configureReportingChatContext(reportingCtx);
+ when(chatModel.chat(any(), any(), any()))
+ .thenReturn(
+ new ChatMessage(
+ MessageRole.ASSISTANT,
+ "partial answ",
+ Map.of("finish_reason", "length")));
+
+ assertThatThrownBy(
+ () ->
+ ChatModelAction.chat(
+ UUID.randomUUID(),
+ "test-model",
+ List.of(new
ChatMessage(MessageRole.USER, "hi")),
+ Map.of(),
+ null,
+ reportingCtx))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("truncated")
+ .hasMessageContaining("token");
+
+ assertThat(sentEvents).isEmpty();
+ }
+
+ @Test
+ void chatRejectsContentFilteredTextResponse() throws Exception {
+ RunnerContext reportingCtx = reportingRunnerContext();
+ BaseChatModelSetup chatModel =
configureReportingChatContext(reportingCtx);
+ when(chatModel.chat(any(), any(), any()))
+ .thenReturn(
+ new ChatMessage(
+ MessageRole.ASSISTANT,
+ "",
+ Map.of("finish_reason", "content_filter")));
+
+ // Both rejection messages interpolate the finish reason, so the
literal
+ // content_filter appears in either one and cannot tell them apart.
These
+ // match prose unique to the filtering message.
+ assertThatThrownBy(
+ () ->
+ ChatModelAction.chat(
+ UUID.randomUUID(),
+ "test-model",
+ List.of(new
ChatMessage(MessageRole.USER, "hi")),
+ Map.of(),
+ null,
+ reportingCtx))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("withheld")
+ .hasMessageContaining("content filter");
+
+ assertThat(sentEvents).isEmpty();
+ }
+
+ @Test
+ void chatRejectsTruncatedToolCallResponseBeforeDispatchingTools() throws
Exception {
+ RunnerContext reportingCtx = reportingRunnerContext();
+ BaseChatModelSetup chatModel =
configureReportingChatContext(reportingCtx);
+ when(chatModel.chat(any(), any(), any()))
+ .thenReturn(
+ new ChatMessage(
+ MessageRole.ASSISTANT,
+ "",
+ List.of(
+ Map.of(
+ "id",
+ "call-1",
+ "function",
+ Map.of("name", "f",
"arguments", ""))),
+ Map.of("finish_reason", "length")));
+
+ assertThatThrownBy(
+ () ->
+ ChatModelAction.chat(
+ UUID.randomUUID(),
+ "test-model",
+ List.of(new
ChatMessage(MessageRole.USER, "hi")),
+ Map.of(),
+ null,
+ reportingCtx))
+ .isInstanceOf(IllegalStateException.class)
+ .hasMessageContaining("truncated");
+
+ // A truncated tool call carries arguments the model never finished
writing,
+ // so no ToolRequestEvent may leave the action.
+ assertThat(sentEvents).isEmpty();
+ }
+
+ @ParameterizedTest
+ @ValueSource(strings = {"length", "content_filter"})
+ void chatRejectedFinishReasonSkipsStructuredOutput(String finishReason)
throws Exception {
+ RunnerContext reportingCtx = reportingRunnerContext();
+ BaseChatModelSetup chatModel =
configureReportingChatContext(reportingCtx);
+ when(chatModel.chat(any(), any(), any()))
+ .thenReturn(
+ new ChatMessage(
+ MessageRole.ASSISTANT,
+ "{\"answer\":\"42\"}",
+ Map.of(
+ "finish_reason",
+ finishReason,
+ "model_name",
+ "provider-model",
+ "promptTokens",
+ 100L,
+ "completionTokens",
+ 50L)));
+
+ assertThatThrownBy(
+ () ->
+ ChatModelAction.chat(
+ UUID.randomUUID(),
+ "test-model",
+ List.of(new
ChatMessage(MessageRole.USER, "hi")),
+ Map.of(),
+ Map.class,
+ reportingCtx))
+ .isInstanceOf(IllegalStateException.class);
+
+ ExecutionReporter reporter = (ExecutionReporter) reportingCtx;
+ // The model call itself succeeded and spent its full token budget, so
both
+ // must be recorded before the response is rejected.
+ verify(reporter)
+ .reportExecutionSucceeded(
+ ExecutionReporter.EntityTypes.LLM, "test-model",
LLM_METADATA);
+ verify(chatModel).recordTokenMetrics(mockActionMetricGroup,
"provider-model", 100L, 50L);
+ // The parse is never attempted, so nothing about it is reported. The
failed
+ // check is unscoped: the rejection must not be reported as a failure
of any
+ // entity, the model call included.
+ verify(reporter, never())
+ .reportExecutionStarted(
+ eq(ExecutionReporter.EntityTypes.PARSER), anyString(),
any());
+ verify(reporter, never())
+ .reportExecutionSucceeded(
+ eq(ExecutionReporter.EntityTypes.PARSER), anyString(),
any());
+ verify(reporter, never())
+ .reportExecutionFailed(anyString(), anyString(), any(), any(),
any());
+ assertThat(sentEvents).isEmpty();
+ }
+
+ @Test
+ void chatIgnoreStrategyDropsRejectedResponseWithoutEvent() throws
Exception {
+ RunnerContext reportingCtx = reportingRunnerContext();
+ BaseChatModelSetup chatModel =
configureReportingChatContext(reportingCtx);
+ when(reportingCtx.getConfig())
+
.thenReturn(readableConfig(Agent.ErrorHandlingStrategy.IGNORE));
+ when(chatModel.chat(any(), any(), any()))
+ .thenReturn(
+ new ChatMessage(
+ MessageRole.ASSISTANT,
+ "partial answ",
+ Map.of("finish_reason", "length")));
+
+ // Under IGNORE the record is dropped: the rejection does not
propagate and no
+ // event carries the truncated content downstream.
+ ChatModelAction.chat(
+ UUID.randomUUID(),
+ "test-model",
+ List.of(new ChatMessage(MessageRole.USER, "hi")),
+ Map.of(),
+ null,
+ reportingCtx);
+
+ assertThat(sentEvents).isEmpty();
+ }
+
+ private static Stream<Map<String, Object>> acceptedFinishReasons() {
+ return Stream.of(
+ Map.of("finish_reason", "stop"),
+ Map.of("finish_reason", "tool_calls"),
+ Map.of("finish_reason", "some_vendor_reason"),
+ Map.of());
+ }
+
+ @ParameterizedTest
+ @MethodSource("acceptedFinishReasons")
+ void chatAcceptedFinishReasonReachesTheResponseEvent(Map<String, Object>
extraArgs)
+ throws Exception {
+ RunnerContext reportingCtx = reportingRunnerContext();
+ BaseChatModelSetup chatModel =
configureReportingChatContext(reportingCtx);
+ when(chatModel.chat(any(), any(), any()))
+ .thenReturn(new ChatMessage(MessageRole.ASSISTANT, "hello",
extraArgs));
+
+ ChatModelAction.chat(
+ UUID.randomUUID(),
+ "test-model",
+ List.of(new ChatMessage(MessageRole.USER, "hi")),
+ Map.of(),
+ null,
+ reportingCtx);
+
+ assertThat(sentEvents).hasSize(1);
+
assertThat(ChatResponseEvent.fromEvent(sentEvents.get(0)).getResponse().getContent())
+ .isEqualTo("hello");
+ }
+
// --- Helper methods ---
private void configureRetryStrategy(int maxRetries, int waitIntervalSec) {
diff --git
a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionTest.java
b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionTest.java
index 46485b5a1..cfdffd03e 100644
---
a/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionTest.java
+++
b/plan/src/test/java/org/apache/flink/agents/plan/actions/ChatModelActionTest.java
@@ -17,16 +17,20 @@
*/
package org.apache.flink.agents.plan.actions;
+import org.apache.flink.agents.api.agents.Agent;
import org.apache.flink.agents.api.chat.messages.ChatMessage;
import org.apache.flink.agents.api.chat.messages.MessageRole;
import org.apache.flink.agents.api.chat.model.BaseChatModelSetup;
+import org.apache.flink.agents.api.context.RunnerContext;
import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup;
+import org.apache.flink.agents.api.trace.ExecutionReporter;
import org.junit.jupiter.api.Test;
import java.util.HashMap;
import java.util.Map;
import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyLong;
import static org.mockito.ArgumentMatchers.anyString;
@@ -34,14 +38,29 @@ import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.verifyNoInteractions;
+import static org.mockito.Mockito.withSettings;
/** Tests for {@link ChatModelAction}. */
class ChatModelActionTest {
+ private static final String PARSEABLE_CONTENT = "{\"answer\":\"42\"}";
+
private static ChatMessage responseWith(Map<String, Object> extraArgs) {
return new ChatMessage(MessageRole.ASSISTANT, "response", extraArgs);
}
+ private static RunnerContext reportingContext() {
+ return mock(RunnerContext.class,
withSettings().extraInterfaces(ExecutionReporter.class));
+ }
+
+ private static ChatMessage generateStructuredOutput(
+ RunnerContext ctx, Map<String, Object> extraArgs) throws Exception
{
+ return ChatModelAction.generateStructuredOutputWithReport(
+ ctx,
+ new ChatMessage(MessageRole.ASSISTANT, PARSEABLE_CONTENT,
extraArgs),
+ Map.class);
+ }
+
@Test
void testRecordChatTokenMetricsRecordsWhenAllKeysPresent() {
BaseChatModelSetup setup = mock(BaseChatModelSetup.class);
@@ -179,4 +198,25 @@ class ChatModelActionTest {
String expected = "{\n \"key\": \"value\"\n}";
assertEquals(expected, ChatModelAction.cleanLlmResponse(input));
}
+
+ @Test
+ void testStructuredOutputParsesToExpectedValue() throws Exception {
+ ChatMessage parsed = generateStructuredOutput(reportingContext(),
Map.of());
+
+ assertEquals(Map.of("answer", "42"),
parsed.getExtraArgs().get(Agent.STRUCTURED_OUTPUT));
+ }
+
+ @Test
+ void testStructuredOutputPreservesInboundExtraArgs() throws Exception {
+ Map<String, Object> extraArgs = new HashMap<>();
+ extraArgs.put("finish_reason", "stop");
+ extraArgs.put("promptTokens", 100L);
+
+ Map<String, Object> parsedArgs =
+ generateStructuredOutput(reportingContext(),
extraArgs).getExtraArgs();
+
+ assertEquals("stop", parsedArgs.get("finish_reason"));
+ assertEquals(100L, parsedArgs.get("promptTokens"));
+ assertTrue(parsedArgs.containsKey(Agent.STRUCTURED_OUTPUT),
parsedArgs.toString());
+ }
}
diff --git
a/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py
b/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py
index 8d264e8bc..65a614c03 100644
---
a/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py
+++
b/python/flink_agents/integrations/chat_models/azure/azure_openai_chat_model.py
@@ -285,7 +285,8 @@ class
AzureOpenAIChatModelConnection(BaseChatModelConnection):
Returns:
-------
ChatMessage
- Model response message
+ Model response message. When the response carries a finish reason,
+ it is available as ``extra_args["finish_reason"]``.
"""
tool_specs = None
if tools is not None:
@@ -362,7 +363,11 @@ class
AzureOpenAIChatModelConnection(BaseChatModelConnection):
extra_args["promptTokens"] = response.usage.prompt_tokens
extra_args["completionTokens"] = response.usage.completion_tokens
- message = response.choices[0].message
+ choice = response.choices[0]
+ if choice.finish_reason is not None:
+ extra_args["finish_reason"] = choice.finish_reason
+
+ message = choice.message
return convert_from_openai_message(message, extra_args)
diff --git
a/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_response_parsing.py
b/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_response_parsing.py
new file mode 100644
index 000000000..e31f72bb1
--- /dev/null
+++
b/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_response_parsing.py
@@ -0,0 +1,140 @@
+################################################################################
+# 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.
+#################################################################################
+from unittest.mock import MagicMock
+
+from openai.types.chat import ChatCompletion
+
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.integrations.chat_models.azure.azure_openai_chat_model
import (
+ AzureOpenAIChatModelConnection,
+)
+
+DEPLOYMENT = "my-deployment"
+
+ASSISTANT_MESSAGE = {
+ "role": "assistant",
+ "content": "ok",
+ "tool_calls": None,
+ "refusal": None,
+}
+USAGE = {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}
+
+OMITTED = object()
+"""Sentinel for a response whose choice carries no finish_reason field at all,
+which is distinct from one carrying an explicit null."""
+
+
+def _connection(
+ finish_reason: object, usage: dict | None = None
+) -> AzureOpenAIChatModelConnection:
+ """Build a connection whose stubbed transport returns one real completion.
+
+ The transport is a mock but the payload is a genuine ``ChatCompletion``, so
+ the assertions exercise the SDK's own attribute access rather than values a
+ mock was told to return.
+
+ Parameters
+ ----------
+ finish_reason : object
+ Value for the choice's ``finish_reason``. Pass ``OMITTED`` to leave the
+ field out of the payload entirely. Required, so a caller cannot omit it
+ by accident and assert against an unintended shape.
+ usage : dict | None
+ Token usage block, or None to build a response carrying none.
+ """
+ choice: dict = {"index": 0, "message": ASSISTANT_MESSAGE}
+ if finish_reason is not OMITTED:
+ choice["finish_reason"] = finish_reason
+ payload: dict = {
+ "id": "chatcmpl-test",
+ "object": "chat.completion",
+ "created": 0,
+ "model": DEPLOYMENT,
+ "choices": [choice],
+ }
+ if usage is not None:
+ payload["usage"] = usage
+
+ conn = AzureOpenAIChatModelConnection(
+ api_key="test-key",
+ azure_endpoint="https://example.openai.azure.com";,
+ api_version="2024-08-01-preview",
+ )
+ mock_client = MagicMock()
+ mock_client.chat.completions.create.return_value =
ChatCompletion.construct(
+ **payload
+ )
+ conn._client = mock_client
+ return conn
+
+
+def _chat(conn: AzureOpenAIChatModelConnection, **kwargs: object) ->
ChatMessage:
+ return conn.chat(
+ [ChatMessage(role=MessageRole.USER, content="hi")], model=DEPLOYMENT,
**kwargs
+ )
+
+
+def test_chat_records_finish_reason_in_extra_args() -> None:
+ """The finish reason survives alongside the token metrics."""
+ result = _chat(
+ _connection("length", usage=USAGE), model_of_azure_deployment="gpt-4o"
+ )
+
+ assert result.extra_args["promptTokens"] == 1
+ assert result.extra_args["finish_reason"] == "length"
+
+
+def test_chat_records_finish_reason_without_usage_or_deployment_model() ->
None:
+ """The finish reason is captured independently of the token metrics.
+
+ The metrics branch needs both model_of_azure_deployment and a usage block;
+ neither is supplied here. promptTokens is asserted absent to confirm that
+ branch did not run, so the finish reason cannot have been recorded by it.
+ """
+ result = _chat(_connection("tool_calls"))
+
+ assert "promptTokens" not in result.extra_args
+ assert result.extra_args["finish_reason"] == "tool_calls"
+
+
+def test_chat_records_unrecognized_finish_reason_verbatim() -> None:
+ """A finish reason outside the documented set is stored as received."""
+ result = _chat(
+ _connection("some_vendor_reason", usage=USAGE),
+ model_of_azure_deployment="gpt-4o",
+ )
+
+ assert result.extra_args["finish_reason"] == "some_vendor_reason"
+
+
+def test_chat_records_empty_finish_reason() -> None:
+ """An empty finish reason is recorded rather than discarded."""
+ # The capture turns on the value being present, not on it being non-empty,
+ # so an empty reason reaches extra_args like any other string.
+ result = _chat(_connection("", usage=USAGE),
model_of_azure_deployment="gpt-4o")
+
+ assert result.extra_args["finish_reason"] == ""
+
+
+def test_chat_omits_finish_reason_when_response_has_none() -> None:
+ """A response whose choice carries no finish reason yields no key."""
+ result = _chat(
+ _connection(OMITTED, usage=USAGE), model_of_azure_deployment="gpt-4o"
+ )
+
+ assert "finish_reason" not in result.extra_args
diff --git
a/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py
b/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py
index 587418dde..2eb1bacbc 100644
--- a/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py
+++ b/python/flink_agents/integrations/chat_models/openai/openai_chat_model.py
@@ -259,7 +259,8 @@ class OpenAIChatModelConnection(BaseChatModelConnection):
Returns:
-------
ChatMessage
- Model response message
+ Model response message. When the response carries a finish reason,
+ it is available as ``extra_args["finish_reason"]``.
"""
tool_specs = None
if tools is not None:
@@ -297,7 +298,11 @@ class OpenAIChatModelConnection(BaseChatModelConnection):
extra_args["promptTokens"] = response.usage.prompt_tokens
extra_args["completionTokens"] = response.usage.completion_tokens
- message = response.choices[0].message
+ choice = response.choices[0]
+ if choice.finish_reason is not None:
+ extra_args["finish_reason"] = choice.finish_reason
+
+ message = choice.message
return convert_from_openai_message(message, extra_args)
diff --git
a/python/flink_agents/integrations/chat_models/openai/openai_utils.py
b/python/flink_agents/integrations/chat_models/openai/openai_utils.py
index 9601d9ebc..0eec2b11b 100644
--- a/python/flink_agents/integrations/chat_models/openai/openai_utils.py
+++ b/python/flink_agents/integrations/chat_models/openai/openai_utils.py
@@ -136,6 +136,11 @@ def convert_to_openai_message(message: ChatMessage) ->
ChatCompletionMessagePara
- ASSISTANT role with tool_calls -> ChatCompletionAssistantMessageParam
- USER role -> ChatCompletionUserMessageParam
- SYSTEM role -> ChatCompletionSystemMessageParam
+
+ Only the fields OpenAI defines for each role are sent. Entries in
+ extra_args are not forwarded as message fields, except that a tool
+ message takes its tool_call_id from extra_args["external_id"], and an
+ assistant message carries extra_args["refusal"] when that value is a str.
"""
role = message.role
@@ -145,7 +150,6 @@ def convert_to_openai_message(message: ChatMessage) ->
ChatCompletionMessagePara
"role": "system",
"content": message.content,
}
- system_message.update(message.extra_args)
return system_message
# Handle USER role messages
@@ -154,7 +158,6 @@ def convert_to_openai_message(message: ChatMessage) ->
ChatCompletionMessagePara
"role": "user",
"content": message.content,
}
- user_message.update(message.extra_args)
return user_message
# Handle ASSISTANT role messages
@@ -172,7 +175,9 @@ def convert_to_openai_message(message: ChatMessage) ->
ChatCompletionMessagePara
]
assistant_message["tool_calls"] = openai_tool_calls
- assistant_message.update(message.extra_args)
+ refusal = message.extra_args.get("refusal")
+ if isinstance(refusal, str):
+ assistant_message["refusal"] = refusal
return assistant_message
# Handle TOOL role messages
diff --git
a/python/flink_agents/integrations/chat_models/openai/tests/test_openai_response_parsing.py
b/python/flink_agents/integrations/chat_models/openai/tests/test_openai_response_parsing.py
index 07c5e3d91..7ec8f384c 100644
---
a/python/flink_agents/integrations/chat_models/openai/tests/test_openai_response_parsing.py
+++
b/python/flink_agents/integrations/chat_models/openai/tests/test_openai_response_parsing.py
@@ -15,11 +15,18 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################
+from unittest.mock import MagicMock
+
import pytest
-from openai.types.chat import ChatCompletionMessage
+from openai.types.chat import ChatCompletion, ChatCompletionMessage
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.integrations.chat_models.openai.openai_chat_model import (
+ OpenAIChatModelConnection,
+)
from flink_agents.integrations.chat_models.openai.openai_utils import (
convert_from_openai_message,
+ convert_to_openai_message,
)
@@ -43,10 +50,200 @@ def test_refusal_is_preserved_in_extra_args(refusal: str)
-> None:
def test_no_refusal_key_when_refusal_absent() -> None:
"""A response that was not refused leaves no refusal key behind."""
- # extra_args is merged back into the outbound assistant message, so a null
- # refusal key here would be echoed to the provider on every later request.
+ # Absence of the key, not a falsy value, is what marks a response that was
+ # never refused.
message = ChatCompletionMessage(role="assistant", content="ok",
refusal=None)
result = convert_from_openai_message(message, {})
assert "refusal" not in result.extra_args
+
+
[email protected](
+ "role", [MessageRole.SYSTEM, MessageRole.USER, MessageRole.ASSISTANT]
+)
+def test_convert_to_openai_message_omits_response_metadata(
+ role: MessageRole,
+) -> None:
+ """Completion metadata held in extra_args never reaches an outbound
param."""
+ message = ChatMessage(
+ role=role,
+ content="hello",
+ extra_args={
+ "model_name": "gpt-4o",
+ "promptTokens": 3,
+ "completionTokens": 5,
+ },
+ )
+
+ param = convert_to_openai_message(message)
+
+ assert param == {"role": role.value, "content": "hello"}
+
+
[email protected]("refusal", ["I cannot help with that", ""])
+def test_convert_to_openai_message_forwards_string_refusal(refusal: str) ->
None:
+ """A string refusal on an assistant message is sent to the provider."""
+ # The outbound guard is a type check rather than a truthiness check, so an
+ # empty reason forwards like any other.
+ message = ChatMessage(
+ role=MessageRole.ASSISTANT,
+ content="",
+ extra_args={"refusal": refusal},
+ )
+
+ param = convert_to_openai_message(message)
+
+ assert param == {"role": "assistant", "content": "", "refusal": refusal}
+
+
[email protected]("refusal", [{"reason": "policy"}, 123, True])
+def test_convert_to_openai_message_omits_non_string_refusal(
+ refusal: object,
+) -> None:
+ """Only a string refusal is forwarded; a value of any other type is
dropped."""
+ message = ChatMessage(
+ role=MessageRole.ASSISTANT,
+ content="",
+ extra_args={"refusal": refusal},
+ )
+
+ param = convert_to_openai_message(message)
+
+ assert param == {"role": "assistant", "content": ""}
+
+
+def test_convert_to_openai_message_assistant_tool_calls() -> None:
+ """An assistant message requesting tool calls sends them with a null
content."""
+ message = ChatMessage(
+ role=MessageRole.ASSISTANT,
+ content="",
+ tool_calls=[
+ {
+ "original_id": "call_abc",
+ "function": {"name": "get_weather", "arguments": {"city":
"Berlin"}},
+ }
+ ],
+ extra_args={"model_name": "gpt-4o", "promptTokens": 3},
+ )
+
+ param = convert_to_openai_message(message)
+
+ assert set(param) == {"role", "content", "tool_calls"}
+ assert param["role"] == "assistant"
+ assert param["content"] is None
+
+
+def test_convert_to_openai_message_tool_role_unchanged() -> None:
+ """A tool result carries its call id and nothing else from extra_args."""
+ message = ChatMessage(
+ role=MessageRole.TOOL,
+ content="42",
+ extra_args={"external_id": "call_abc", "promptTokens": 7},
+ )
+
+ param = convert_to_openai_message(message)
+
+ assert param == {"role": "tool", "content": "42", "tool_call_id":
"call_abc"}
+
+
+ASSISTANT_MESSAGE = {
+ "role": "assistant",
+ "content": "ok",
+ "tool_calls": None,
+ "refusal": None,
+}
+USAGE = {"prompt_tokens": 1, "completion_tokens": 2, "total_tokens": 3}
+
+OMITTED = object()
+"""Sentinel for a response whose choice carries no finish_reason field at all,
+which is distinct from one carrying an explicit null."""
+
+
+def _connection(
+ finish_reason: object, usage: dict | None = None
+) -> OpenAIChatModelConnection:
+ """Build a connection whose stubbed transport returns one real completion.
+
+ The transport is a mock but the payload is a genuine ``ChatCompletion``, so
+ the assertions exercise the SDK's own attribute access rather than values a
+ mock was told to return.
+
+ Parameters
+ ----------
+ finish_reason : object
+ Value for the choice's ``finish_reason``. Pass ``OMITTED`` to leave the
+ field out of the payload entirely. Required, so a caller cannot omit it
+ by accident and assert against an unintended shape.
+ usage : dict | None
+ Token usage block, or None to build a response carrying none.
+ """
+ choice: dict = {"index": 0, "message": ASSISTANT_MESSAGE}
+ if finish_reason is not OMITTED:
+ choice["finish_reason"] = finish_reason
+ payload: dict = {
+ "id": "chatcmpl-test",
+ "object": "chat.completion",
+ "created": 0,
+ "model": "gpt-4o",
+ "choices": [choice],
+ }
+ if usage is not None:
+ payload["usage"] = usage
+
+ conn = OpenAIChatModelConnection(
+ api_key="test-key", api_base_url="http://localhost";
+ )
+ mock_client = MagicMock()
+ mock_client.chat.completions.create.return_value =
ChatCompletion.construct(
+ **payload
+ )
+ conn._client = mock_client
+ return conn
+
+
+def _chat(conn: OpenAIChatModelConnection) -> ChatMessage:
+ return conn.chat([ChatMessage(role=MessageRole.USER, content="hi")],
model="gpt-4o")
+
+
+def test_chat_records_finish_reason_in_extra_args() -> None:
+ """The finish reason survives alongside the token metrics."""
+ result = _chat(_connection("length", usage=USAGE))
+
+ assert result.extra_args["promptTokens"] == 1
+ assert result.extra_args["finish_reason"] == "length"
+
+
+def test_chat_records_finish_reason_when_usage_is_absent() -> None:
+ """The finish reason is captured independently of the token metrics.
+
+ promptTokens is asserted absent to confirm the metrics branch did not run,
+ so the finish reason cannot have been recorded by it.
+ """
+ result = _chat(_connection("tool_calls"))
+
+ assert "promptTokens" not in result.extra_args
+ assert result.extra_args["finish_reason"] == "tool_calls"
+
+
+def test_chat_records_unrecognized_finish_reason_verbatim() -> None:
+ """A finish reason outside the documented set is stored as received."""
+ result = _chat(_connection("some_vendor_reason", usage=USAGE))
+
+ assert result.extra_args["finish_reason"] == "some_vendor_reason"
+
+
+def test_chat_records_empty_finish_reason() -> None:
+ """An empty finish reason is recorded rather than discarded."""
+ # The capture turns on the value being present, not on it being non-empty,
+ # so an empty reason reaches extra_args like any other string.
+ result = _chat(_connection("", usage=USAGE))
+
+ assert result.extra_args["finish_reason"] == ""
+
+
+def test_chat_omits_finish_reason_when_response_has_none() -> None:
+ """A response whose choice carries no finish reason yields no key."""
+ result = _chat(_connection(OMITTED, usage=USAGE))
+
+ assert "finish_reason" not in result.extra_args
diff --git a/python/flink_agents/plan/actions/chat_model_action.py
b/python/flink_agents/plan/actions/chat_model_action.py
index 0b6e7f4ae..b802fda3b 100644
--- a/python/flink_agents/plan/actions/chat_model_action.py
+++ b/python/flink_agents/plan/actions/chat_model_action.py
@@ -57,6 +57,9 @@ _TOOL_CALL_CONTEXT = "_TOOL_CALL_CONTEXT"
_TOOL_REQUEST_EVENT_CONTEXT = "_TOOL_REQUEST_EVENT_CONTEXT"
_RETRY_STATS_CONTEXT = "_RETRY_STATS_CONTEXT"
_PROMPT_ARGS = "prompt_args"
+_FINISH_REASON = "finish_reason"
+_TRUNCATED_FINISH_REASON = "length"
+_CONTENT_FILTERED_FINISH_REASON = "content_filter"
_logger = logging.getLogger(__name__)
@@ -272,6 +275,38 @@ def _generate_structured_output(
return response
+def _reject_incomplete_response(response: ChatMessage) -> None:
+ """Reject a response the provider did not finish emitting.
+
+ Evaluated once per chat response, before it is dispatched as text,
+ structured output, or tool calls.
+
+ Args:
+ response: The chat response whose ``finish_reason`` is inspected. Any
+ other reason, and an absent one, are accepted.
+
+ Raises:
+ ValueError: If the finish reason reports the response as cut off by the
+ token budget or withheld by content filtering.
+ """
+ finish_reason = response.extra_args.get(_FINISH_REASON)
+ if finish_reason == _TRUNCATED_FINISH_REASON:
+ error_message = (
+ f"ChatModel response is truncated
(finish_reason={finish_reason!r}): "
+ "it exhausted the completion token budget before the model
finished, "
+ "so the content is incomplete. Raise the model's max output
tokens, "
+ "or ask for a smaller output."
+ )
+ raise ValueError(error_message)
+ if finish_reason == _CONTENT_FILTERED_FINISH_REASON:
+ error_message = (
+ "ChatModel response was withheld by the provider's content filter "
+ f"(finish_reason={finish_reason!r}), so the content is incomplete.
"
+ "Adjust the prompt or the provider's content filtering
configuration."
+ )
+ raise ValueError(error_message)
+
+
def _generate_structured_output_with_report(
ctx: RunnerContext, response: ChatMessage, output_schema: OutputSchema
) -> ChatMessage:
@@ -393,6 +428,10 @@ async def chat(
response.extra_args["completionTokens"],
request_metric_group,
)
+ # A truncated response consumed its full token budget, so the token
+ # metrics above are recorded before this rejects and abandons the
+ # response.
+ _reject_incomplete_response(response)
if output_schema is not None and len(response.tool_calls) == 0:
response = _generate_structured_output_with_report(
ctx, response, output_schema
diff --git a/python/flink_agents/plan/tests/actions/test_chat_model_action.py
b/python/flink_agents/plan/tests/actions/test_chat_model_action.py
index f0bba58e4..42c42b3e5 100644
--- a/python/flink_agents/plan/tests/actions/test_chat_model_action.py
+++ b/python/flink_agents/plan/tests/actions/test_chat_model_action.py
@@ -15,6 +15,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#################################################################################
+from unittest.mock import MagicMock
from uuid import uuid4
from pydantic import BaseModel
@@ -23,10 +24,12 @@ from pyflink.common.typeinfo import BasicTypeInfo,
RowTypeInfo
from flink_agents.api.agents.react_agent import OutputSchema
from flink_agents.api.chat_message import ChatMessage, MessageRole
from flink_agents.api.memory_object import MemoryType
+from flink_agents.api.trace import ExecutionReporter
from flink_agents.plan.actions.chat_model_action import (
_TOOL_CALL_CONTEXT,
_TOOL_REQUEST_EVENT_CONTEXT,
_clean_llm_response,
+ _generate_structured_output_with_report,
_get_tool_request_event_context,
_save_tool_request_event_context,
_update_tool_call_context,
@@ -207,3 +210,30 @@ def test_save_get_preserves_model_and_prompt_args():
context = _get_tool_request_event_context(mem, event_id)
assert context["model"] == "ollama"
assert context["prompt_args"] == prompt_args
+
+
+_PARSEABLE_CONTENT = '{"result": 42}'
+
+
+def _response(extra_args) -> ChatMessage:
+ return ChatMessage(
+ role=MessageRole.ASSISTANT,
+ content=_PARSEABLE_CONTENT,
+ extra_args=extra_args,
+ )
+
+
+def _parse(ctx, extra_args) -> ChatMessage:
+ return _generate_structured_output_with_report(
+ ctx, _response(extra_args), OutputSchema(output_schema=_Result)
+ )
+
+
+def test_accepted_finish_reason_reports_parser_execution():
+ ctx = MagicMock(spec=ExecutionReporter)
+
+ _parse(ctx, {"finish_reason": "stop"})
+
+ ctx.report_execution_started.assert_called_once()
+ ctx.report_execution_succeeded.assert_called_once()
+ ctx.report_execution_failed.assert_not_called()
diff --git
a/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py
b/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py
index f73648aa8..a5cad1a34 100644
--- a/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py
+++ b/python/flink_agents/plan/tests/actions/test_chat_model_action_retry.py
@@ -34,7 +34,7 @@ from flink_agents.api.core_options import (
ErrorHandlingStrategy,
)
from flink_agents.api.events.chat_event import ChatResponseEvent
-from flink_agents.api.events.tool_event import ToolResponseEvent
+from flink_agents.api.events.tool_event import ToolRequestEvent,
ToolResponseEvent
from flink_agents.api.metric_group import Counter, MetricGroup
from flink_agents.api.trace import (
ExecutionEntityTypes,
@@ -119,6 +119,7 @@ def _create_mock_runner_context(
chat_model: Any,
max_retries: int = 3,
retry_wait_interval_sec: int = 1,
+ error_handling_strategy: ErrorHandlingStrategy =
ErrorHandlingStrategy.RETRY,
) -> tuple[MagicMock, list, _MockMetricGroup, _MockMemoryObject]:
"""Create a mock RunnerContext with configurable retry settings.
@@ -131,7 +132,7 @@ def _create_mock_runner_context(
config = MagicMock()
option_values = {
- id(AgentExecutionOptions.ERROR_HANDLING_STRATEGY):
ErrorHandlingStrategy.RETRY,
+ id(AgentExecutionOptions.ERROR_HANDLING_STRATEGY):
error_handling_strategy,
id(AgentExecutionOptions.MAX_RETRIES): max_retries,
id(AgentExecutionOptions.RETRY_WAIT_INTERVAL): retry_wait_interval_sec,
id(AgentExecutionOptions.CHAT_ASYNC): False,
@@ -343,6 +344,204 @@ class TestChatModelActionRetry:
)
+class TestChatModelActionFinishReason:
+ """Tests for the finish-reason gate on the common chat-response path."""
+
+ def _run(self, ctx, output_schema=None) -> None:
+ asyncio.run(
+ chat(
+ uuid4(),
+ "test-model",
+ [ChatMessage(role=MessageRole.USER, content="hi")],
+ {},
+ output_schema,
+ ctx,
+ )
+ )
+
+ def test_truncated_text_response_rejected(self) -> None:
+ chat_model = MagicMock()
+ chat_model.chat = MagicMock(
+ return_value=ChatMessage(
+ role=MessageRole.ASSISTANT,
+ content="partial answ",
+ extra_args={"finish_reason": "length"},
+ )
+ )
+ ctx, sent_events, _, _ = _create_mock_runner_context(
+ chat_model, max_retries=0, retry_wait_interval_sec=0
+ )
+
+ with pytest.raises(ValueError, match="(?i)truncat") as exc_info:
+ self._run(ctx)
+
+ assert "token" in str(exc_info.value).lower()
+ assert len(sent_events) == 0
+
+ def test_content_filtered_text_response_rejected(self) -> None:
+ # Matches a word unique to the filtering message. Both messages
+ # interpolate the finish reason, so "content_filter" appears in either
+ # one and cannot tell them apart.
+ chat_model = MagicMock()
+ chat_model.chat = MagicMock(
+ return_value=ChatMessage(
+ role=MessageRole.ASSISTANT,
+ content="",
+ extra_args={"finish_reason": "content_filter"},
+ )
+ )
+ ctx, sent_events, _, _ = _create_mock_runner_context(
+ chat_model, max_retries=0, retry_wait_interval_sec=0
+ )
+
+ with pytest.raises(ValueError, match="(?i)withheld"):
+ self._run(ctx)
+
+ assert len(sent_events) == 0
+
+ def test_truncated_tool_call_response_rejected_before_tool_dispatch(self)
-> None:
+ chat_model = MagicMock()
+ chat_model.chat = MagicMock(
+ return_value=ChatMessage(
+ role=MessageRole.ASSISTANT,
+ content="",
+ tool_calls=[
+ {
+ "id": "call-1",
+ "function": {"name": "f", "arguments": ""},
+ }
+ ],
+ extra_args={"finish_reason": "length"},
+ )
+ )
+ ctx, sent_events, _, _ = _create_mock_runner_context(
+ chat_model, max_retries=0, retry_wait_interval_sec=0
+ )
+
+ with pytest.raises(ValueError, match="(?i)truncat"):
+ self._run(ctx)
+
+ # A truncated tool call carries arguments the model never finished
+ # writing, so no ToolRequestEvent may leave the action.
+ assert len(sent_events) == 0
+
+ @pytest.mark.parametrize(
+ "extra_args",
+ [
+ {"finish_reason": "stop"},
+ {"finish_reason": "tool_calls"},
+ {"finish_reason": "some_vendor_reason"},
+ {},
+ ],
+ ids=["stop", "tool_calls", "unrecognized", "absent"],
+ )
+ def test_accepted_finish_reason_reaches_the_response_event(
+ self, extra_args: dict
+ ) -> None:
+ chat_model = MagicMock()
+ chat_model.chat = MagicMock(
+ return_value=ChatMessage(
+ role=MessageRole.ASSISTANT,
+ content="hello",
+ extra_args=extra_args,
+ )
+ )
+ ctx, sent_events, _, _ = _create_mock_runner_context(
+ chat_model, max_retries=0, retry_wait_interval_sec=0
+ )
+
+ self._run(ctx)
+
+ assert len(sent_events) == 1
+ assert isinstance(sent_events[0], ChatResponseEvent)
+ assert sent_events[0].response.content == "hello"
+
+ def test_accepted_finish_reason_dispatches_tool_request_event(self) ->
None:
+ # A response carrying tool calls passes the same finish-reason gate as
a
+ # text response, so an accepted reason must reach tool dispatch.
+ tool_calls = [{"id": "call-1", "function": {"name": "f", "arguments":
{}}}]
+ chat_model = MagicMock()
+ chat_model.chat = MagicMock(
+ return_value=ChatMessage(
+ role=MessageRole.ASSISTANT,
+ content="",
+ tool_calls=tool_calls,
+ extra_args={"finish_reason": "tool_calls"},
+ )
+ )
+ ctx, sent_events, _, _ = _create_mock_runner_context(
+ chat_model, max_retries=0, retry_wait_interval_sec=0
+ )
+
+ self._run(ctx)
+
+ assert len(sent_events) == 1
+ assert isinstance(sent_events[0], ToolRequestEvent)
+ assert sent_events[0].tool_calls == tool_calls
+
+ def test_ignore_strategy_drops_rejected_response_without_event(self) ->
None:
+ # Under IGNORE the record is dropped: the rejection does not propagate
+ # and no event carries the truncated content downstream.
+ chat_model = MagicMock()
+ chat_model.chat = MagicMock(
+ return_value=ChatMessage(
+ role=MessageRole.ASSISTANT,
+ content="partial answ",
+ extra_args={"finish_reason": "length"},
+ )
+ )
+ ctx, sent_events, _, _ = _create_mock_runner_context(
+ chat_model,
+ max_retries=0,
+ retry_wait_interval_sec=0,
+ error_handling_strategy=ErrorHandlingStrategy.IGNORE,
+ )
+
+ self._run(ctx)
+
+ assert len(sent_events) == 0
+
+ @pytest.mark.parametrize("finish_reason", ["length", "content_filter"])
+ def test_rejected_finish_reason_skips_structured_output(
+ self, finish_reason: str
+ ) -> None:
+ chat_model = MagicMock()
+ chat_model.chat = MagicMock(
+ return_value=ChatMessage(
+ role=MessageRole.ASSISTANT,
+ content='{"result": 42}',
+ extra_args={
+ "finish_reason": finish_reason,
+ "model_name": "provider-model",
+ "promptTokens": 100,
+ "completionTokens": 50,
+ },
+ )
+ )
+ ctx, sent_events, metric_group, _ = _create_mock_runner_context(
+ chat_model, max_retries=0, retry_wait_interval_sec=0
+ )
+
+ with pytest.raises(ValueError):
+ self._run(ctx, OutputSchema(output_schema=_StructuredResult))
+
+ # The model call itself succeeded and spent its full token budget, so
+ # both must be recorded before the response is rejected.
+ ctx.report_execution_succeeded.assert_called_once_with(
+ ExecutionEntityTypes.LLM, "test-model", _LLM_METADATA
+ )
+ chat_model._record_token_metrics.assert_called_once_with(
+ "provider-model", 100, 50, metric_group
+ )
+ # The parse is never attempted, so nothing about it is reported and no
+ # response leaves the action.
+ ctx.report_execution_started.assert_called_once_with(
+ ExecutionEntityTypes.LLM, "test-model", _LLM_METADATA
+ )
+ ctx.report_execution_failed.assert_not_called()
+ assert len(sent_events) == 0
+
+
class TestChatResponseEventRetryFields:
"""Tests for ChatResponseEvent retry fields."""
