weiqingy commented on code in PR #922: URL: https://github.com/apache/flink-agents/pull/922#discussion_r3697192791
########## python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py: ########## @@ -0,0 +1,476 @@ +################################################################################ +# 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. +################################################################################# +import ast +import contextlib +import json +import logging +import os +import time +import uuid +from typing import Any, Dict, List, Sequence + +import httpx +from ibm_watsonx_ai import APIClient, Credentials +from ibm_watsonx_ai.foundation_models import ModelInference +from ibm_watsonx_ai.wml_client_error import ApiRequestFailure +from pydantic import Field, PrivateAttr +from typing_extensions import override + +from flink_agents.api.chat_message import ChatMessage, MessageRole +from flink_agents.api.chat_models.chat_model import ( + BaseChatModelConnection, + BaseChatModelSetup, +) +from flink_agents.api.tools.tool import Tool +from flink_agents.integrations.chat_models.chat_model_utils import to_openai_tool + +logger = logging.getLogger(__name__) + +DEFAULT_MODEL = "ibm/granite-4-h-small" +DEFAULT_REQUEST_TIMEOUT = 120.0 +DEFAULT_MAX_RETRIES = 3 +RETRYABLE_STATUS_CODES = frozenset({408, 429, 500, 502, 503, 504}) +RESERVED_ADDITIONAL_KWARGS = frozenset( + { + "model", + "temperature", + "max_tokens", + "extract_reasoning", + "tool_choice", + "tool_choice_option", + } +) + + +def _normalize(value: str | None) -> str | None: + if value is None or not value.strip(): + return None + return value.strip() + + +def _retry_delay_seconds(attempt: int, response: httpx.Response | None) -> float: + """Return capped exponential backoff, honoring a numeric Retry-After header.""" + backoff = min(2**attempt, 10) + if response is not None: + retry_after = response.headers.get("Retry-After") + if retry_after is not None: + with contextlib.suppress(ValueError): + return max(backoff, min(float(retry_after.strip()), 30)) + return backoff + + +def convert_to_watsonx_messages( + messages: Sequence[ChatMessage], +) -> List[Dict[str, Any]]: + """Convert framework messages to the watsonx.ai chat format.""" + watsonx_messages: List[Dict[str, Any]] = [] + for message in messages: + role = message.role + + if role == MessageRole.ASSISTANT: + assistant_message: Dict[str, Any] = {"role": "assistant"} + if message.content: + assistant_message["content"] = message.content + if message.tool_calls: + assistant_message["tool_calls"] = [ + _convert_to_watsonx_tool_call(tool_call) + for tool_call in message.tool_calls + ] + watsonx_messages.append(assistant_message) + elif role == MessageRole.TOOL: + tool_call_id = message.extra_args.get("external_id") + if not tool_call_id or not isinstance(tool_call_id, str): + msg = "Tool message must have 'external_id' as a string in extra_args" + raise ValueError(msg) + watsonx_messages.append( + { + "role": "tool", + "content": message.content, + "tool_call_id": tool_call_id, + } + ) + else: + watsonx_messages.append({"role": role.value, "content": message.content}) + return watsonx_messages + + +def _parse_tool_arguments(args: Any) -> Dict[str, Any]: + """Parse model-emitted tool arguments, including common malformed variants.""" + if args is None or args == "": + return {} + raw = args + for _ in range(3): + if not isinstance(args, str): + break + try: + args = json.loads(args) + except ValueError: + with contextlib.suppress(Exception): + literal = ast.literal_eval(args) + if isinstance(literal, dict): + args = literal + break + if not isinstance(args, dict): + msg = ( + "Failed to parse tool call arguments returned by watsonx.ai " + f"as a JSON object: {raw!r}" + ) + raise TypeError(msg) + return args + + +def _convert_to_watsonx_tool_call(tool_call: Dict[str, Any]) -> Dict[str, Any]: + """Convert a framework tool call to watsonx.ai format.""" + watsonx_tool_call_id = tool_call.get("original_id") + if watsonx_tool_call_id is None: + tool_call_id = tool_call.get("id") + if tool_call_id is None: + msg = "Tool call must have either 'original_id' or 'id' field" + raise ValueError(msg) + watsonx_tool_call_id = str(tool_call_id) + + arguments = tool_call["function"]["arguments"] + return { + "id": watsonx_tool_call_id, + "type": "function", + "function": { + "name": tool_call["function"]["name"], + "arguments": json.dumps(arguments) + if isinstance(arguments, dict) + else arguments, + }, + } + + +class WatsonxChatModelConnection(BaseChatModelConnection): + """Connection to the IBM watsonx.ai chat API.""" + + url: str = Field(description="The watsonx.ai service endpoint.") + api_key: str | None = Field(default=None, description="The IBM Cloud API key.") + token: str | None = Field( + default=None, description="A bearer token, as an alternative to api_key." + ) + project_id: str | None = Field( + default=None, description="The watsonx.ai project id." + ) + space_id: str | None = Field( + default=None, description="The watsonx.ai deployment space id." + ) + request_timeout: float = Field( + default=DEFAULT_REQUEST_TIMEOUT, + description="The timeout, in seconds, for chat requests to watsonx.ai.", + gt=0, + allow_inf_nan=False, + ) + max_retries: int = Field( + default=DEFAULT_MAX_RETRIES, + description="Maximum number of retries for transient failures.", + ge=0, + ) + + _client: APIClient | None = PrivateAttr(default=None) + _http_client: httpx.Client | None = PrivateAttr(default=None) + _models: Dict[str, ModelInference] = PrivateAttr(default_factory=dict) + + def __init__( + self, + *, + url: str | None = None, + api_key: str | None = None, + token: str | None = None, + project_id: str | None = None, + space_id: str | None = None, + request_timeout: float = DEFAULT_REQUEST_TIMEOUT, + max_retries: int = DEFAULT_MAX_RETRIES, + **kwargs: Any, + ) -> None: + """Initialize the connection.""" + resolved_url = _normalize(url) or _normalize(os.environ.get("WATSONX_URL")) + resolved_api_key = _normalize(api_key) or _normalize( + os.environ.get("WATSONX_API_KEY") + ) + resolved_token = _normalize(token) or _normalize( + os.environ.get("WATSONX_TOKEN") + ) + resolved_project_id = _normalize(project_id) or _normalize( + os.environ.get("WATSONX_PROJECT_ID") + ) + resolved_space_id = _normalize(space_id) or _normalize( + os.environ.get("WATSONX_SPACE_ID") + ) + + if not resolved_url: + msg = ( + "watsonx.ai url is not provided. Please pass it as an argument " + "or set the 'WATSONX_URL' environment variable." + ) + raise ValueError(msg) + if not resolved_api_key and not resolved_token: + msg = ( + "watsonx.ai credentials are not provided. Please pass 'api_key' " + "or 'token' as an argument, or set the 'WATSONX_API_KEY' or " + "'WATSONX_TOKEN' environment variable." + ) + raise ValueError(msg) + if resolved_api_key and resolved_token: + msg = ( + "watsonx.ai api_key and token cannot both be provided. Please configure " + "exactly one credential source." + ) + raise ValueError(msg) + if not resolved_project_id and not resolved_space_id: + msg = ( + "watsonx.ai project or space is not provided. Please pass " + "'project_id' or 'space_id' as an argument, or set the " + "'WATSONX_PROJECT_ID' or 'WATSONX_SPACE_ID' environment variable." + ) + raise ValueError(msg) + if resolved_project_id and resolved_space_id: + msg = ( + "watsonx.ai project and space cannot both be provided. Please configure " + "exactly one of 'project_id' or 'space_id'." + ) + raise ValueError(msg) + + super().__init__( + url=resolved_url, + api_key=resolved_api_key, + token=resolved_token, + project_id=resolved_project_id, + space_id=resolved_space_id, + request_timeout=request_timeout, + max_retries=max_retries, + **kwargs, + ) + + @property + def client(self) -> APIClient: + """Return the lazily initialized API client.""" + if self._client is None: + credential_kwargs: Dict[str, Any] = {"url": self.url} + if self.api_key: + credential_kwargs["api_key"] = self.api_key + if self.token: + credential_kwargs["token"] = self.token + self._http_client = httpx.Client(timeout=self.request_timeout) + self._client = APIClient( + credentials=Credentials(**credential_kwargs), + project_id=self.project_id, + space_id=self.space_id, + httpx_client=self._http_client, + ) + return self._client + + @override + def close(self) -> None: + """Close the underlying HTTP client.""" + self._models = {} + self._client = None + if self._http_client is not None: + with contextlib.suppress(Exception): + self._http_client.close() + self._http_client = None + + def _get_model(self, model: str) -> ModelInference: + if model not in self._models: + self._models[model] = ModelInference( + model_id=model, + api_client=self.client, + project_id=self.project_id, + space_id=self.space_id, + ) + return self._models[model] + + def _chat_with_retry(self, model_name: str, **chat_kwargs: Any) -> Dict[str, Any]: + """Call chat with retries for selected HTTP statuses and transport failures. + + Retries use capped exponential backoff and honor a numeric ``Retry-After`` + response header. + """ + attempt = 0 + while True: + try: + return self._get_model(model_name).chat(**chat_kwargs) + except (ApiRequestFailure, httpx.TransportError) as e: # noqa: PERF203 + response = getattr(e, "response", None) + status = getattr(response, "status_code", None) + retryable = isinstance(e, httpx.TransportError) or ( + status in RETRYABLE_STATUS_CODES + ) + if attempt >= self.max_retries or not retryable: + raise + delay = _retry_delay_seconds(attempt, response) + logger.warning( + "watsonx.ai chat request for model %s failed with %s; " + "retry %d/%d in %ds", + model_name, + status if status is not None else type(e).__name__, + attempt + 1, + self.max_retries, + delay, + ) + time.sleep(delay) + attempt += 1 + + def chat( Review Comment: This one is a rebase follow-up rather than something done wrong. #843 ("[api][python] Add explicit output schema parameter to the chat path (structured-output foundation)") added `output_schema: OutputSchema | None = None` to `BaseChatModelConnection.chat` (`python/flink_agents/api/chat_models/chat_model.py:230`) on 2026-07-25, after this branch was cut. The docstring at `:241-249` asks for a named parameter, because "`**kwargs` is forwarded to the provider SDK, so a schema landing there would reach the request body." That happens literally here: `watsonx_chat_model.py:362` forwards `params={**kwargs, **additional_kwargs}` into `ModelInference.chat(params=...)`, and the SDK does `payload.update(parameters)` (`ibm_watsonx_ai/foundation_models/inference/fm_model_inference.py:597-598`). The same commit added `python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py`. After a rebase, `test_every_connection_declares_output_schema_param` (`:90`) and the rejection test at `:143` fail on a whole-tree `pytest` run, where collection imports the watsonx test module. That file run on its own still passes, since "`__subclasses__()` only sees classes that have been imported" and `watsonx_chat_model` isn't in its import list at `:41-51`. The green checks predate all of it: the last CI run was created 2026-07-24, a day before #843 merged. `openai_chat_model.py:142` and `:167` are the closest template. The leak stays latent even after the rebase, since `BaseChatModelSetup.chat` doesn't forward `output_schema` yet, so only a direct `connection.chat(..., output_schema=...)` caller reaches the body. Easier to pick the parameter and the import-list entry up as part of the rebase, or would you rather handle them separately once the branch is current? ########## integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java: ########## @@ -0,0 +1,678 @@ +/* + * 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.watsonx; + +import com.fasterxml.jackson.core.json.JsonReadFeature; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +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.BaseChatModelConnection; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.annotation.VisibleForTesting; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.net.URI; +import java.net.URLEncoder; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Function; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +/** Chat model connection for the IBM watsonx.ai text chat REST API. */ +public class WatsonxChatModelConnection extends BaseChatModelConnection { + + private static final Logger LOG = LoggerFactory.getLogger(WatsonxChatModelConnection.class); + + static final String DEFAULT_IAM_URL = "https://iam.cloud.ibm.com";; + static final String DEFAULT_API_VERSION = "2025-04-23"; + static final long DEFAULT_REQUEST_TIMEOUT_SEC = 120; + static final int DEFAULT_MAX_RETRIES = 3; + private static final Set<Integer> RETRYABLE_STATUS_CODES = Set.of(408, 429, 500, 502, 503, 504); + + private static final Set<String> CONTROL_PARAMS = + Set.of( + "model", + "tool_choice", + "tool_choice_option", + "extract_reasoning", + "additional_kwargs"); + private static final Set<String> RESERVED_ADDITIONAL_KWARGS = + Set.of( + "model", + "temperature", + "max_tokens", + "extract_reasoning", + "tool_choice", + "tool_choice_option"); + + private static final Pattern[] REASONING_PATTERNS = { + Pattern.compile("<think>(.*?)</think>", Pattern.DOTALL | Pattern.CASE_INSENSITIVE), + Pattern.compile("<analysis>(.*?)</analysis>", Pattern.DOTALL | Pattern.CASE_INSENSITIVE), + Pattern.compile("<reasoning>(.*?)</reasoning>", Pattern.DOTALL | Pattern.CASE_INSENSITIVE), + Pattern.compile( + "```(?:think|reasoning|thought)\\s*\\n(.*?)\\n```", + Pattern.DOTALL | Pattern.CASE_INSENSITIVE), + Pattern.compile( + "(?:^|\\n)Reasoning:\\s*(.*?)(?:\\n{2,}|$)", + Pattern.DOTALL | Pattern.CASE_INSENSITIVE), + }; + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private static final ObjectMapper LENIENT_MAPPER = + JsonMapper.builder() + .enable(JsonReadFeature.ALLOW_SINGLE_QUOTES) + .enable(JsonReadFeature.ALLOW_UNQUOTED_FIELD_NAMES) + .build(); + + private final String url; + private final String apiKey; + private final String staticToken; + private final String projectId; + private final String spaceId; + private final String apiVersion; + private final String iamUrl; + private final Duration requestTimeout; + private final int maxRetries; + + private final HttpClient httpClient; + + private transient String cachedIamToken; + private transient long iamTokenExpirationEpochSec; + + public WatsonxChatModelConnection( + ResourceDescriptor descriptor, ResourceContext resourceContext) { + this(descriptor, resourceContext, System::getenv); + } + + @VisibleForTesting + WatsonxChatModelConnection( + ResourceDescriptor descriptor, + ResourceContext resourceContext, + Function<String, String> environmentLookup) { + super(descriptor, resourceContext); + + this.url = + trimTrailingSlash( + argumentOrEnv(descriptor, "url", "WATSONX_URL", environmentLookup)); + this.apiKey = argumentOrEnv(descriptor, "api_key", "WATSONX_API_KEY", environmentLookup); + this.staticToken = argumentOrEnv(descriptor, "token", "WATSONX_TOKEN", environmentLookup); + this.projectId = + argumentOrEnv(descriptor, "project_id", "WATSONX_PROJECT_ID", environmentLookup); + this.spaceId = argumentOrEnv(descriptor, "space_id", "WATSONX_SPACE_ID", environmentLookup); + + String apiVersion = normalize(descriptor.getArgument("api_version")); + this.apiVersion = apiVersion != null ? apiVersion : DEFAULT_API_VERSION; + String iamUrl = normalize(descriptor.getArgument("iam_url")); + this.iamUrl = trimTrailingSlash(iamUrl != null ? iamUrl : DEFAULT_IAM_URL); + Number requestTimeout = descriptor.getArgument("request_timeout"); + double requestTimeoutSeconds = + requestTimeout != null ? requestTimeout.doubleValue() : DEFAULT_REQUEST_TIMEOUT_SEC; + if (!Double.isFinite(requestTimeoutSeconds) || requestTimeoutSeconds <= 0) { + throw new IllegalArgumentException("request_timeout must be a positive finite number."); + } + this.requestTimeout = + Duration.ofMillis(Math.max(1L, Math.round(requestTimeoutSeconds * 1000.0))); + Number maxRetries = descriptor.getArgument("max_retries"); + this.maxRetries = + maxRetries != null + ? requireInteger(maxRetries, "max_retries", 0) + : DEFAULT_MAX_RETRIES; + + if (this.url == null || this.url.isEmpty()) { + throw new IllegalArgumentException( + "watsonx.ai url is not provided. Please pass the 'url' argument or set the" + + " 'WATSONX_URL' environment variable."); + } + if ((this.apiKey == null || this.apiKey.isEmpty()) + && (this.staticToken == null || this.staticToken.isEmpty())) { + throw new IllegalArgumentException( + "watsonx.ai credentials are not provided. Please pass the 'api_key' or 'token'" + + " argument, or set the 'WATSONX_API_KEY' or 'WATSONX_TOKEN'" + + " environment variable."); + } + if (this.apiKey != null && this.staticToken != null) { + throw new IllegalArgumentException( + "watsonx.ai api_key and token cannot both be provided. Please configure" + + " exactly one credential source."); + } + if ((this.projectId == null || this.projectId.isEmpty()) + && (this.spaceId == null || this.spaceId.isEmpty())) { + throw new IllegalArgumentException( + "watsonx.ai project or space is not provided. Please pass the 'project_id' or" + + " 'space_id' argument, or set the 'WATSONX_PROJECT_ID' or" + + " 'WATSONX_SPACE_ID' environment variable."); + } + if (this.projectId != null + && !this.projectId.isEmpty() + && this.spaceId != null + && !this.spaceId.isEmpty()) { + throw new IllegalArgumentException( + "watsonx.ai project and space cannot both be provided. Please configure" + + " exactly one of 'project_id' or 'space_id'."); + } + + this.httpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(30)).build(); + } + + static int requireInteger(Number value, String argumentName, int minimum) { + double numericValue = value.doubleValue(); + if (!Double.isFinite(numericValue) + || numericValue != Math.rint(numericValue) + || numericValue < minimum + || numericValue > Integer.MAX_VALUE) { + throw new IllegalArgumentException( + argumentName + + " must be " + + (minimum == 0 ? "a non-negative" : "a positive") + + " integer."); + } + return (int) numericValue; + } + + private static String argumentOrEnv( + ResourceDescriptor descriptor, + String argumentName, + String envName, + Function<String, String> environmentLookup) { + String value = normalize(descriptor.getArgument(argumentName)); + if (value == null) { + value = normalize(environmentLookup.apply(envName)); + } + return value; + } + + private static String normalize(String value) { + if (value == null || value.isBlank()) { + return null; + } + return value.trim(); + } + + private static String trimTrailingSlash(String value) { + if (value != null && value.endsWith("/")) { + return value.substring(0, value.length() - 1); + } + return value; + } + + @Override + public ChatMessage chat( + List<ChatMessage> messages, List<Tool> tools, Map<String, Object> modelParams) { + try { + final String modelName = (String) modelParams.get("model"); + final boolean extractReasoning = + Boolean.TRUE.equals(modelParams.get("extract_reasoning")); + final ObjectNode payload = buildPayload(messages, tools, modelParams); + if (projectId != null && !projectId.isEmpty()) { + payload.put("project_id", projectId); + } else { + payload.put("space_id", spaceId); + } + + final HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create(url + "/ml/v1/text/chat?version=" + apiVersion)) + .timeout(requestTimeout) + .header("Authorization", "Bearer " + getBearerToken()) + .header("Content-Type", "application/json") + .header("Accept", "application/json") + .POST( + HttpRequest.BodyPublishers.ofString( + MAPPER.writeValueAsString(payload))) + .build(); + + final HttpResponse<String> response = sendWithRetry(request); + if (response.statusCode() / 100 != 2) { + throw new RuntimeException( + String.format( + "watsonx.ai chat request failed with status %d: %s", + response.statusCode(), response.body())); + } + + final ChatMessage chatMessage = + parseResponse(MAPPER.readTree(response.body()), modelName); + if (extractReasoning) { + final String[] parts = extractReasoning(chatMessage.getContent()); + chatMessage.setContent(parts[0]); + if (parts[1] != null) { + chatMessage.getExtraArgs().put("reasoning", parts[1]); + } + } + return chatMessage; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted while calling watsonx.ai.", e); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + /** + * Sends the request, retrying HTTP 408, 429, 500, 502, 503, and 504 responses and I/O errors up + * to {@code max_retries} times with capped exponential backoff, honoring {@code Retry-After}. + */ + private HttpResponse<String> sendWithRetry(HttpRequest request) + throws IOException, InterruptedException { + for (int attempt = 0; ; attempt++) { + try { + final HttpResponse<String> response = + httpClient.send(request, HttpResponse.BodyHandlers.ofString()); + if (attempt >= maxRetries || !isRetryableStatus(response.statusCode())) { + return response; + } + final long delayMillis = + retryDelayMillis( + attempt, response.headers().firstValue("Retry-After").orElse(null)); + LOG.warn( + "watsonx.ai request to {} returned status {}; retry {}/{} in {} ms", + request.uri().getPath(), + response.statusCode(), + attempt + 1, + maxRetries, + delayMillis); + Thread.sleep(delayMillis); + } catch (IOException e) { + if (attempt >= maxRetries) { + throw e; + } + final long delayMillis = retryDelayMillis(attempt, null); + LOG.warn( + "watsonx.ai request to {} failed ({}); retry {}/{} in {} ms", + request.uri().getPath(), + e.toString(), + attempt + 1, + maxRetries, + delayMillis); + Thread.sleep(delayMillis); + } + } + } + + @VisibleForTesting + static boolean isRetryableStatus(int status) { + return RETRYABLE_STATUS_CODES.contains(status); + } + + @VisibleForTesting + static long retryDelayMillis(int attempt, String retryAfterHeader) { + long backoffMillis = Math.min(1000L << attempt, 10_000L); + if (retryAfterHeader != null) { + try { + long retryAfterMillis = + Math.min(Long.parseLong(retryAfterHeader.trim()) * 1000L, 30_000L); + backoffMillis = Math.max(backoffMillis, retryAfterMillis); + } catch (NumberFormatException ignored) { + // Retry-After may be an HTTP date; fall back to exponential backoff. + } + } + return backoffMillis; + } + + @VisibleForTesting + static String[] extractReasoning(String content) { + if (content == null || content.isEmpty()) { + return new String[] {"", null}; + } + final List<String> reasoningChunks = new ArrayList<>(); + String cleaned = content; + for (Pattern pattern : REASONING_PATTERNS) { + final Matcher matcher = pattern.matcher(cleaned); + final StringBuilder rest = new StringBuilder(); + boolean found = false; + int position = 0; + while (matcher.find()) { + final String chunk = matcher.group(1).trim(); + if (!chunk.isEmpty()) { + reasoningChunks.add(chunk); + } + rest.append(cleaned, position, matcher.start()); + position = matcher.end(); + found = true; + } + if (found) { + rest.append(cleaned, position, cleaned.length()); + cleaned = rest.toString(); + } + } + final String reasoning = + reasoningChunks.isEmpty() ? null : String.join("\n\n", reasoningChunks); + cleaned = cleaned.replaceAll("\\n{3,}", "\n\n").replaceAll(" {2,}", " ").trim(); + return new String[] {cleaned, reasoning}; + } + + @VisibleForTesting + static ObjectNode buildPayload( + List<ChatMessage> messages, List<Tool> tools, Map<String, Object> modelParams) { + final ObjectNode payload = MAPPER.createObjectNode(); + payload.put("model_id", (String) modelParams.get("model")); + payload.set("messages", convertMessages(messages)); + + if (tools != null && !tools.isEmpty()) { + payload.set("tools", convertTools(tools)); + } + final Object toolChoice = modelParams.get("tool_choice"); + if (toolChoice != null) { + payload.set("tool_choice", MAPPER.valueToTree(toolChoice)); + } + final Object toolChoiceOption = modelParams.get("tool_choice_option"); + if (toolChoiceOption != null) { + payload.put("tool_choice_option", toolChoiceOption.toString()); + } + + @SuppressWarnings("unchecked") + final Map<String, Object> additionalKwargs = + (Map<String, Object>) modelParams.get("additional_kwargs"); + if (additionalKwargs != null) { + final Set<String> collisions = new java.util.HashSet<>(additionalKwargs.keySet()); Review Comment: `buildPayload` sets the structural fields first, `model_id` at `:383`, `messages` at `:384` and `tools` at `:387`, then spills `additional_kwargs` over the top at `:410-415`. This collision check only guards `RESERVED_ADDITIONAL_KWARGS`, which is `model`, `temperature`, `max_tokens`, `extract_reasoning`, `tool_choice` and `tool_choice_option` (`:72-79`). So `additional_kwargs: {messages: [...]}` passes the check and replaces the real conversation, and `tools` goes the same way. `model_id` is the sharper case, because reserving `model` doesn't cover it: the payload key is `model_id`, so `additional_kwargs: {model_id: "other-model"}` silently overrides the configured model while the reserved-name check sees nothing. Python has the same shape through `params={**kwargs, **additional_kwargs}` at `watsonx_chat_model.py:362`, where the SDK's `payload.update(parameters)` lets those keys into the body. Would it make sense to reserve the request-owned keys too, so static setup configuration can't replace the model, the messages or the tools? ########## python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py: ########## @@ -0,0 +1,476 @@ +################################################################################ +# 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. +################################################################################# +import ast +import contextlib +import json +import logging +import os +import time +import uuid +from typing import Any, Dict, List, Sequence + +import httpx +from ibm_watsonx_ai import APIClient, Credentials +from ibm_watsonx_ai.foundation_models import ModelInference +from ibm_watsonx_ai.wml_client_error import ApiRequestFailure +from pydantic import Field, PrivateAttr +from typing_extensions import override + +from flink_agents.api.chat_message import ChatMessage, MessageRole +from flink_agents.api.chat_models.chat_model import ( + BaseChatModelConnection, + BaseChatModelSetup, +) +from flink_agents.api.tools.tool import Tool +from flink_agents.integrations.chat_models.chat_model_utils import to_openai_tool + +logger = logging.getLogger(__name__) + +DEFAULT_MODEL = "ibm/granite-4-h-small" +DEFAULT_REQUEST_TIMEOUT = 120.0 +DEFAULT_MAX_RETRIES = 3 +RETRYABLE_STATUS_CODES = frozenset({408, 429, 500, 502, 503, 504}) +RESERVED_ADDITIONAL_KWARGS = frozenset( + { + "model", + "temperature", + "max_tokens", + "extract_reasoning", + "tool_choice", + "tool_choice_option", + } +) + + +def _normalize(value: str | None) -> str | None: + if value is None or not value.strip(): + return None + return value.strip() + + +def _retry_delay_seconds(attempt: int, response: httpx.Response | None) -> float: + """Return capped exponential backoff, honoring a numeric Retry-After header.""" + backoff = min(2**attempt, 10) + if response is not None: + retry_after = response.headers.get("Retry-After") + if retry_after is not None: + with contextlib.suppress(ValueError): + return max(backoff, min(float(retry_after.strip()), 30)) + return backoff + + +def convert_to_watsonx_messages( + messages: Sequence[ChatMessage], +) -> List[Dict[str, Any]]: + """Convert framework messages to the watsonx.ai chat format.""" + watsonx_messages: List[Dict[str, Any]] = [] + for message in messages: + role = message.role + + if role == MessageRole.ASSISTANT: + assistant_message: Dict[str, Any] = {"role": "assistant"} + if message.content: + assistant_message["content"] = message.content + if message.tool_calls: + assistant_message["tool_calls"] = [ + _convert_to_watsonx_tool_call(tool_call) + for tool_call in message.tool_calls + ] + watsonx_messages.append(assistant_message) + elif role == MessageRole.TOOL: + tool_call_id = message.extra_args.get("external_id") + if not tool_call_id or not isinstance(tool_call_id, str): + msg = "Tool message must have 'external_id' as a string in extra_args" + raise ValueError(msg) + watsonx_messages.append( + { + "role": "tool", + "content": message.content, + "tool_call_id": tool_call_id, + } + ) + else: + watsonx_messages.append({"role": role.value, "content": message.content}) + return watsonx_messages + + +def _parse_tool_arguments(args: Any) -> Dict[str, Any]: + """Parse model-emitted tool arguments, including common malformed variants.""" + if args is None or args == "": + return {} + raw = args + for _ in range(3): + if not isinstance(args, str): + break + try: + args = json.loads(args) + except ValueError: + with contextlib.suppress(Exception): + literal = ast.literal_eval(args) + if isinstance(literal, dict): + args = literal + break + if not isinstance(args, dict): + msg = ( + "Failed to parse tool call arguments returned by watsonx.ai " + f"as a JSON object: {raw!r}" + ) + raise TypeError(msg) + return args + + +def _convert_to_watsonx_tool_call(tool_call: Dict[str, Any]) -> Dict[str, Any]: + """Convert a framework tool call to watsonx.ai format.""" + watsonx_tool_call_id = tool_call.get("original_id") + if watsonx_tool_call_id is None: + tool_call_id = tool_call.get("id") + if tool_call_id is None: + msg = "Tool call must have either 'original_id' or 'id' field" + raise ValueError(msg) + watsonx_tool_call_id = str(tool_call_id) + + arguments = tool_call["function"]["arguments"] + return { + "id": watsonx_tool_call_id, + "type": "function", + "function": { + "name": tool_call["function"]["name"], + "arguments": json.dumps(arguments) + if isinstance(arguments, dict) + else arguments, + }, + } + + +class WatsonxChatModelConnection(BaseChatModelConnection): + """Connection to the IBM watsonx.ai chat API.""" + + url: str = Field(description="The watsonx.ai service endpoint.") + api_key: str | None = Field(default=None, description="The IBM Cloud API key.") + token: str | None = Field( + default=None, description="A bearer token, as an alternative to api_key." + ) + project_id: str | None = Field( + default=None, description="The watsonx.ai project id." + ) + space_id: str | None = Field( + default=None, description="The watsonx.ai deployment space id." + ) + request_timeout: float = Field( + default=DEFAULT_REQUEST_TIMEOUT, + description="The timeout, in seconds, for chat requests to watsonx.ai.", + gt=0, + allow_inf_nan=False, + ) + max_retries: int = Field( + default=DEFAULT_MAX_RETRIES, + description="Maximum number of retries for transient failures.", + ge=0, + ) + + _client: APIClient | None = PrivateAttr(default=None) + _http_client: httpx.Client | None = PrivateAttr(default=None) + _models: Dict[str, ModelInference] = PrivateAttr(default_factory=dict) + + def __init__( + self, + *, + url: str | None = None, + api_key: str | None = None, + token: str | None = None, + project_id: str | None = None, + space_id: str | None = None, + request_timeout: float = DEFAULT_REQUEST_TIMEOUT, + max_retries: int = DEFAULT_MAX_RETRIES, + **kwargs: Any, + ) -> None: + """Initialize the connection.""" + resolved_url = _normalize(url) or _normalize(os.environ.get("WATSONX_URL")) + resolved_api_key = _normalize(api_key) or _normalize( + os.environ.get("WATSONX_API_KEY") + ) + resolved_token = _normalize(token) or _normalize( + os.environ.get("WATSONX_TOKEN") + ) + resolved_project_id = _normalize(project_id) or _normalize( + os.environ.get("WATSONX_PROJECT_ID") + ) + resolved_space_id = _normalize(space_id) or _normalize( + os.environ.get("WATSONX_SPACE_ID") + ) + + if not resolved_url: + msg = ( + "watsonx.ai url is not provided. Please pass it as an argument " + "or set the 'WATSONX_URL' environment variable." + ) + raise ValueError(msg) + if not resolved_api_key and not resolved_token: + msg = ( + "watsonx.ai credentials are not provided. Please pass 'api_key' " + "or 'token' as an argument, or set the 'WATSONX_API_KEY' or " + "'WATSONX_TOKEN' environment variable." + ) + raise ValueError(msg) + if resolved_api_key and resolved_token: + msg = ( + "watsonx.ai api_key and token cannot both be provided. Please configure " + "exactly one credential source." + ) + raise ValueError(msg) + if not resolved_project_id and not resolved_space_id: + msg = ( + "watsonx.ai project or space is not provided. Please pass " + "'project_id' or 'space_id' as an argument, or set the " + "'WATSONX_PROJECT_ID' or 'WATSONX_SPACE_ID' environment variable." + ) + raise ValueError(msg) + if resolved_project_id and resolved_space_id: + msg = ( + "watsonx.ai project and space cannot both be provided. Please configure " + "exactly one of 'project_id' or 'space_id'." + ) + raise ValueError(msg) + + super().__init__( + url=resolved_url, + api_key=resolved_api_key, + token=resolved_token, + project_id=resolved_project_id, + space_id=resolved_space_id, + request_timeout=request_timeout, + max_retries=max_retries, + **kwargs, + ) + + @property + def client(self) -> APIClient: + """Return the lazily initialized API client.""" + if self._client is None: + credential_kwargs: Dict[str, Any] = {"url": self.url} + if self.api_key: + credential_kwargs["api_key"] = self.api_key + if self.token: + credential_kwargs["token"] = self.token + self._http_client = httpx.Client(timeout=self.request_timeout) + self._client = APIClient( + credentials=Credentials(**credential_kwargs), + project_id=self.project_id, + space_id=self.space_id, + httpx_client=self._http_client, + ) + return self._client + + @override + def close(self) -> None: + """Close the underlying HTTP client.""" + self._models = {} + self._client = None + if self._http_client is not None: + with contextlib.suppress(Exception): + self._http_client.close() + self._http_client = None + + def _get_model(self, model: str) -> ModelInference: + if model not in self._models: + self._models[model] = ModelInference( + model_id=model, + api_client=self.client, + project_id=self.project_id, + space_id=self.space_id, + ) + return self._models[model] + + def _chat_with_retry(self, model_name: str, **chat_kwargs: Any) -> Dict[str, Any]: Review Comment: `_chat_with_retry` retries over `RETRYABLE_STATUS_CODES = {408, 429, 500, 502, 503, 504}` with `max_retries` defaulting to 3. The SDK already retries underneath: `ibm_watsonx_ai/_wrappers/httpx/rate_limited_retry/rate_limited_retry_decorator.py:40-44` sets `MAX_RETRIES = 10` over `(429, 503, 504, 520)`, wired into every chat POST at `base_model_inference.py:63-76`, same on the 1.3.42 pin. A sustained 429 then costs up to 44 HTTP requests and minutes of blocking `time.sleep` (`:327`) on a Flink async-executor thread for one record, with a third layer above both at `plan/actions/chat_model_action.py:304-360`. Java's `sendWithRetry` (`WatsonxChatModelConnection.java:290-325`) is the only layer there, so the same `max_retries=3` means about 3 retries on Java and up to about 43 on Python, while the `max_retries` row in `chat_models.md` describes both identically ("Maximum retries for transport failures and HTTP 408, 429, 500, 502, 503, and 504 responses"). `AGENTS.md` asks for semantic alignment, and here the key, the default and the docs all match while behavior differs by an order of magnitude. Was the SDK's built-in retry known about when the wrapper was written? `ModelInference.__init__` takes `max_retries`, `delay_time` and `retry_status_codes` directly (`model_inference.py:162-164`, both pins), so the configured policy could be handled by one layer with the same user-visible semantics. ########## integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java: ########## @@ -0,0 +1,678 @@ +/* + * 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.watsonx; + +import com.fasterxml.jackson.core.json.JsonReadFeature; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +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.BaseChatModelConnection; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.annotation.VisibleForTesting; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.net.URI; +import java.net.URLEncoder; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Function; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +/** Chat model connection for the IBM watsonx.ai text chat REST API. */ +public class WatsonxChatModelConnection extends BaseChatModelConnection { + + private static final Logger LOG = LoggerFactory.getLogger(WatsonxChatModelConnection.class); + + static final String DEFAULT_IAM_URL = "https://iam.cloud.ibm.com";; + static final String DEFAULT_API_VERSION = "2025-04-23"; + static final long DEFAULT_REQUEST_TIMEOUT_SEC = 120; + static final int DEFAULT_MAX_RETRIES = 3; + private static final Set<Integer> RETRYABLE_STATUS_CODES = Set.of(408, 429, 500, 502, 503, 504); + + private static final Set<String> CONTROL_PARAMS = + Set.of( + "model", + "tool_choice", + "tool_choice_option", + "extract_reasoning", + "additional_kwargs"); + private static final Set<String> RESERVED_ADDITIONAL_KWARGS = + Set.of( + "model", + "temperature", + "max_tokens", + "extract_reasoning", + "tool_choice", + "tool_choice_option"); + + private static final Pattern[] REASONING_PATTERNS = { + Pattern.compile("<think>(.*?)</think>", Pattern.DOTALL | Pattern.CASE_INSENSITIVE), + Pattern.compile("<analysis>(.*?)</analysis>", Pattern.DOTALL | Pattern.CASE_INSENSITIVE), + Pattern.compile("<reasoning>(.*?)</reasoning>", Pattern.DOTALL | Pattern.CASE_INSENSITIVE), + Pattern.compile( + "```(?:think|reasoning|thought)\\s*\\n(.*?)\\n```", + Pattern.DOTALL | Pattern.CASE_INSENSITIVE), + Pattern.compile( + "(?:^|\\n)Reasoning:\\s*(.*?)(?:\\n{2,}|$)", + Pattern.DOTALL | Pattern.CASE_INSENSITIVE), + }; + + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private static final ObjectMapper LENIENT_MAPPER = + JsonMapper.builder() + .enable(JsonReadFeature.ALLOW_SINGLE_QUOTES) + .enable(JsonReadFeature.ALLOW_UNQUOTED_FIELD_NAMES) + .build(); + + private final String url; + private final String apiKey; + private final String staticToken; + private final String projectId; + private final String spaceId; + private final String apiVersion; + private final String iamUrl; + private final Duration requestTimeout; + private final int maxRetries; + + private final HttpClient httpClient; + + private transient String cachedIamToken; + private transient long iamTokenExpirationEpochSec; + + public WatsonxChatModelConnection( + ResourceDescriptor descriptor, ResourceContext resourceContext) { + this(descriptor, resourceContext, System::getenv); + } + + @VisibleForTesting + WatsonxChatModelConnection( + ResourceDescriptor descriptor, + ResourceContext resourceContext, + Function<String, String> environmentLookup) { + super(descriptor, resourceContext); + + this.url = + trimTrailingSlash( + argumentOrEnv(descriptor, "url", "WATSONX_URL", environmentLookup)); + this.apiKey = argumentOrEnv(descriptor, "api_key", "WATSONX_API_KEY", environmentLookup); + this.staticToken = argumentOrEnv(descriptor, "token", "WATSONX_TOKEN", environmentLookup); + this.projectId = + argumentOrEnv(descriptor, "project_id", "WATSONX_PROJECT_ID", environmentLookup); + this.spaceId = argumentOrEnv(descriptor, "space_id", "WATSONX_SPACE_ID", environmentLookup); + + String apiVersion = normalize(descriptor.getArgument("api_version")); + this.apiVersion = apiVersion != null ? apiVersion : DEFAULT_API_VERSION; + String iamUrl = normalize(descriptor.getArgument("iam_url")); + this.iamUrl = trimTrailingSlash(iamUrl != null ? iamUrl : DEFAULT_IAM_URL); + Number requestTimeout = descriptor.getArgument("request_timeout"); + double requestTimeoutSeconds = + requestTimeout != null ? requestTimeout.doubleValue() : DEFAULT_REQUEST_TIMEOUT_SEC; + if (!Double.isFinite(requestTimeoutSeconds) || requestTimeoutSeconds <= 0) { + throw new IllegalArgumentException("request_timeout must be a positive finite number."); + } + this.requestTimeout = + Duration.ofMillis(Math.max(1L, Math.round(requestTimeoutSeconds * 1000.0))); + Number maxRetries = descriptor.getArgument("max_retries"); + this.maxRetries = + maxRetries != null + ? requireInteger(maxRetries, "max_retries", 0) + : DEFAULT_MAX_RETRIES; + + if (this.url == null || this.url.isEmpty()) { + throw new IllegalArgumentException( + "watsonx.ai url is not provided. Please pass the 'url' argument or set the" + + " 'WATSONX_URL' environment variable."); + } + if ((this.apiKey == null || this.apiKey.isEmpty()) + && (this.staticToken == null || this.staticToken.isEmpty())) { + throw new IllegalArgumentException( + "watsonx.ai credentials are not provided. Please pass the 'api_key' or 'token'" + + " argument, or set the 'WATSONX_API_KEY' or 'WATSONX_TOKEN'" + + " environment variable."); + } + if (this.apiKey != null && this.staticToken != null) { + throw new IllegalArgumentException( + "watsonx.ai api_key and token cannot both be provided. Please configure" + + " exactly one credential source."); + } + if ((this.projectId == null || this.projectId.isEmpty()) + && (this.spaceId == null || this.spaceId.isEmpty())) { + throw new IllegalArgumentException( + "watsonx.ai project or space is not provided. Please pass the 'project_id' or" + + " 'space_id' argument, or set the 'WATSONX_PROJECT_ID' or" + + " 'WATSONX_SPACE_ID' environment variable."); + } + if (this.projectId != null + && !this.projectId.isEmpty() + && this.spaceId != null + && !this.spaceId.isEmpty()) { + throw new IllegalArgumentException( + "watsonx.ai project and space cannot both be provided. Please configure" + + " exactly one of 'project_id' or 'space_id'."); + } + + this.httpClient = HttpClient.newBuilder().connectTimeout(Duration.ofSeconds(30)).build(); + } + + static int requireInteger(Number value, String argumentName, int minimum) { + double numericValue = value.doubleValue(); + if (!Double.isFinite(numericValue) + || numericValue != Math.rint(numericValue) + || numericValue < minimum + || numericValue > Integer.MAX_VALUE) { + throw new IllegalArgumentException( + argumentName + + " must be " + + (minimum == 0 ? "a non-negative" : "a positive") + + " integer."); + } + return (int) numericValue; + } + + private static String argumentOrEnv( + ResourceDescriptor descriptor, + String argumentName, + String envName, + Function<String, String> environmentLookup) { + String value = normalize(descriptor.getArgument(argumentName)); + if (value == null) { + value = normalize(environmentLookup.apply(envName)); + } + return value; + } + + private static String normalize(String value) { + if (value == null || value.isBlank()) { + return null; + } + return value.trim(); + } + + private static String trimTrailingSlash(String value) { + if (value != null && value.endsWith("/")) { + return value.substring(0, value.length() - 1); + } + return value; + } + + @Override + public ChatMessage chat( + List<ChatMessage> messages, List<Tool> tools, Map<String, Object> modelParams) { + try { + final String modelName = (String) modelParams.get("model"); + final boolean extractReasoning = + Boolean.TRUE.equals(modelParams.get("extract_reasoning")); + final ObjectNode payload = buildPayload(messages, tools, modelParams); + if (projectId != null && !projectId.isEmpty()) { + payload.put("project_id", projectId); + } else { + payload.put("space_id", spaceId); + } + + final HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create(url + "/ml/v1/text/chat?version=" + apiVersion)) + .timeout(requestTimeout) + .header("Authorization", "Bearer " + getBearerToken()) + .header("Content-Type", "application/json") + .header("Accept", "application/json") + .POST( + HttpRequest.BodyPublishers.ofString( + MAPPER.writeValueAsString(payload))) + .build(); + + final HttpResponse<String> response = sendWithRetry(request); + if (response.statusCode() / 100 != 2) { + throw new RuntimeException( + String.format( + "watsonx.ai chat request failed with status %d: %s", + response.statusCode(), response.body())); + } + + final ChatMessage chatMessage = + parseResponse(MAPPER.readTree(response.body()), modelName); + if (extractReasoning) { + final String[] parts = extractReasoning(chatMessage.getContent()); + chatMessage.setContent(parts[0]); + if (parts[1] != null) { + chatMessage.getExtraArgs().put("reasoning", parts[1]); + } + } + return chatMessage; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted while calling watsonx.ai.", e); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + /** + * Sends the request, retrying HTTP 408, 429, 500, 502, 503, and 504 responses and I/O errors up + * to {@code max_retries} times with capped exponential backoff, honoring {@code Retry-After}. + */ + private HttpResponse<String> sendWithRetry(HttpRequest request) + throws IOException, InterruptedException { + for (int attempt = 0; ; attempt++) { + try { + final HttpResponse<String> response = + httpClient.send(request, HttpResponse.BodyHandlers.ofString()); + if (attempt >= maxRetries || !isRetryableStatus(response.statusCode())) { + return response; + } + final long delayMillis = + retryDelayMillis( + attempt, response.headers().firstValue("Retry-After").orElse(null)); + LOG.warn( + "watsonx.ai request to {} returned status {}; retry {}/{} in {} ms", + request.uri().getPath(), + response.statusCode(), + attempt + 1, + maxRetries, + delayMillis); + Thread.sleep(delayMillis); + } catch (IOException e) { + if (attempt >= maxRetries) { + throw e; + } + final long delayMillis = retryDelayMillis(attempt, null); + LOG.warn( + "watsonx.ai request to {} failed ({}); retry {}/{} in {} ms", + request.uri().getPath(), + e.toString(), + attempt + 1, + maxRetries, + delayMillis); + Thread.sleep(delayMillis); + } + } + } + + @VisibleForTesting + static boolean isRetryableStatus(int status) { + return RETRYABLE_STATUS_CODES.contains(status); + } + + @VisibleForTesting + static long retryDelayMillis(int attempt, String retryAfterHeader) { + long backoffMillis = Math.min(1000L << attempt, 10_000L); + if (retryAfterHeader != null) { + try { + long retryAfterMillis = + Math.min(Long.parseLong(retryAfterHeader.trim()) * 1000L, 30_000L); + backoffMillis = Math.max(backoffMillis, retryAfterMillis); + } catch (NumberFormatException ignored) { + // Retry-After may be an HTTP date; fall back to exponential backoff. + } + } + return backoffMillis; + } + + @VisibleForTesting + static String[] extractReasoning(String content) { + if (content == null || content.isEmpty()) { + return new String[] {"", null}; + } + final List<String> reasoningChunks = new ArrayList<>(); + String cleaned = content; + for (Pattern pattern : REASONING_PATTERNS) { + final Matcher matcher = pattern.matcher(cleaned); + final StringBuilder rest = new StringBuilder(); + boolean found = false; + int position = 0; + while (matcher.find()) { + final String chunk = matcher.group(1).trim(); + if (!chunk.isEmpty()) { + reasoningChunks.add(chunk); + } + rest.append(cleaned, position, matcher.start()); + position = matcher.end(); + found = true; + } + if (found) { + rest.append(cleaned, position, cleaned.length()); + cleaned = rest.toString(); + } + } + final String reasoning = + reasoningChunks.isEmpty() ? null : String.join("\n\n", reasoningChunks); + cleaned = cleaned.replaceAll("\\n{3,}", "\n\n").replaceAll(" {2,}", " ").trim(); + return new String[] {cleaned, reasoning}; + } + + @VisibleForTesting + static ObjectNode buildPayload( + List<ChatMessage> messages, List<Tool> tools, Map<String, Object> modelParams) { + final ObjectNode payload = MAPPER.createObjectNode(); + payload.put("model_id", (String) modelParams.get("model")); + payload.set("messages", convertMessages(messages)); + + if (tools != null && !tools.isEmpty()) { + payload.set("tools", convertTools(tools)); + } + final Object toolChoice = modelParams.get("tool_choice"); + if (toolChoice != null) { + payload.set("tool_choice", MAPPER.valueToTree(toolChoice)); + } + final Object toolChoiceOption = modelParams.get("tool_choice_option"); + if (toolChoiceOption != null) { + payload.put("tool_choice_option", toolChoiceOption.toString()); + } + + @SuppressWarnings("unchecked") + final Map<String, Object> additionalKwargs = + (Map<String, Object>) modelParams.get("additional_kwargs"); + if (additionalKwargs != null) { + final Set<String> collisions = new java.util.HashSet<>(additionalKwargs.keySet()); + collisions.retainAll(RESERVED_ADDITIONAL_KWARGS); + if (!collisions.isEmpty()) { + throw new IllegalArgumentException( + "additional_kwargs must not contain reserved typed fields: " + + collisions + + ". Set these via the corresponding Setup field instead."); + } + additionalKwargs.forEach( + (key, value) -> { + if (value != null) { + payload.set(key, MAPPER.valueToTree(value)); + } + }); + } + + for (Map.Entry<String, Object> entry : modelParams.entrySet()) { + if (!CONTROL_PARAMS.contains(entry.getKey()) && entry.getValue() != null) { + payload.set(entry.getKey(), MAPPER.valueToTree(entry.getValue())); + } + } + return payload; + } + + /** + * Converts framework chat messages to the watsonx.ai (OpenAI-compatible) message format. + * + * <ul> + * <li>SYSTEM/USER messages carry {@code role} and {@code content}. + * <li>ASSISTANT messages may carry {@code tool_calls} with JSON string arguments. + * <li>TOOL messages carry {@code tool_call_id} referencing the original call, taken from the + * {@code externalId} entry of the message extra args. + * </ul> + */ + @VisibleForTesting + static ArrayNode convertMessages(List<ChatMessage> messages) { + final ArrayNode result = MAPPER.createArrayNode(); + for (ChatMessage message : messages) { + final ObjectNode node = MAPPER.createObjectNode(); + final MessageRole role = message.getRole(); + switch (role) { + case SYSTEM: + case USER: + node.put("role", role.name().toLowerCase()); + node.put("content", message.getContent()); + break; + case ASSISTANT: + node.put("role", "assistant"); + if (message.getContent() != null && !message.getContent().isEmpty()) { + node.put("content", message.getContent()); + } + final List<Map<String, Object>> toolCalls = message.getToolCalls(); + if (toolCalls != null && !toolCalls.isEmpty()) { + node.set("tool_calls", convertToolCalls(toolCalls)); + } + break; + case TOOL: + final Object externalId = message.getExtraArgs().get("externalId"); + if (externalId == null) { + throw new IllegalArgumentException( + "Tool message must have 'externalId' in extra args."); + } + node.put("role", "tool"); + node.put("content", message.getContent()); + node.put("tool_call_id", externalId.toString()); + break; + default: + throw new IllegalArgumentException("Unsupported role: " + role); + } + result.add(node); + } + return result; + } + + private static ArrayNode convertToolCalls(List<Map<String, Object>> toolCalls) { + final ArrayNode result = MAPPER.createArrayNode(); + for (Map<String, Object> toolCall : toolCalls) { + final Object originalId = toolCall.get("original_id"); + final Object id = originalId != null ? originalId : toolCall.get("id"); + if (id == null) { + throw new IllegalArgumentException( + "Tool call must have either 'original_id' or 'id' field."); + } + + @SuppressWarnings("unchecked") + final Map<String, Object> function = (Map<String, Object>) toolCall.get("function"); + final Object arguments = function.get("arguments"); + final String argumentsJson; + try { + argumentsJson = + arguments instanceof String + ? (String) arguments + : MAPPER.writeValueAsString(arguments); + } catch (Exception e) { + throw new RuntimeException(e); + } + + final ObjectNode node = MAPPER.createObjectNode(); + node.put("id", id.toString()); + node.put("type", "function"); + final ObjectNode functionNode = node.putObject("function"); + functionNode.put("name", (String) function.get("name")); + functionNode.put("arguments", argumentsJson); + result.add(node); + } + return result; + } + + @VisibleForTesting + static ArrayNode convertTools(List<Tool> tools) { + final ArrayNode result = MAPPER.createArrayNode(); + try { + for (Tool tool : tools) { + final ObjectNode node = MAPPER.createObjectNode(); + node.put("type", "function"); + final ObjectNode functionNode = node.putObject("function"); + functionNode.put("name", tool.getName()); + functionNode.put("description", tool.getDescription()); + functionNode.set( + "parameters", MAPPER.readTree(tool.getMetadata().getInputSchema())); + result.add(node); + } + return result; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + @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"); + if (finishReasonNode != null && !finishReasonNode.isNull()) { + final String finishReason = finishReasonNode.asText(); + if (!"stop".equals(finishReason) && !"tool_calls".equals(finishReason)) { + LOG.warn( + "watsonx.ai chat for model {} finished with reason '{}'; the response" + + " may be truncated or incomplete", + modelName, + finishReason); + } + } + + final JsonNode contentNode = responseMessage.get("content"); + final String content = + contentNode != null && !contentNode.isNull() ? contentNode.asText() : ""; + final ChatMessage chatMessage = ChatMessage.assistant(content); + + final JsonNode toolCallsNode = responseMessage.get("tool_calls"); + if (toolCallsNode != null && toolCallsNode.isArray() && !toolCallsNode.isEmpty()) { + final List<Map<String, Object>> toolCalls = new java.util.ArrayList<>(); + for (JsonNode toolCallNode : toolCallsNode) { + final String id = toolCallNode.required("id").asText(); + final JsonNode functionNode = toolCallNode.required("function"); + final Map<String, Object> arguments = + parseToolArguments(functionNode.get("arguments")); + toolCalls.add( + Map.of( + "id", + id, + "original_id", + id, + "type", + "function", + "function", + Map.of( + "name", + functionNode.required("name").asText(), + "arguments", + arguments))); + } + chatMessage.setToolCalls(toolCalls); + } + + final JsonNode usage = response.get("usage"); + if (modelName != null && !modelName.isBlank() && usage != null && !usage.isNull()) { + 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)); + chatMessage.setExtraArgs(extraArgs); + } + + return chatMessage; + } + + /** + * Parses model-emitted tool call arguments into a map, which is the format the framework's tool + * execution expects. + * + * <p>Models do not always return arguments as a clean JSON object string: some double-encode + * the JSON, and some (notably smaller models) emit single-quoted or unquoted pseudo-JSON. This + * method tolerates those variants and throws a descriptive error (including the raw value) when + * the arguments cannot be interpreted as an object. + */ + @VisibleForTesting + static Map<String, Object> parseToolArguments(JsonNode argumentsNode) { + if (argumentsNode == null || argumentsNode.isNull()) { + return Map.of(); + } + if (argumentsNode.isObject()) { + return MAPPER.convertValue(argumentsNode, new TypeReference<Map<String, Object>>() {}); + } + if (argumentsNode.isTextual()) { + String text = argumentsNode.asText().trim(); + if (text.isEmpty()) { + return Map.of(); + } + // Unwrap up to a few levels of string-encoding ("{\"a\": 1}" or "\"{\\\"a\\\": 1}\""). + for (int i = 0; i < 3; i++) { + final JsonNode parsed; + try { + parsed = LENIENT_MAPPER.readTree(text); + } catch (Exception e) { + break; + } + if (parsed.isObject()) { + return MAPPER.convertValue(parsed, new TypeReference<Map<String, Object>>() {}); + } + if (parsed.isTextual()) { + text = parsed.asText().trim(); + continue; + } + break; + } + } + throw new RuntimeException( + "Failed to parse tool call arguments returned by watsonx.ai as a JSON object: " + + argumentsNode); + } + + private synchronized String getBearerToken() { Review Comment: `cachedIamToken` is only ever replaced by the time check at `:641`. 401 isn't in `RETRYABLE_STATUS_CODES` (`:63`) and nothing in the non-2xx branch clears the cache, so if the token stops being accepted before its computed local expiry (key rotated, session invalidated, or the wall clock at `:639` jumping backwards) every subsequent `chat()` fails with a 401 until local time crosses the expiry. The normal-expiry path is right, with the 60s margin and `nowEpochSec` captured before the network call. `:636-638` returns a user-supplied `staticToken` forever. IBM Cloud IAM tokens live about an hour, so that path only works for jobs shorter than the TTL, and the docs describe `token` only as "Bearer token; configure exactly one of `api_key` and `token`". Python shares this, since its SDK docs say passing `token=` disables automatic refresh. How often the first one bites I genuinely don't know, since the 60s margin covers the common case. Does it make sense for a 401 to invalidate the cache, or is the time check meant to be the only trigger? One shape, if invalidating is the direction you'd want: clearing the cache and retrying once on 401 or 403 when `apiKey != null`, with a docs note that `token` doesn't refresh. ########## integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java: ########## @@ -0,0 +1,395 @@ +/* + * 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.watsonx; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +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.BaseChatModelConnection; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import java.util.List; +import java.util.Map; +import java.util.function.Function; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Unit tests for {@link WatsonxChatModelConnection}. These exercise the protocol-conversion logic + * with no network access, so they run in CI without any API key. + */ +class WatsonxChatModelConnectionTest { + + private static final ResourceContext NOOP = ResourceContext.fromGetResource((a, b) -> null); + private static final Function<String, String> NO_ENVIRONMENT = ignored -> null; + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private static ResourceDescriptor descriptor(String url, String apiKey, String projectId) { + ResourceDescriptor.Builder b = + ResourceDescriptor.Builder.newBuilder(WatsonxChatModelConnection.class.getName()); + if (url != null) { + b.addInitialArgument("url", url); + } + if (apiKey != null) { + b.addInitialArgument("api_key", apiKey); + } + if (projectId != null) { + b.addInitialArgument("project_id", projectId); + } + return b.build(); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("missingRequiredConfiguration") + void testConstructorRejectsMissingRequiredConfiguration( + String ignoredCaseName, + String url, + String apiKey, + String projectId, + String expectedMessage) { + assertThatThrownBy( + () -> + new WatsonxChatModelConnection( + descriptor(url, apiKey, projectId), NOOP, NO_ENVIRONMENT)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining(expectedMessage); + } + + private static Stream<Arguments> missingRequiredConfiguration() { + return Stream.of( + Arguments.of("missing url", null, "test-key", "test-project", "url"), + Arguments.of( + "missing credentials", + "https://us-south.ml.cloud.ibm.com";, + null, + "test-project", + "credentials"), + Arguments.of( + "missing project or space", + "https://us-south.ml.cloud.ibm.com";, + "test-key", + null, + "project or space")); + } + + @Test + @DisplayName("Constructor accepts space_id without project_id") + void testConstructorWithSpaceId() { + ResourceDescriptor descriptor = + ResourceDescriptor.Builder.newBuilder(WatsonxChatModelConnection.class.getName()) + .addInitialArgument("url", " https://us-south.ml.cloud.ibm.com ") + .addInitialArgument("api_key", " test-key ") + .addInitialArgument("space_id", " test-space ") + .build(); + + assertThat(new WatsonxChatModelConnection(descriptor, NOOP, NO_ENVIRONMENT)) + .isInstanceOf(BaseChatModelConnection.class); + } + + @Test + @DisplayName("Constructor rejects ambiguous scope and credentials") + void testConstructorRejectsAmbiguousConfiguration() { + ResourceDescriptor descriptor = + ResourceDescriptor.Builder.newBuilder(WatsonxChatModelConnection.class.getName()) + .addInitialArgument("url", "https://us-south.ml.cloud.ibm.com";) + .addInitialArgument("api_key", "test-key") + .addInitialArgument("project_id", "test-project") + .addInitialArgument("space_id", "test-space") + .build(); + + assertThatThrownBy(() -> new WatsonxChatModelConnection(descriptor, NOOP, NO_ENVIRONMENT)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("cannot both be provided") + .hasMessageContaining("exactly one"); + + ResourceDescriptor credentials = + ResourceDescriptor.Builder.newBuilder(WatsonxChatModelConnection.class.getName()) + .addInitialArgument("url", " https://us-south.ml.cloud.ibm.com ") + .addInitialArgument("api_key", " test-key ") + .addInitialArgument("token", " test-token ") + .addInitialArgument("project_id", " test-project ") + .build(); + assertThatThrownBy(() -> new WatsonxChatModelConnection(credentials, NOOP, NO_ENVIRONMENT)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("api_key and token") + .hasMessageContaining("exactly one"); + } + + @Test + @DisplayName("Request timeout accepts positive fractional seconds and rejects invalid values") + void testRequestTimeoutValidation() { + ResourceDescriptor fractionalTimeout = + ResourceDescriptor.Builder.newBuilder(WatsonxChatModelConnection.class.getName()) + .addInitialArgument("url", "https://us-south.ml.cloud.ibm.com";) + .addInitialArgument("api_key", "test-key") + .addInitialArgument("project_id", "test-project") + .addInitialArgument("request_timeout", 0.5) + .build(); + assertThat(new WatsonxChatModelConnection(fractionalTimeout, NOOP, NO_ENVIRONMENT)) + .isInstanceOf(BaseChatModelConnection.class); + + for (double invalidTimeout : List.of(0.0, -1.0, Double.NaN, Double.POSITIVE_INFINITY)) { + ResourceDescriptor invalid = + ResourceDescriptor.Builder.newBuilder( + WatsonxChatModelConnection.class.getName()) + .addInitialArgument("url", "https://us-south.ml.cloud.ibm.com";) + .addInitialArgument("api_key", "test-key") + .addInitialArgument("project_id", "test-project") + .addInitialArgument("request_timeout", invalidTimeout) + .build(); + assertThatThrownBy(() -> new WatsonxChatModelConnection(invalid, NOOP, NO_ENVIRONMENT)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("request_timeout"); + } + } + + @Test + @DisplayName("max_retries must be a non-negative integer") + void testMaxRetriesValidation() { + for (Number invalidMaxRetries : + new Number[] {-1, 0.9, Double.NaN, Double.POSITIVE_INFINITY}) { + ResourceDescriptor invalid = + ResourceDescriptor.Builder.newBuilder( + WatsonxChatModelConnection.class.getName()) + .addInitialArgument("url", "https://us-south.ml.cloud.ibm.com";) + .addInitialArgument("api_key", "test-key") + .addInitialArgument("project_id", "test-project") + .addInitialArgument("max_retries", invalidMaxRetries) + .build(); + assertThatThrownBy(() -> new WatsonxChatModelConnection(invalid, NOOP, NO_ENVIRONMENT)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("max_retries"); + } + } + + @Test + @DisplayName("System, user, assistant and tool messages convert to the watsonx format") + void testConvertMessages() { + ChatMessage assistant = new ChatMessage(MessageRole.ASSISTANT, ""); + assistant.setToolCalls( + List.of( + Map.of( + "id", "internal-uuid", + "original_id", "call_abc123", + "type", "function", + "function", + Map.of( + "name", + "add", + "arguments", + Map.of("a", 1, "b", 2))))); + ChatMessage toolResult = + new ChatMessage(MessageRole.TOOL, "3", Map.of("externalId", "call_abc123")); + + ArrayNode converted = + WatsonxChatModelConnection.convertMessages( + List.of( + new ChatMessage(MessageRole.SYSTEM, "You are helpful."), + new ChatMessage(MessageRole.USER, "What is 1 + 2?"), + assistant, + toolResult)); + + assertThat(converted).hasSize(4); + assertThat(converted.get(0).get("role").asText()).isEqualTo("system"); + assertThat(converted.get(0).get("content").asText()).isEqualTo("You are helpful."); + assertThat(converted.get(1).get("role").asText()).isEqualTo("user"); + + JsonNode assistantNode = converted.get(2); + assertThat(assistantNode.get("role").asText()).isEqualTo("assistant"); + assertThat(assistantNode.has("content")).isFalse(); + JsonNode toolCall = assistantNode.get("tool_calls").get(0); + assertThat(toolCall.get("id").asText()).isEqualTo("call_abc123"); + assertThat(toolCall.get("function").get("name").asText()).isEqualTo("add"); + // arguments must be serialized as a JSON string + assertThat(toolCall.get("function").get("arguments").isTextual()).isTrue(); + + JsonNode toolNode = converted.get(3); + assertThat(toolNode.get("role").asText()).isEqualTo("tool"); + assertThat(toolNode.get("tool_call_id").asText()).isEqualTo("call_abc123"); + assertThat(toolNode.get("content").asText()).isEqualTo("3"); + } + + @Test + @DisplayName("Tool message without externalId is rejected") + void testConvertToolMessageWithoutExternalId() { + assertThatThrownBy( + () -> + WatsonxChatModelConnection.convertMessages( + List.of(new ChatMessage(MessageRole.TOOL, "3")))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("externalId"); + } + + @Test + @DisplayName("Model params are copied top-level into the payload") + void testBuildPayload() { + ObjectNode payload = + WatsonxChatModelConnection.buildPayload( + List.of(new ChatMessage(MessageRole.USER, "Hello!")), + List.of(), + Map.of( + "model", + "ibm/granite-3-3-8b-instruct", + "temperature", + 0.5, + "max_tokens", + 256, + "extract_reasoning", + true, + "additional_kwargs", + Map.of("top_p", 0.9))); + + assertThat(payload.get("model_id").asText()).isEqualTo("ibm/granite-3-3-8b-instruct"); + assertThat(payload.get("temperature").asDouble()).isEqualTo(0.5); + assertThat(payload.get("max_tokens").asInt()).isEqualTo(256); + assertThat(payload.get("top_p").asDouble()).isEqualTo(0.9); + assertThat(payload.get("messages")).hasSize(1); + // framework control params must not leak into the request + assertThat(payload.has("model")).isFalse(); + assertThat(payload.has("extract_reasoning")).isFalse(); + assertThat(payload.has("tools")).isFalse(); + + assertThatThrownBy( + () -> + WatsonxChatModelConnection.buildPayload( + List.of(new ChatMessage(MessageRole.USER, "Hello!")), + List.of(), + Map.of( + "model", + "ibm/granite-3-3-8b-instruct", + "additional_kwargs", + Map.of("temperature", 5.0)))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("additional_kwargs") + .hasMessageContaining("temperature"); + } + + @Test + @DisplayName("Chat response with content and usage parses into a ChatMessage") + void testParseResponse() throws Exception { + JsonNode response = + MAPPER.readTree( + "{\"choices\": [{\"index\": 0, \"message\": {\"role\": \"assistant\"," + + " \"content\": \"Hello there!\"}, \"finish_reason\": \"stop\"}]," + + " \"usage\": {\"prompt_tokens\": 100, \"completion_tokens\": 50," + + " \"total_tokens\": 150}}"); + + ChatMessage message = + WatsonxChatModelConnection.parseResponse(response, "ibm/granite-3-3-8b-instruct"); + + assertThat(message.getRole()).isEqualTo(MessageRole.ASSISTANT); + assertThat(message.getContent()).isEqualTo("Hello there!"); + assertThat(message.getExtraArgs().get("model_name")) + .isEqualTo("ibm/granite-3-3-8b-instruct"); + assertThat(message.getExtraArgs().get("promptTokens")).isEqualTo(100L); + assertThat(message.getExtraArgs().get("completionTokens")).isEqualTo(50L); + } + + @Test + @DisplayName("Transient HTTP statuses are retryable and backoff honors Retry-After") + void testRetryPolicy() { + assertThat(WatsonxChatModelConnection.isRetryableStatus(408)).isTrue(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(429)).isTrue(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(500)).isTrue(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(502)).isTrue(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(503)).isTrue(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(504)).isTrue(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(200)).isFalse(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(400)).isFalse(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(401)).isFalse(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(501)).isFalse(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(505)).isFalse(); + + // exponential backoff, capped + assertThat(WatsonxChatModelConnection.retryDelayMillis(0, null)).isEqualTo(1000L); + assertThat(WatsonxChatModelConnection.retryDelayMillis(1, null)).isEqualTo(2000L); + assertThat(WatsonxChatModelConnection.retryDelayMillis(10, null)).isEqualTo(10_000L); + // Retry-After wins when larger, is capped, and non-numeric values are ignored + assertThat(WatsonxChatModelConnection.retryDelayMillis(0, "5")).isEqualTo(5000L); + assertThat(WatsonxChatModelConnection.retryDelayMillis(0, "600")).isEqualTo(30_000L); + assertThat(WatsonxChatModelConnection.retryDelayMillis(0, "not-a-number")).isEqualTo(1000L); + } Review Comment: `getBearerToken` (`WatsonxChatModelConnection.java:635-677`) has no test in either test class, and #902's rationale for a separate module is exactly that watsonx needs IAM token caching and refresh. Nothing verifies that a second `chat()` reuses the cached token, that it's re-fetched past the 60s margin, and so on. `sendWithRetry` (`:290-325`) is untested as a loop: the tests at `:316-337` cover only `isRetryableStatus` and `retryDelayMillis`, so nothing checks that a 503 is retried or a 400 isn't. `extractReasoning` (`:348-377`) is untested too, and as a pure `@VisibleForTesting` static it's the cheapest to cover. None of this needs a production seam: `url` and `iam_url` are ordinary connection arguments, so a test can point both at a local stub. The repo has that fixture already, `com.sun.net.httpserver.HttpServer` on port 0 at `runtime/src/test/java/org/apache/flink/agents/runtime/skill/URLSkillRepositoryTest.java:64-66`, JDK-provided with no new test dependency. One covers the IAM cache, the retry loop and the error branches at once. Is a stub server along those lines workable here, or does something about this setup rule it out? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
