weiqingy commented on code in PR #922:
URL: https://github.com/apache/flink-agents/pull/922#discussion_r3763168454
##########
python/flink_agents/api/chat_models/chat_model.py:
##########
@@ -216,7 +216,10 @@ def _extract_reasoning(
reasoning_chunks.extend(m.strip() for m in matches if
m.strip())
cleaned = pat.sub("", cleaned)
- reasoning = "\n\n".join(reasoning_chunks) if reasoning_chunks else None
+ if not reasoning_chunks:
+ return content, None
Review Comment:
Returning `content` here means an empty or whitespace-only block keeps its
tags: `<think></think>Answer` comes back unchanged at head, where before it
returned `Answer`. Being on the shared base, this reaches
`OllamaChatModelConnection` too (`ollama_chat_model.py:138`).
Would `return cleaned, None` work? `cleaned` still equals `content` when
nothing matched, so the no-tag case you're protecting is untouched.
##########
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 no-pattern case is fixed. One case still slips through though: the early
return hands back `content` rather than `cleaned`, so an empty or
whitespace-only block keeps its tags. `<think></think>Answer` returns
`<think></think>Answer` at head, where before it was `Answer`. I ran both
versions side by side to check.
Would returning `cleaned` work?
```java
if (reasoningChunks.isEmpty()) {
return new String[] {cleaned, null};
}
```
`cleaned` still equals `content` when nothing matched, so the case you just
fixed is unaffected.
The assertion at `:531-533` only covers the no-tag path. Worth an
empty-block case alongside it?
--
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]
