weiqingy commented on code in PR #922: URL: https://github.com/apache/flink-agents/pull/922#discussion_r3708918445
########## integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java: ########## @@ -0,0 +1,701 @@ +/* + * 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", + "model_id", + "messages", + "tools", + "project_id", + "space_id", + "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 String requestBody = MAPPER.writeValueAsString(payload); + String bearerToken = getBearerToken(); + HttpResponse<String> response = + sendWithRetry(buildChatRequest(requestBody, bearerToken)); + if ((response.statusCode() == 401 || response.statusCode() == 403) && apiKey != null) { + LOG.warn( + "watsonx.ai returned status {}; refreshing the cached IAM token and" + + " retrying once", + response.statusCode()); + invalidateCachedIamToken(bearerToken); + bearerToken = getBearerToken(); + response = sendWithRetry(buildChatRequest(requestBody, bearerToken)); + } + 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); + } + } + + private HttpRequest buildChatRequest(String requestBody, String bearerToken) { + return HttpRequest.newBuilder() + .uri(URI.create(url + "/ml/v1/text/chat?version=" + apiVersion)) + .timeout(requestTimeout) + .header("Authorization", "Bearer " + bearerToken) + .header("Content-Type", "application/json") + .header("Accept", "application/json") + .POST(HttpRequest.BodyPublishers.ofString(requestBody)) + .build(); + } + + /** + * 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(); Review Comment: The test added this round at `WatsonxChatModelConnectionTest.java:508-515` is named "Reasoning blocks are extracted without changing plain content", but that property doesn't hold. This line ends `extractReasoning` with `cleaned.replaceAll("\\n{3,}", "\n\n").replaceAll(" {2,}", " ").trim()`, and it runs every time, including when no reasoning pattern matched at all. The test only passes because its sample `"Plain answer"` has no double space in it. At head, `"| 1 | 2 |"` comes back as `"| 1 | 2 |"`, and a four-space indent collapses to one. `chat()` applies this to the whole response whenever `extract_reasoning` is set (`:273-279`), so content the extractor never touched gets reformatted too. The scope is narrow, since `extract_reasoning` defaults to `false` (`WatsonxChatModelSetup.java:56`). Could the cleanup be limited to where a reasoning block was actually removed, so the rest stays as the model wrote it? ########## python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py: ########## @@ -0,0 +1,491 @@ +################################################################################ +# 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.agents.types import OutputSchema +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", + "model_id", + "messages", + "tools", + "project_id", + "space_id", + "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, + max_retries=0, + ) + 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( + self, + messages: Sequence[ChatMessage], + tools: List[Tool] | None = None, + output_schema: OutputSchema | None = None, + **kwargs: Any, + ) -> ChatMessage: + """Process a sequence of messages, and return a response. + + A non-``None`` ``output_schema`` is rejected: this connection has no native + 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. + """ + self._reject_unsupported_output_schema(output_schema) + model_name = kwargs.pop("model", DEFAULT_MODEL) + extract_reasoning = bool(kwargs.pop("extract_reasoning", False)) + tool_choice = kwargs.pop("tool_choice", None) + tool_choice_option = kwargs.pop("tool_choice_option", None) + additional_kwargs = kwargs.pop("additional_kwargs", None) or {} + collisions = RESERVED_ADDITIONAL_KWARGS & additional_kwargs.keys() + if collisions: + msg = ( + "additional_kwargs must not contain reserved typed fields: " + f"{sorted(collisions)}. Set these via the corresponding Setup field instead." + ) + raise ValueError(msg) + + tool_specs: List[Dict[str, Any]] | None = ( + [to_openai_tool(metadata=tool.metadata) for tool in tools] + if tools + else None + ) + + response = self._chat_with_retry( + model_name, + messages=convert_to_watsonx_messages(messages), + tools=tool_specs, + tool_choice=tool_choice, + tool_choice_option=tool_choice_option, + params={**kwargs, **additional_kwargs} or None, Review Comment: Two things about this merge, and neither is reachable today, which is what makes them cheap to settle now. The precedence runs opposite in the two languages. Here `additional_kwargs` expands last, so it wins. Java spills `additional_kwargs` first (`WatsonxChatModelConnection.java:426-431`), then puts every non-control `modelParams` entry on top (`:434-438`), so there the per-call value wins. With `additional_kwargs={"top_p": 0.9}` and a caller passing `top_p=0.5`, Python sends `0.9` and Java sends `0.5` from the same configuration. I reproduced both at head. The new guard also covers only half of this dict. It checks `additional_kwargs` but not the sibling `**kwargs`, so `chat(..., model_id="X")` still reaches `params` and quietly replaces the model, since the SDK does `payload.update(parameters)` after seeding `model_id`. Nothing supplies per-call model parameters today: `ChatModelAction.java:365` passes `Map.of()` and `chat_model_action.py:320-325` passes none, both on `origin/main`. Which precedence did you have in mind? If per-call should win, reversing the order here would line Python up with Java, and running the guard over the merged dict would close the `**kwargs` half at the same time. The Java spill loop has the same unguarded shape (`CONTROL_PARAMS` at `:65-71`), so it may be worth a look in the same pass. ########## 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: Connector-level in both makes sense to me, for the reason you gave. It keeps the two languages on one retry set and one `Retry-After` policy. I checked that `max_retries=0` really does switch the SDK's own retry off, and it does on every version in the range pinned at `python/pyproject.toml:59-60`. Worth checking on more than the newest, since that code was reworked between releases. The one thing still open is that nothing tests it. `test_watsonx_chat_model.py` replaces `_get_model` outright (`:89-91`, `:146-148`, `:184-186`), so the `ModelInference` constructor is never seen, and deleting `max_retries=0` would leave the suite green. The five new reserved keys sit in the same spot, since `test_configuration_contract` (`:355-360`) only exercises `temperature`. Would asserting the constructor kwargs in that mocked test be enough to pin both, or would you rather keep this for a follow-up? ########## 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: The stub-server approach is right, and most of what you listed holds up when I ran it. The 401/403 test is the one gap: it can't fail on the bug it exists to catch. The chat handler only counts requests and never looks at the `Authorization` header (`WatsonxChatModelConnectionTest.java:444-452`), so it catches an invalidate that does nothing, but not a retry that re-sends the token the server just rejected. I checked by breaking it on purpose. Changing `WatsonxChatModelConnection.java:261` from `bearerToken = getBearerToken();` to `getBearerToken();` leaves all 20 tests green. Would recording the header on each request, then asserting the first was `Bearer token-1` and the second `Bearer token-2`, close that? -- 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]
