This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 5ce2746d0 [integrations][java][python] Preserve Watsonx finish_reason
in extra_args (#1113)
5ce2746d0 is described below
commit 5ce2746d0ecb95e388798994f69993ce9d734ef5
Author: Ashfaq <[email protected]>
AuthorDate: Mon Sep 14 11:41:33 2026 +0530
[integrations][java][python] Preserve Watsonx finish_reason in extra_args
(#1113)
Generated-by: Claude Code 2.1.226 (Claude Sonnet 5)
---
.../watsonx/WatsonxChatModelConnection.java | 23 +++++++--
.../watsonx/WatsonxChatModelConnectionTest.java | 58 ++++++++++++++++++++++
.../watsonx/tests/test_watsonx_chat_model.py | 58 ++++++++++++++++++++++
.../chat_models/watsonx/watsonx_chat_model.py | 6 +++
4 files changed, 140 insertions(+), 5 deletions(-)
diff --git
a/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
b/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
index 627ccd969..1c1606c59 100644
---
a/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
+++
b/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
@@ -553,14 +553,20 @@ public class WatsonxChatModelConnection extends
BaseChatModelConnection {
}
}
+ /**
+ * Parses the watsonx.ai chat 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.
+ */
@VisibleForTesting
static ChatMessage parseResponse(JsonNode response, String modelName) {
final JsonNode choice = response.required("choices").get(0);
final JsonNode responseMessage = choice.required("message");
final JsonNode finishReasonNode = choice.get("finish_reason");
+ String finishReason = null;
if (finishReasonNode != null && !finishReasonNode.isNull()) {
- final String finishReason = finishReasonNode.asText();
+ finishReason = finishReasonNode.asText();
if (!"stop".equals(finishReason) &&
!"tool_calls".equals(finishReason)) {
LOG.warn(
"watsonx.ai chat for model {} finished with reason
'{}'; the response"
@@ -602,11 +608,18 @@ public class WatsonxChatModelConnection extends
BaseChatModelConnection {
}
final JsonNode usage = response.get("usage");
- if (modelName != null && !modelName.isBlank() && usage != null &&
!usage.isNull()) {
+ final boolean hasUsageMetadata =
+ modelName != null && !modelName.isBlank() && usage != null &&
!usage.isNull();
+ if (hasUsageMetadata || finishReason != null) {
final Map<String, Object> extraArgs = new
HashMap<>(chatMessage.getExtraArgs());
- extraArgs.put("model_name", modelName);
- extraArgs.put("promptTokens",
usage.path("prompt_tokens").asLong(0));
- extraArgs.put("completionTokens",
usage.path("completion_tokens").asLong(0));
+ if (hasUsageMetadata) {
+ extraArgs.put("model_name", modelName);
+ extraArgs.put("promptTokens",
usage.path("prompt_tokens").asLong(0));
+ extraArgs.put("completionTokens",
usage.path("completion_tokens").asLong(0));
+ }
+ if (finishReason != null) {
+ extraArgs.put("finish_reason", finishReason);
+ }
chatMessage.setExtraArgs(extraArgs);
}
diff --git
a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
index 9b1cd5b6d..f594e9329 100644
---
a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
+++
b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
@@ -380,6 +380,64 @@ class WatsonxChatModelConnectionTest {
.isEqualTo("ibm/granite-3-3-8b-instruct");
assertThat(message.getExtraArgs().get("promptTokens")).isEqualTo(100L);
assertThat(message.getExtraArgs().get("completionTokens")).isEqualTo(50L);
+ assertThat(message.getExtraArgs()).containsEntry("finish_reason",
"stop");
+ }
+
+ @Test
+ @DisplayName("A finish reason outside the documented set is stored as
received")
+ void testParseResponseCarriesUnknownFinishReasonVerbatim() throws
Exception {
+ JsonNode response =
+ MAPPER.readTree(
+ "{\"choices\": [{\"index\": 0, \"message\": {\"role\":
\"assistant\","
+ + " \"content\": \"hi\"}, \"finish_reason\":"
+ + " \"some_vendor_reason\"}]}");
+
+ ChatMessage message =
WatsonxChatModelConnection.parseResponse(response, null);
+
+ assertThat(message.getExtraArgs()).containsEntry("finish_reason",
"some_vendor_reason");
+ }
+
+ @Test
+ @DisplayName("A response with no finish_reason member yields no key and no
error")
+ void testParseResponseNoFinishReasonKeyWhenMemberAbsent() throws Exception
{
+ JsonNode response =
+ MAPPER.readTree(
+ "{\"choices\": [{\"index\": 0, \"message\": {\"role\":
\"assistant\","
+ + " \"content\": \"hi\"}}]}");
+
+ ChatMessage message =
WatsonxChatModelConnection.parseResponse(response, null);
+
+ assertThat(message.getExtraArgs()).doesNotContainKey("finish_reason");
+ }
+
+ @Test
+ @DisplayName("A response whose finish_reason is JSON null yields no key
and no error")
+ void testParseResponseNoFinishReasonKeyWhenJsonNull() throws Exception {
+ JsonNode response =
+ MAPPER.readTree(
+ "{\"choices\": [{\"index\": 0, \"message\": {\"role\":
\"assistant\","
+ + " \"content\": \"hi\"}, \"finish_reason\":
null}]}");
+
+ ChatMessage message =
WatsonxChatModelConnection.parseResponse(response, null);
+
+ assertThat(message.getExtraArgs()).doesNotContainKey("finish_reason");
+ }
+
+ @Test
+ @DisplayName("The finish reason is captured independently of the token
metrics")
+ void testParseResponseCarriesFinishReasonWithoutUsage() throws Exception {
+ // modelName is null here, so the usage-metadata branch cannot run;
this proves
+ // finish_reason capture does not depend on it.
+ JsonNode response =
+ MAPPER.readTree(
+ "{\"choices\": [{\"index\": 0, \"message\": {\"role\":
\"assistant\","
+ + " \"content\": \"hi\"}, \"finish_reason\":
\"tool_calls\"}]}");
+
+ ChatMessage message =
WatsonxChatModelConnection.parseResponse(response, null);
+
+ assertThat(message.getExtraArgs())
+ .containsEntry("finish_reason", "tool_calls")
+ .doesNotContainKey("promptTokens");
}
@Test
diff --git
a/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_chat_model.py
b/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_chat_model.py
index ab79674c6..6d9a38fcf 100644
---
a/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_chat_model.py
+++
b/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_chat_model.py
@@ -130,6 +130,7 @@ def test_watsonx_chat_mocked(monkeypatch:
pytest.MonkeyPatch) -> None:
assert response.extra_args["model_name"] == test_model
assert response.extra_args["promptTokens"] == 100
assert response.extra_args["completionTokens"] == 50
+ assert response.extra_args["finish_reason"] == "stop"
model_inference.assert_called_once_with(
model_id=test_model,
api_client=api_client,
@@ -139,6 +140,63 @@ def test_watsonx_chat_mocked(monkeypatch:
pytest.MonkeyPatch) -> None:
)
+def _setup_llm(
+ mock_model: MagicMock, monkeypatch: pytest.MonkeyPatch
+) -> WatsonxChatModelSetup:
+ model_inference = MagicMock(return_value=mock_model)
+ monkeypatch.setattr(
+
"flink_agents.integrations.chat_models.watsonx.watsonx_chat_model.ModelInference",
+ model_inference,
+ )
+
+ connection = _fake_connection()
+ connection._client = MagicMock()
+
+ def get_resource(name: str, type: ResourceType) -> Resource:
+ return connection
+
+ mock_ctx = MagicMock(spec=ResourceContext)
+ mock_ctx.get_resource = get_resource
+
+ llm = WatsonxChatModelSetup(
+ model=test_model,
+ connection="watsonx",
+ resource_context=mock_ctx,
+ )
+ llm.open()
+ return llm
+
+
+def test_watsonx_chat_carries_unknown_finish_reason_verbatim(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """A finish reason outside the documented set is stored as received."""
+ mock_model = MagicMock()
+ mock_model.chat.return_value = _mock_chat_response(
+ {"role": "assistant", "content": "hi"},
finish_reason="some_vendor_reason"
+ )
+ llm = _setup_llm(mock_model, monkeypatch)
+
+ response = llm.chat([ChatMessage(role=MessageRole.USER, content="Hello!")])
+
+ assert response.extra_args["finish_reason"] == "some_vendor_reason"
+
+
+def test_watsonx_chat_no_finish_reason_key_when_none(
+ monkeypatch: pytest.MonkeyPatch,
+) -> None:
+ """A response with no finish reason yields no key and no error."""
+ mock_model = MagicMock()
+ mock_model.chat.return_value = _mock_chat_response(
+ {"role": "assistant", "content": "hi"}, finish_reason=None
+ )
+ llm = _setup_llm(mock_model, monkeypatch)
+
+ response = llm.chat([ChatMessage(role=MessageRole.USER, content="Hello!")])
+
+ assert "finish_reason" not in response.extra_args
+
+
def test_watsonx_tool_call_response_mocked(monkeypatch: pytest.MonkeyPatch) ->
None:
"""Test that tool call responses are converted to the framework format."""
mock_model = MagicMock()
diff --git
a/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
b/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
index 6544d2218..b553e3f95 100644
--- a/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
+++ b/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
@@ -345,6 +345,10 @@ class WatsonxChatModelConnection(BaseChatModelConnection):
structured-output translation, so callers stay on the
prompt-engineering
fallback. Declaring the parameter keeps a caller-supplied schema out of
``**kwargs``, which is forwarded to the provider SDK.
+
+ When the response carries a finish reason, it is available verbatim as
+ ``extra_args["finish_reason"]``; the key is absent when the provider
+ reports none.
"""
self._reject_unsupported_output_schema(output_schema)
model_name = kwargs.pop("model", DEFAULT_MODEL)
@@ -401,6 +405,8 @@ class WatsonxChatModelConnection(BaseChatModelConnection):
model_name,
finish_reason,
)
+ if finish_reason is not None:
+ extra_args["finish_reason"] = finish_reason
response_message: Dict[str, Any] = choice["message"]
