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 @@ static ArrayNode convertTools(List tools) { } } + /** + * 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 @@ static ChatMessage parseResponse(JsonNode response, String modelName) { } 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 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 @@ void testParseResponse() throws Exception { .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 get_resource(name: str, type: ResourceType) -> Resource: 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 get_resource(name: str, type: ResourceType) -> Resource: ) +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 @@ def chat( 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 @@ def chat( model_name, finish_reason, ) + if finish_reason is not None: + extra_args["finish_reason"] = finish_reason response_message: Dict[str, Any] = choice["message"]