From 2aef0fc148a93d3656f5e438db7b1973942bcb43 Mon Sep 17 00:00:00 2001 From: weiqingy Date: Tue, 18 Aug 2026 00:23:50 -0700 Subject: [PATCH 1/7] [python][integrations] Allowlist outbound OpenAI message fields convert_to_openai_message merged the whole extra_args dict into outbound system, user and assistant message params. The inbound path writes completion metadata into that same dict, and chat_model_action replays the stored assistant message on the next turn of a tool loop, so model_name, promptTokens, completionTokens and structured_output were sent back to the provider as top-level message fields. structured_output holds a Pydantic model or a pyflink Row, which is not serialisable as a message field at all. Send only the fields OpenAI defines for each role, matching what Java's OpenAIChatCompletionsUtils.convertToOpenAIMessage already produces: content alone for system and user; content, tool calls and a string refusal for assistant; the tool branch unchanged. The isinstance guard on refusal mirrors Java's instanceof String on the same outbound path. Azure OpenAI and vLLM share this converter and pick up the same fix. The function had no test coverage before. The new tests were driven by mutation testing and kill every mutant in scope, including reinstating the blanket merge. Generated-by: Claude Code 2.1.234 (Claude Opus 5) --- .../chat_models/openai/openai_utils.py | 11 ++- .../tests/test_openai_response_parsing.py | 94 ++++++++++++++++++- 2 files changed, 100 insertions(+), 5 deletions(-) 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..9620331374 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 @@ -18,8 +18,10 @@ import pytest from openai.types.chat import ChatCompletionMessage +from flink_agents.api.chat_message import ChatMessage, MessageRole from flink_agents.integrations.chat_models.openai.openai_utils import ( convert_from_openai_message, + convert_to_openai_message, ) @@ -43,10 +45,98 @@ 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"} From db007d28a4d256b7825035327e5b8fbe98f3648a Mon Sep 17 00:00:00 2001 From: weiqingy Date: Tue, 18 Aug 2026 02:03:36 -0700 Subject: [PATCH 2/7] [python][integrations] Capture finish_reason on OpenAI and Azure responses The provider's finish reason was never recorded, so a completion truncated by a token limit came back as ordinary partial content with nothing marking it incomplete. Record it on the returned message as extra_args["finish_reason"] in both OpenAI-family connections. Nothing reads it yet; acting on it under structured output is a separate change. The capture sits outside the token-metrics guard in both connections. A response can carry a finish reason with no usage block, and Azure's guard additionally depends on a deployment name unrelated to the reason, so writing it inside either guard would drop it silently. The value is stored verbatim, including reasons outside OpenAI's documented set, and an absent or null reason writes no key and raises nothing. OpenAI-compatible servers omit the field and emit vendor-specific values, and VLLMChatModelConnection subclasses the OpenAI connection, so this path serves them too. Tests mock only the transport and drive real SDK response objects. Choice declares finish_reason as a required Literal, so strict validation would reject an unknown vendor value before this code ran; the client itself defaults to lenient response construction, which is what the tests use. Generated-by: Claude Code 2.1.234 (Claude Opus 5) --- .../azure/azure_openai_chat_model.py | 9 +- .../test_azure_openai_response_parsing.py | 141 ++++++++++++++++++ .../chat_models/openai/openai_chat_model.py | 9 +- .../tests/test_openai_response_parsing.py | 109 +++++++++++++- 4 files changed, 263 insertions(+), 5 deletions(-) create mode 100644 python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_response_parsing.py 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..2707d5ccf8 --- /dev/null +++ b/python/flink_agents/integrations/chat_models/azure/tests/test_azure_openai_response_parsing.py @@ -0,0 +1,141 @@ +################################################################################ +# 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( + name="azure_openai", + 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/tests/test_openai_response_parsing.py b/python/flink_agents/integrations/chat_models/openai/tests/test_openai_response_parsing.py index 9620331374..10a1a67365 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,10 +15,15 @@ # 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, @@ -140,3 +145,105 @@ def test_convert_to_openai_message_tool_role_unchanged() -> None: 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( + name="openai", 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 From c997549bd1e7694eae9709a1ed873bd51e2298cc Mon Sep 17 00:00:00 2001 From: weiqingy Date: Sat, 22 Aug 2026 15:49:08 -0700 Subject: [PATCH 3/7] [python][plan] Fail structured output on a truncated or filtered response A response truncated by a token limit was parsed like any other. When the partial content was invalid JSON the caller got a syntax error pointing at the parser rather than at the truncation, and when it happened to still parse, the caller silently received wrong structured output. Check why the model stopped before parsing against an output schema. A length or content_filter reason raises with the cause named; stop and tool_calls parse as before; an unrecognized reason or an absent key also parses, so providers that never record one are unaffected. The check runs before the parser execution is reported, not inside the parse. The reporting wrapper records any exception from the parse as model_output_parse_error, so gating inside it would report a truncation as a parse failure, which is the same misleading diagnosis this change removes. Rejecting beforehand records no parser execution at all, which is accurate because none was attempted. ValueError keeps callers that already catch the JSONDecodeError symptom working, and adds no public API. Truncated content that still parses previously produced wrong output and now raises, so callers using the RETRY strategy will see such responses retried. Generated-by: Claude Code 2.1.234 (Claude Opus 5) --- .../plan/actions/chat_model_action.py | 37 +++++++++ .../tests/actions/test_chat_model_action.py | 76 +++++++++++++++++++ 2 files changed, 113 insertions(+) diff --git a/python/flink_agents/plan/actions/chat_model_action.py b/python/flink_agents/plan/actions/chat_model_action.py index 0b6e7f4ae7..3642fcb6f8 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,9 +275,43 @@ def _generate_structured_output( return response +def _reject_incomplete_response(response: ChatMessage) -> None: + """Reject a response the provider did not finish emitting. + + 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 and cannot yield structured output. " + "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 " + "and cannot yield structured output. 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: + # Precedes the start report: when this rejects, parsing is never attempted, + # so there is no parser execution to report. + _reject_incomplete_response(response) + ExecutionReporters.started(ctx, ExecutionEntityTypes.PARSER, STRUCTURED_OUTPUT) try: structured_response = _generate_structured_output(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..1ac8f1e47d 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,18 +15,23 @@ # See the License for the specific language governing permissions and # limitations under the License. ################################################################################# +from unittest.mock import MagicMock from uuid import uuid4 +import pytest from pydantic import BaseModel from pyflink.common.typeinfo import BasicTypeInfo, RowTypeInfo +from flink_agents.api.agents.agent import STRUCTURED_OUTPUT 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 +212,74 @@ 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_length_finish_reason_rejects_parsing(): + with pytest.raises(ValueError, match="(?i)truncat") as exc_info: + _parse(MagicMock(spec=ExecutionReporter), {"finish_reason": "length"}) + assert "token" in str(exc_info.value).lower() + + +def test_content_filter_finish_reason_rejects_parsing(): + # 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. + with pytest.raises(ValueError, match="(?i)withheld"): + _parse(MagicMock(spec=ExecutionReporter), {"finish_reason": "content_filter"}) + + +@pytest.mark.parametrize("finish_reason", ["length", "content_filter"]) +def test_rejected_finish_reason_reports_no_parser_execution(finish_reason): + ctx = MagicMock(spec=ExecutionReporter) + + with pytest.raises(ValueError): + _parse(ctx, {"finish_reason": finish_reason}) + + ctx.report_execution_started.assert_not_called() + ctx.report_execution_succeeded.assert_not_called() + ctx.report_execution_failed.assert_not_called() + + +def test_accepted_finish_reason_reports_parser_execution(): + # Reports only reach a context that satisfies ExecutionReporter, so this + # also pins that the fixture still does. Without it, the assert_not_called + # checks above would hold for any context and stop testing the gate. + 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() + + +@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_parses_structured_output(extra_args): + response = _parse(MagicMock(spec=ExecutionReporter), extra_args) + assert response.extra_args[STRUCTURED_OUTPUT].result == 42 From a52790931d1132cbde225a7be5fb4a863343b904 Mon Sep 17 00:00:00 2001 From: weiqingy Date: Sat, 22 Aug 2026 16:47:27 -0700 Subject: [PATCH 4/7] [java][integrations] Capture finish_reason on OpenAI and Azure responses Mirror the Python capture on the Java side: record the provider's finish reason on the returned message as extraArgs["finish_reason"] in both OpenAI-family connections. Nothing reads it yet; acting on it under structured output is a separate change. Read the value through _finishReason().asKnown() rather than finishReason(). The latter resolves through JsonField.getRequired and throws OpenAIInvalidDataException when the member is absent or JSON-null. VLLMChatModelConnection extends OpenAICompletionsConnection, so OpenAI-compatible servers that omit the member use this path, and the plain accessor would fail a call that previously succeeded. FinishReason is an open value wrapper rather than an enum, so the wire string is taken with asString() and vendor-specific reasons are stored verbatim. known() throws on an unrecognized value and value() collapses it, so neither is used. The capture sits outside the token-usage guard in both connections. A response can carry a finish reason with no usage block, and the Azure guard additionally requires a deployment name unrelated to the reason. Azure is tested through its existing package-private seam. The OpenAI connection stashes inline in a private method, so its tests drive a loopback endpoint, which also covers a vendor reason surviving real deserialization. An omitted member and a null member are distinct wire shapes and both are covered. Generated-by: Claude Code 2.1.234 (Claude Opus 5) --- .../AzureOpenAIChatModelConnection.java | 20 +++- .../openai/OpenAICompletionsConnection.java | 21 +++- .../AzureOpenAIChatModelConnectionTest.java | 112 +++++++++++++++++- .../openai/FakeOpenAICompletionsEndpoint.java | 106 +++++++++++++++++ .../OpenAICompletionsConnectionTest.java | 94 +++++++++++++++ 5 files changed, 348 insertions(+), 5 deletions(-) create mode 100644 integrations/chat-models/openai/src/test/java/org/apache/flink/agents/integrations/chatmodels/openai/FakeOpenAICompletionsEndpoint.java 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() { From 9e72b51d9fd20f32ce1fce3f8d83aadfcabbdda5 Mon Sep 17 00:00:00 2001 From: weiqingy Date: Sat, 22 Aug 2026 18:11:13 -0700 Subject: [PATCH 5/7] [java][plan] Fail structured output on a truncated or filtered response Java mirror of the Python gate. Check why the model stopped before parsing against an output schema: a length or content_filter reason throws with the cause named, while stop, tool_calls, an unrecognized reason and an absent key all parse as before. The check runs before the parser execution is reported, not inside the parse. The reporting wrapper records any exception from the parse as model_output_parse_error, so gating inside it would report a truncation as a parse failure, which is the same misleading diagnosis this change removes. Rejecting beforehand records no parser execution at all, which is accurate because none was attempted. IllegalStateException is what plan already uses for this class of failure, and adds no public API. The two messages are kept individually identifiable so a truncation cannot be reported as content filtering or the reverse. Also seed the structured-output result from the inbound extra args instead of building a fresh map. The previous behavior discarded everything the connection recorded, so Java dropped finish_reason exactly when an output schema was in play, while Python preserved it. Token metrics are recorded before this point from the only call site, and the result cannot re-enter the parser because the path is guarded on tool calls being empty. Generated-by: Claude Code 2.1.234 (Claude Opus 5) --- .../agents/plan/actions/ChatModelAction.java | 37 ++++++- .../plan/actions/ChatModelActionTest.java | 104 ++++++++++++++++++ 2 files changed, 140 insertions(+), 1 deletion(-) 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..89270a5d8c 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,8 +579,40 @@ private static Map takeRoutingMetadata( return routing; } + /** + * Rejects a response the provider did not finish emitting. 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. + */ + private 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 and cannot yield" + + " structured output. 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" + + " and cannot yield structured output. Adjust the prompt or" + + " the provider's content filtering configuration.", + finishReason)); + } + } + static ChatMessage generateStructuredOutputWithReport( RunnerContext ctx, ChatMessage response, Object outputSchema) throws Exception { + // Precedes the start report: when this rejects, parsing is never attempted, so there is no + // parser execution to report. + rejectIncompleteResponse(response); + ExecutionReporters.started(ctx, ExecutionReporter.EntityTypes.PARSER, STRUCTURED_OUTPUT); try { ChatMessage structuredResponse = generateStructuredOutput(response, outputSchema); 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..19ee34ddd5 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,26 @@ */ 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 org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.MethodSource; +import org.junit.jupiter.params.provider.ValueSource; import java.util.HashMap; +import java.util.Locale; import java.util.Map; +import java.util.stream.Stream; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; +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 +44,37 @@ 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); + } + + 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()); + } + @Test void testRecordChatTokenMetricsRecordsWhenAllKeysPresent() { BaseChatModelSetup setup = mock(BaseChatModelSetup.class); @@ -179,4 +212,75 @@ void testCleanLlmResponseWithMultipleLinesInBlock() { String expected = "{\n \"key\": \"value\"\n}"; assertEquals(expected, ChatModelAction.cleanLlmResponse(input)); } + + @Test + void testStructuredOutputRejectsTruncatedResponse() { + IllegalStateException error = + assertThrows( + IllegalStateException.class, + () -> + generateStructuredOutput( + reportingContext(), Map.of("finish_reason", "length"))); + + String message = error.getMessage().toLowerCase(Locale.ROOT); + assertTrue(message.contains("truncat"), message); + assertTrue(message.contains("token"), message); + } + + @Test + void testStructuredOutputRejectsContentFilteredResponse() { + IllegalStateException error = + assertThrows( + IllegalStateException.class, + () -> + generateStructuredOutput( + reportingContext(), + 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. + String message = error.getMessage().toLowerCase(Locale.ROOT); + assertTrue(message.contains("withheld"), message); + assertTrue(message.contains("content filter"), message); + } + + @ParameterizedTest + @MethodSource("acceptedFinishReasons") + void testAcceptedFinishReasonParsesStructuredOutput(Map extraArgs) + throws Exception { + ChatMessage parsed = generateStructuredOutput(reportingContext(), extraArgs); + + assertEquals(Map.of("answer", "42"), parsed.getExtraArgs().get(Agent.STRUCTURED_OUTPUT)); + } + + @ParameterizedTest + @ValueSource(strings = {"length", "content_filter"}) + void testRejectedFinishReasonReportsNoParserExecution(String finishReason) throws Exception { + RunnerContext ctx = reportingContext(); + + assertThrows( + IllegalStateException.class, + () -> generateStructuredOutput(ctx, Map.of("finish_reason", finishReason))); + + ExecutionReporter reporter = (ExecutionReporter) ctx; + verify(reporter, never()).reportExecutionStarted(anyString(), anyString(), any()); + verify(reporter, never()).reportExecutionSucceeded(anyString(), anyString(), any()); + verify(reporter, never()) + .reportExecutionFailed(anyString(), anyString(), any(), any(), any()); + } + + @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()); + } } From e00ef82b0c11b0c8a27d4f0b3cc092953b165471 Mon Sep 17 00:00:00 2001 From: weiqingy Date: Sun, 30 Aug 2026 15:06:30 -0700 Subject: [PATCH 6/7] [plan] Reject a truncated or filtered response on the common chat path The finish_reason gate ran only when a response was parsed against an output schema, so a truncated text or tool-call response was still treated as successful. Move it to the common chat-response path in both languages, after the LLM execution is reported and its token metrics are recorded, and before the response is dispatched as text, structured output, or tool calls. The reasons that raise are unchanged: length and content_filter raise, while stop, tool_calls, any other value and an absent key proceed. Both messages drop their structured-output wording, which no longer holds for a text or tool-call rejection. A rejected response no longer reaches the caller under any error handling strategy, and a deterministic truncation is retried until each router candidate's attempts are exhausted. Generated-by: Claude Code 2.1.234 (Claude Opus 5) --- .../agents/plan/actions/ChatModelAction.java | 22 +- .../agents/plan/actions/ChatModelInvoker.java | 3 + .../actions/ChatModelActionRetryTest.java | 208 +++++++++++++++++- .../plan/actions/ChatModelActionTest.java | 68 +----- .../plan/actions/chat_model_action.py | 20 +- .../tests/actions/test_chat_model_action.py | 46 ---- .../actions/test_chat_model_action_retry.py | 203 ++++++++++++++++- 7 files changed, 433 insertions(+), 137 deletions(-) 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 89270a5d8c..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 @@ -580,39 +580,35 @@ private static Map takeRoutingMetadata( } /** - * Rejects a response the provider did not finish emitting. A finish reason reporting the - * content as cut off by the token budget or withheld by content filtering raises {@link + * 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. */ - private static void rejectIncompleteResponse(ChatMessage response) { + 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 and cannot yield" - + " structured output. Raise the model's max output tokens," - + " or ask for a smaller output.", + + " 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" - + " and cannot yield structured output. Adjust the prompt or" - + " the provider's content filtering configuration.", + + " 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 { - // Precedes the start report: when this rejects, parsing is never attempted, so there is no - // parser execution to report. - rejectIncompleteResponse(response); - ExecutionReporters.started(ctx, ExecutionReporter.EntityTypes.PARSER, STRUCTURED_OUTPUT); try { ChatMessage structuredResponse = generateStructuredOutput(response, outputSchema); 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 19ee34ddd5..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 @@ -25,17 +25,11 @@ import org.apache.flink.agents.api.metrics.FlinkAgentsMetricGroup; import org.apache.flink.agents.api.trace.ExecutionReporter; 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 java.util.HashMap; -import java.util.Locale; import java.util.Map; -import java.util.stream.Stream; import static org.junit.jupiter.api.Assertions.assertEquals; -import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; @@ -67,14 +61,6 @@ private static ChatMessage generateStructuredOutput( Map.class); } - 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()); - } - @Test void testRecordChatTokenMetricsRecordsWhenAllKeysPresent() { BaseChatModelSetup setup = mock(BaseChatModelSetup.class); @@ -214,62 +200,12 @@ void testCleanLlmResponseWithMultipleLinesInBlock() { } @Test - void testStructuredOutputRejectsTruncatedResponse() { - IllegalStateException error = - assertThrows( - IllegalStateException.class, - () -> - generateStructuredOutput( - reportingContext(), Map.of("finish_reason", "length"))); - - String message = error.getMessage().toLowerCase(Locale.ROOT); - assertTrue(message.contains("truncat"), message); - assertTrue(message.contains("token"), message); - } - - @Test - void testStructuredOutputRejectsContentFilteredResponse() { - IllegalStateException error = - assertThrows( - IllegalStateException.class, - () -> - generateStructuredOutput( - reportingContext(), - 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. - String message = error.getMessage().toLowerCase(Locale.ROOT); - assertTrue(message.contains("withheld"), message); - assertTrue(message.contains("content filter"), message); - } - - @ParameterizedTest - @MethodSource("acceptedFinishReasons") - void testAcceptedFinishReasonParsesStructuredOutput(Map extraArgs) - throws Exception { - ChatMessage parsed = generateStructuredOutput(reportingContext(), extraArgs); + void testStructuredOutputParsesToExpectedValue() throws Exception { + ChatMessage parsed = generateStructuredOutput(reportingContext(), Map.of()); assertEquals(Map.of("answer", "42"), parsed.getExtraArgs().get(Agent.STRUCTURED_OUTPUT)); } - @ParameterizedTest - @ValueSource(strings = {"length", "content_filter"}) - void testRejectedFinishReasonReportsNoParserExecution(String finishReason) throws Exception { - RunnerContext ctx = reportingContext(); - - assertThrows( - IllegalStateException.class, - () -> generateStructuredOutput(ctx, Map.of("finish_reason", finishReason))); - - ExecutionReporter reporter = (ExecutionReporter) ctx; - verify(reporter, never()).reportExecutionStarted(anyString(), anyString(), any()); - verify(reporter, never()).reportExecutionSucceeded(anyString(), anyString(), any()); - verify(reporter, never()) - .reportExecutionFailed(anyString(), anyString(), any(), any(), any()); - } - @Test void testStructuredOutputPreservesInboundExtraArgs() throws Exception { Map extraArgs = new HashMap<>(); diff --git a/python/flink_agents/plan/actions/chat_model_action.py b/python/flink_agents/plan/actions/chat_model_action.py index 3642fcb6f8..b802fda3ba 100644 --- a/python/flink_agents/plan/actions/chat_model_action.py +++ b/python/flink_agents/plan/actions/chat_model_action.py @@ -278,6 +278,9 @@ def _generate_structured_output( 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. @@ -291,16 +294,15 @@ def _reject_incomplete_response(response: ChatMessage) -> None: 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 and cannot yield structured output. " - "Raise the model's max output tokens, or ask for a smaller output." + "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 " - "and cannot yield structured output. Adjust the prompt or the " - "provider's content filtering configuration." + 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) @@ -308,10 +310,6 @@ def _reject_incomplete_response(response: ChatMessage) -> None: def _generate_structured_output_with_report( ctx: RunnerContext, response: ChatMessage, output_schema: OutputSchema ) -> ChatMessage: - # Precedes the start report: when this rejects, parsing is never attempted, - # so there is no parser execution to report. - _reject_incomplete_response(response) - ExecutionReporters.started(ctx, ExecutionEntityTypes.PARSER, STRUCTURED_OUTPUT) try: structured_response = _generate_structured_output(response, output_schema) @@ -430,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 1ac8f1e47d..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 @@ -18,11 +18,9 @@ from unittest.mock import MagicMock from uuid import uuid4 -import pytest from pydantic import BaseModel from pyflink.common.typeinfo import BasicTypeInfo, RowTypeInfo -from flink_agents.api.agents.agent import STRUCTURED_OUTPUT 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 @@ -231,36 +229,7 @@ def _parse(ctx, extra_args) -> ChatMessage: ) -def test_length_finish_reason_rejects_parsing(): - with pytest.raises(ValueError, match="(?i)truncat") as exc_info: - _parse(MagicMock(spec=ExecutionReporter), {"finish_reason": "length"}) - assert "token" in str(exc_info.value).lower() - - -def test_content_filter_finish_reason_rejects_parsing(): - # 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. - with pytest.raises(ValueError, match="(?i)withheld"): - _parse(MagicMock(spec=ExecutionReporter), {"finish_reason": "content_filter"}) - - -@pytest.mark.parametrize("finish_reason", ["length", "content_filter"]) -def test_rejected_finish_reason_reports_no_parser_execution(finish_reason): - ctx = MagicMock(spec=ExecutionReporter) - - with pytest.raises(ValueError): - _parse(ctx, {"finish_reason": finish_reason}) - - ctx.report_execution_started.assert_not_called() - ctx.report_execution_succeeded.assert_not_called() - ctx.report_execution_failed.assert_not_called() - - def test_accepted_finish_reason_reports_parser_execution(): - # Reports only reach a context that satisfies ExecutionReporter, so this - # also pins that the fixture still does. Without it, the assert_not_called - # checks above would hold for any context and stop testing the gate. ctx = MagicMock(spec=ExecutionReporter) _parse(ctx, {"finish_reason": "stop"}) @@ -268,18 +237,3 @@ def test_accepted_finish_reason_reports_parser_execution(): ctx.report_execution_started.assert_called_once() ctx.report_execution_succeeded.assert_called_once() ctx.report_execution_failed.assert_not_called() - - -@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_parses_structured_output(extra_args): - response = _parse(MagicMock(spec=ExecutionReporter), extra_args) - assert response.extra_args[STRUCTURED_OUTPUT].result == 42 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.""" From b913d8c7d9c0453e5b8d1e746fdf8548ef4a2791 Mon Sep 17 00:00:00 2001 From: weiqingy Date: Sun, 30 Aug 2026 16:24:06 -0700 Subject: [PATCH 7/7] [python][integrations] Drop the connection name argument from the finish_reason tests BaseChatModelConnection now rejects unknown fields, so the name argument these two tests passed, which was previously ignored, fails validation. Generated-by: Claude Code 2.1.234 (Claude Opus 5) --- .../azure/tests/test_azure_openai_response_parsing.py | 1 - .../chat_models/openai/tests/test_openai_response_parsing.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) 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 index 2707d5ccf8..e31f72bb12 100644 --- 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 @@ -71,7 +71,6 @@ def _connection( payload["usage"] = usage conn = AzureOpenAIChatModelConnection( - name="azure_openai", api_key="test-key", azure_endpoint="https://example.openai.azure.com", api_version="2024-08-01-preview", 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 10a1a67365..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 @@ -192,7 +192,7 @@ def _connection( payload["usage"] = usage conn = OpenAIChatModelConnection( - name="openai", api_key="test-key", api_base_url="http://localhost" + api_key="test-key", api_base_url="http://localhost" ) mock_client = MagicMock() mock_client.chat.completions.create.return_value = ChatCompletion.construct(