cansuk99 commented on code in PR #922: URL: https://github.com/apache/flink-agents/pull/922#discussion_r3700388491
########## 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: I have addressed this as per @wenjin272 comment https://github.com/apache/flink-agents/pull/922#discussion_r3698436566 -- 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]
