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 db3270db0e..4ce122ad04 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 @@ 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 messages, List tools, Map 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. * + *

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 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() 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 3af78b72c4..4864769e55 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 @@ 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 messages, List tools, Map 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 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()) { 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 6b7b2d3b1c..c15f553df9 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 @@ -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 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 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 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 finishReason) { + return completionBuilder(finishReason) + .usage( + CompletionUsage.builder() + .promptTokens(11L) + .completionTokens(7L) + .totalTokens(18L) + .build()) + .build(); + } + + private static ChatCompletion completionWithFinishReasonWithoutUsage( + JsonField 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 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) 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 0000000000..6e1e633b93 --- /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 269dee2966..c5775dc7b7 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 @@ -39,8 +39,10 @@ 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; /** @@ -147,6 +149,98 @@ void testProviderErrorPropagatesUnwrapped() throws IOException { } } + @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 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 57c191c75a..2e8c2476d6 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 @@ static ChatMessage generateStructuredOutput(ChatMessage response, Object outputS throw new RuntimeException( String.format("Unsupported output schema %s.", outputSchema)); } - Map extraArgs = new HashMap<>(); + Map extraArgs = new HashMap<>(response.getExtraArgs()); extraArgs.put(STRUCTURED_OUTPUT, structuredOutput); return new ChatMessage(response.getRole(), output, extraArgs); } @@ -576,6 +579,34 @@ private static Map takeRoutingMetadata( 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 8dfacee4e0..4f49feac7f 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 @@ public ChatMessage call() throws Exception { 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 a1d7b0cf31..450ef22dc1 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.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.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.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 LLM_METADATA = @@ -425,6 +432,205 @@ void processToolResponseForwardsSavedArgumentsToChat() throws Exception { 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> 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 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 46485b5a12..cfdffd03eb 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.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 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 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 @@ void testCleanLlmResponseWithMultipleLinesInBlock() { 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 extraArgs = new HashMap<>(); + extraArgs.put("finish_reason", "stop"); + extraArgs.put("promptTokens", 100L); + + Map 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 50884bf681..88c05f378b 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 @@ -264,7 +264,8 @@ def chat( 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: @@ -336,7 +337,11 @@ def chat( 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 0000000000..e31f72bb12 --- /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 e810613f78..8019e0cd68 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 @@ -253,7 +253,8 @@ def chat( 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: @@ -291,7 +292,11 @@ def chat( 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 9601d9ebc1..0eec2b11b3 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 07c5e3d915..7ec8f384cd 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 + + +@pytest.mark.parametrize( + "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"} + + +@pytest.mark.parametrize("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} + + +@pytest.mark.parametrize("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 0b6e7f4ae7..b802fda3ba 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_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 f0bba58e4c..42c42b3e5b 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 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 f73648aa8e..a5cad1a340 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 @@ 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 @@ def test_structured_output_parse_error_retries_without_failing_llm( ) +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."""