This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new 0e36efd3 [java][python] IBM watsonx.ai chat model integration (#922)
0e36efd3 is described below
commit 0e36efd34a7fd0543fdcc8ce233d7a87a8835f98
Author: Cansu Kertmen <[email protected]>
AuthorDate: Mon Aug 24 11:46:11 2026 +0200
[java][python] IBM watsonx.ai chat model integration (#922)
---
.../flink/agents/api/resource/ResourceName.java | 12 +
.../org/apache/flink/agents/api/yaml/Aliases.java | 4 +
.../apache/flink/agents/api/yaml/AliasesTest.java | 16 +
dist/pom.xml | 5 +
docs/content/docs/development/chat_models.md | 157 +++++
docs/content/docs/development/yaml.md | 1 +
docs/content/docs/faq/faq.md | 1 +
integrations/chat-models/pom.xml | 1 +
integrations/chat-models/watsonx/pom.xml | 64 ++
.../watsonx/WatsonxChatModelConnection.java | 711 +++++++++++++++++++++
.../chatmodels/watsonx/WatsonxChatModelSetup.java | 75 +++
.../watsonx/WatsonxChatModelConnectionTest.java | 617 ++++++++++++++++++
.../watsonx/WatsonxChatModelLiveTest.java | 86 +++
.../watsonx/WatsonxChatModelSetupTest.java | 114 ++++
python/flink_agents/api/chat_models/chat_model.py | 5 +-
python/flink_agents/api/resource.py | 8 +
python/flink_agents/api/yaml/aliases.py | 4 +
python/flink_agents/api/yaml/tests/test_aliases.py | 21 +-
.../chat_models/tests/test_ollama_chat_model.py | 2 +-
.../tests/test_output_schema_param_declared.py | 2 +
.../integrations/chat_models/watsonx/__init__.py | 17 +
.../chat_models/watsonx/tests/__init__.py | 17 +
.../watsonx/tests/test_watsonx_chat_model.py | 393 ++++++++++++
.../chat_models/watsonx/watsonx_chat_model.py | 498 +++++++++++++++
python/pyproject.toml | 2 +
25 files changed, 2830 insertions(+), 3 deletions(-)
diff --git
a/api/src/main/java/org/apache/flink/agents/api/resource/ResourceName.java
b/api/src/main/java/org/apache/flink/agents/api/resource/ResourceName.java
index 3382ee6a..19dc86de 100644
--- a/api/src/main/java/org/apache/flink/agents/api/resource/ResourceName.java
+++ b/api/src/main/java/org/apache/flink/agents/api/resource/ResourceName.java
@@ -95,6 +95,12 @@ public final class ResourceName {
public static final String VLLM_SETUP =
"org.apache.flink.agents.integrations.chatmodels.openai.VLLMChatModelSetup";
+ // IBM watsonx.ai
+ public static final String WATSONX_CONNECTION =
+
"org.apache.flink.agents.integrations.chatmodels.watsonx.WatsonxChatModelConnection";
+ public static final String WATSONX_SETUP =
+
"org.apache.flink.agents.integrations.chatmodels.watsonx.WatsonxChatModelSetup";
+
// Python Wrapper
public static final String PYTHON_WRAPPER_CONNECTION =
"org.apache.flink.agents.api.chat.model.python.PythonChatModelConnection";
@@ -140,6 +146,12 @@ public final class ResourceName {
public static final String VLLM_SETUP =
"flink_agents.integrations.chat_models.vllm.vllm_chat_model.VLLMChatModelSetup";
+ // IBM watsonx.ai
+ public static final String WATSONX_CONNECTION =
+
"flink_agents.integrations.chat_models.watsonx.watsonx_chat_model.WatsonxChatModelConnection";
+ public static final String WATSONX_SETUP =
+
"flink_agents.integrations.chat_models.watsonx.watsonx_chat_model.WatsonxChatModelSetup";
+
private Python() {}
}
diff --git a/api/src/main/java/org/apache/flink/agents/api/yaml/Aliases.java
b/api/src/main/java/org/apache/flink/agents/api/yaml/Aliases.java
index ec2bb306..ee70e3fe 100644
--- a/api/src/main/java/org/apache/flink/agents/api/yaml/Aliases.java
+++ b/api/src/main/java/org/apache/flink/agents/api/yaml/Aliases.java
@@ -90,6 +90,7 @@ public final class Aliases {
chatConnJava.put("azure_openai",
ResourceName.ChatModel.AZURE_OPENAI_CONNECTION);
chatConnJava.put("bedrock", ResourceName.ChatModel.BEDROCK_CONNECTION);
chatConnJava.put("vllm", ResourceName.ChatModel.VLLM_CONNECTION);
+ chatConnJava.put("watsonx", ResourceName.ChatModel.WATSONX_CONNECTION);
Map<String, String> chatConnPython = new HashMap<>();
chatConnPython.put("ollama",
ResourceName.ChatModel.Python.OLLAMA_CONNECTION);
chatConnPython.put("openai",
ResourceName.ChatModel.Python.OPENAI_COMPLETIONS_CONNECTION);
@@ -97,6 +98,7 @@ public final class Aliases {
chatConnPython.put("tongyi",
ResourceName.ChatModel.Python.TONGYI_CONNECTION);
chatConnPython.put("azure_openai",
ResourceName.ChatModel.Python.AZURE_OPENAI_CONNECTION);
chatConnPython.put("vllm",
ResourceName.ChatModel.Python.VLLM_CONNECTION);
+ chatConnPython.put("watsonx",
ResourceName.ChatModel.Python.WATSONX_CONNECTION);
ca.put(ResourceType.CHAT_MODEL_CONNECTION,
buildLangBuckets(chatConnJava, chatConnPython));
// CHAT_MODEL
@@ -109,6 +111,7 @@ public final class Aliases {
chatJava.put("azure_openai",
ResourceName.ChatModel.AZURE_OPENAI_SETUP);
chatJava.put("bedrock", ResourceName.ChatModel.BEDROCK_SETUP);
chatJava.put("vllm", ResourceName.ChatModel.VLLM_SETUP);
+ chatJava.put("watsonx", ResourceName.ChatModel.WATSONX_SETUP);
Map<String, String> chatPython = new HashMap<>();
chatPython.put("ollama", ResourceName.ChatModel.Python.OLLAMA_SETUP);
chatPython.put("openai",
ResourceName.ChatModel.Python.OPENAI_COMPLETIONS_SETUP);
@@ -116,6 +119,7 @@ public final class Aliases {
chatPython.put("tongyi", ResourceName.ChatModel.Python.TONGYI_SETUP);
chatPython.put("azure_openai",
ResourceName.ChatModel.Python.AZURE_OPENAI_SETUP);
chatPython.put("vllm", ResourceName.ChatModel.Python.VLLM_SETUP);
+ chatPython.put("watsonx", ResourceName.ChatModel.Python.WATSONX_SETUP);
ca.put(ResourceType.CHAT_MODEL, buildLangBuckets(chatJava,
chatPython));
// EMBEDDING_MODEL_CONNECTION
diff --git
a/api/src/test/java/org/apache/flink/agents/api/yaml/AliasesTest.java
b/api/src/test/java/org/apache/flink/agents/api/yaml/AliasesTest.java
index ab8d723e..9d4508c1 100644
--- a/api/src/test/java/org/apache/flink/agents/api/yaml/AliasesTest.java
+++ b/api/src/test/java/org/apache/flink/agents/api/yaml/AliasesTest.java
@@ -103,6 +103,22 @@ class AliasesTest {
assertThat(pythonSetup).isEqualTo(ResourceName.ChatModel.Python.VLLM_SETUP);
}
+ @Test
+ void watsonxAliasesResolveForJavaAndPython() {
+ assertThat(
+ Aliases.resolveClazz(
+ "watsonx", ResourceType.CHAT_MODEL_CONNECTION,
Language.JAVA))
+ .isEqualTo(ResourceName.ChatModel.WATSONX_CONNECTION);
+ assertThat(Aliases.resolveClazz("watsonx", ResourceType.CHAT_MODEL,
Language.JAVA))
+ .isEqualTo(ResourceName.ChatModel.WATSONX_SETUP);
+ assertThat(
+ Aliases.resolveClazz(
+ "watsonx", ResourceType.CHAT_MODEL_CONNECTION,
Language.PYTHON))
+ .isEqualTo(ResourceName.ChatModel.Python.WATSONX_CONNECTION);
+ assertThat(Aliases.resolveClazz("watsonx", ResourceType.CHAT_MODEL,
Language.PYTHON))
+ .isEqualTo(ResourceName.ChatModel.Python.WATSONX_SETUP);
+ }
+
@Test
void clazzAliasMissPassesThrough() {
String fqn =
diff --git a/dist/pom.xml b/dist/pom.xml
index 01c24343..7cb74295 100644
--- a/dist/pom.xml
+++ b/dist/pom.xml
@@ -87,6 +87,11 @@ under the License.
<artifactId>flink-agents-integrations-chat-models-gemini</artifactId>
<version>${project.version}</version>
</dependency>
+ <dependency>
+ <groupId>org.apache.flink</groupId>
+
<artifactId>flink-agents-integrations-chat-models-watsonx</artifactId>
+ <version>${project.version}</version>
+ </dependency>
<dependency>
<groupId>org.apache.flink</groupId>
<artifactId>flink-agents-integrations-embedding-models-ollama</artifactId>
diff --git a/docs/content/docs/development/chat_models.md
b/docs/content/docs/development/chat_models.md
index 8f9962aa..9152136f 100644
--- a/docs/content/docs/development/chat_models.md
+++ b/docs/content/docs/development/chat_models.md
@@ -1241,6 +1241,163 @@ public class MyAgent extends Agent {
A vLLM server serves the model(s) it was started with. Query `GET /v1/models`
on the server to list them; the `model` value in the setup must match one of
the returned names.
+### Watsonx (IBM watsonx.ai)
+
+IBM watsonx.ai provides cloud-based chat models, including the IBM Granite
series, with enterprise-grade governance and deployment options on IBM Cloud.
+
+#### Prerequisites
+
+1. Create an account on [IBM Cloud](https://cloud.ibm.com/) and provision a
[watsonx.ai](https://www.ibm.com/products/watsonx-ai) instance
+2. Create a watsonx.ai project or deployment space
+2. On Developer access, select your Project ID and create API Key
+4. Note the generated API Key, Project ID and watsonx.ai URL for configuring
your connection
+
+#### WatsonxChatModelConnection Parameters
+
+{{< tabs "WatsonxChatModelConnection Parameters" >}}
+
+{{< tab "Python" >}}
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `url` | str | `$WATSONX_URL` | watsonx.ai service endpoint of your region |
+| `api_key` | str | `$WATSONX_API_KEY` | IBM Cloud API key; configure exactly
one of `api_key` and `token` |
+| `token` | str | `$WATSONX_TOKEN` | Caller-provided bearer token; it is not
refreshed automatically and must remain valid for the job lifetime; configure
exactly one of `api_key` and `token` |
+| `project_id` | str | `$WATSONX_PROJECT_ID` | watsonx.ai project id (or use
`space_id`) |
+| `space_id` | str | `$WATSONX_SPACE_ID` | Deployment space id, as an
alternative to `project_id` |
+| `request_timeout` | float | `120.0` | HTTP request timeout in seconds |
+| `max_retries` | int | `3` | Maximum retries for transport failures and HTTP
408, 429, 500, 502, 503, and 504 responses |
+
+{{< /tab >}}
+
+{{< tab "Java" >}}
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `url` | String | `$WATSONX_URL` | watsonx.ai service endpoint of your region
|
+| `api_key` | String | `$WATSONX_API_KEY` | IBM Cloud API key; configure
exactly one of `api_key` and `token` |
+| `token` | String | `$WATSONX_TOKEN` | Caller-provided bearer token; it is
not refreshed automatically and must remain valid for the job lifetime;
configure exactly one of `api_key` and `token` |
+| `project_id` | String | `$WATSONX_PROJECT_ID` | watsonx.ai project id (or
use `space_id`) |
+| `space_id` | String | `$WATSONX_SPACE_ID` | Deployment space id, as an
alternative to `project_id` |
+| `api_version` | String | `"2025-04-23"` | watsonx.ai REST API version date |
+| `iam_url` | String | `"https://iam.cloud.ibm.com"` | IAM endpoint used to
exchange the API key for a token |
+| `request_timeout` | double | `120.0` | HTTP request timeout in seconds |
+| `max_retries` | int | `3` | Maximum retries for transport failures and HTTP
408, 429, 500, 502, 503, and 504 responses |
+
+{{< /tab >}}
+
+{{< /tabs >}}
+
+#### WatsonxChatModelSetup Parameters
+
+{{< tabs "WatsonxChatModelSetup Parameters" >}}
+
+{{< tab "Python" >}}
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `connection` | str | Required | Reference to connection method name |
+| `model` | str | `"ibm/granite-4-h-small"` | watsonx.ai model id to use |
+| `prompt` | Prompt \| str | None | Prompt template or reference to prompt
resource |
+| `tools` | List[str] | None | List of tool names available to the model |
+| `temperature` | float | `0.1` | Sampling temperature (0.0 to 2.0) |
+| `max_tokens` | int | None | Maximum number of tokens to generate |
+| `extract_reasoning` | bool | `False` | Extract reasoning content (e.g.
`<think>` blocks) from response |
+| `additional_kwargs` | dict | `{}` | Additional watsonx.ai chat parameters
(e.g. `top_p`, `time_limit`, `seed`) |
+
+{{< /tab >}}
+
+{{< tab "Java" >}}
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `connection` | String | Required | Reference to connection method name |
+| `model` | String | `"ibm/granite-4-h-small"` | watsonx.ai model id to use |
+| `prompt` | Prompt \| String | None | Prompt template or reference to prompt
resource |
+| `tools` | List[String] | None | List of tool names available to the model |
+| `temperature` | Double | `0.1` | Sampling temperature (0.0 to 2.0) |
+| `max_tokens` | Integer | None | Maximum number of tokens to generate |
+| `extract_reasoning` | Boolean | `false` | Extract reasoning content (e.g.
`<think>` blocks) from response |
+| `additional_kwargs` | Map | `{}` | Additional watsonx.ai chat parameters
(e.g. `top_p`, `time_limit`, `seed`) |
+
+{{< /tab >}}
+
+{{< /tabs >}}
+
+#### Usage Example
+
+{{< tabs "Watsonx Usage Example" >}}
+
+{{< tab "Python" >}}
+```python
+class MyAgent(Agent):
+
+ @chat_model_connection
+ @staticmethod
+ def watsonx_connection() -> ResourceDescriptor:
+ return ResourceDescriptor(
+ clazz=ResourceName.ChatModel.WATSONX_CONNECTION,
+ url="https://us-south.ml.cloud.ibm.com";, # set WATSONX_URL env var
+ api_key="your-api-key-here", # set WATSONX_API_KEY env var
+ project_id="your-project-id", # set WATSONX_PROJECT_ID env var
+ )
+
+ @chat_model_setup
+ @staticmethod
+ def watsonx_chat_model() -> ResourceDescriptor:
+ return ResourceDescriptor(
+ clazz=ResourceName.ChatModel.WATSONX_SETUP,
+ connection="watsonx_connection",
+ model="ibm/granite-4-h-small",
+ temperature=0.1,
+ max_tokens=1024
+ )
+
+ ...
+```
+{{< /tab >}}
+
+{{< tab "Java" >}}
+```java
+public class MyAgent extends Agent {
+ @ChatModelConnection
+ public static ResourceDescriptor watsonxConnection() {
+ return
ResourceDescriptor.Builder.newBuilder(ResourceName.ChatModel.WATSONX_CONNECTION)
+ .addInitialArgument("url", "https://us-south.ml.cloud.ibm.com";)
+ .addInitialArgument("api_key", "your-api-key-here")
+ .addInitialArgument("project_id", "your-project-id")
+ .build();
+ }
+
+ @ChatModelSetup
+ public static ResourceDescriptor watsonxChatModel() {
+ return
ResourceDescriptor.Builder.newBuilder(ResourceName.ChatModel.WATSONX_SETUP)
+ .addInitialArgument("connection", "watsonxConnection")
+ .addInitialArgument("model", "ibm/granite-4-h-small")
+ .addInitialArgument("temperature", 0.1)
+ .build();
+ }
+
+ ...
+}
+```
+{{< /tab >}}
+
+{{< /tabs >}}
+
+#### Available Models
+
+Visit the [watsonx.ai foundation models
documentation](https://www.ibm.com/products/watsonx-ai/foundation-models) for
the complete and up-to-date list of available chat models.
+
+Some popular options include:
+- **ibm/granite** series (ibm/granite-4-h-small, ibm/granite-3-3-8b-instruct)
+- **meta-llama** series (meta-llama/llama-3-3-70b-instruct)
+- **mistralai** series (mistralai/mistral-large)
+
+{{< hint warning >}}
+Model availability and specifications may change. Always check the official
IBM watsonx.ai documentation for the latest information before implementing in
production.
+{{< /hint >}}
+
## Using Cross-Language Providers
Flink Agents supports cross-language chat model integration, allowing you to
use chat models implemented in one language (Java or Python) from agents
written in the other language. This is particularly useful when a chat model
provider is only available in one language (e.g., Tongyi is currently
Python-only).
diff --git a/docs/content/docs/development/yaml.md
b/docs/content/docs/development/yaml.md
index b4837566..a053873b 100644
--- a/docs/content/docs/development/yaml.md
+++ b/docs/content/docs/development/yaml.md
@@ -558,6 +558,7 @@ Common chat-model aliases:
| `bedrock` | — | Bedrock (Java)
|
| `tongyi` | Tongyi (Python) | —
|
| `vllm` | vLLM (Python) | vLLM (Java)
|
+| `watsonx` | IBM watsonx.ai (Python) | IBM watsonx.ai (Java)
|
Embedding-model aliases (apply to both `embedding_model_connections` and
`embedding_model_setups`):
diff --git a/docs/content/docs/faq/faq.md b/docs/content/docs/faq/faq.md
index 2e0159a0..3f4fe812 100644
--- a/docs/content/docs/faq/faq.md
+++ b/docs/content/docs/faq/faq.md
@@ -105,6 +105,7 @@ Flink Agents provides built-in integrations for many
ecosystem providers. Some i
| [OpenAI]({{< ref "docs/development/chat_models#openai" >}}) | ✅ | ✅ |
| [Tongyi (DashScope)]({{< ref "docs/development/chat_models#tongyi-dashscope"
>}}) | ✅ | ❌ |
| [vLLM]({{< ref "docs/development/chat_models#vllm" >}}) | ✅ | ✅ |
+| [Watsonx (IBM watsonx.ai)]({{< ref
"docs/development/chat_models#watsonx-ibm-watsonxai" >}}) | ✅ | ✅ |
**Embedding Models**
diff --git a/integrations/chat-models/pom.xml b/integrations/chat-models/pom.xml
index d83457ae..37840ae7 100644
--- a/integrations/chat-models/pom.xml
+++ b/integrations/chat-models/pom.xml
@@ -36,6 +36,7 @@ under the License.
<module>gemini</module>
<module>ollama</module>
<module>openai</module>
+ <module>watsonx</module>
</modules>
</project>
diff --git a/integrations/chat-models/watsonx/pom.xml
b/integrations/chat-models/watsonx/pom.xml
new file mode 100644
index 00000000..864426ff
--- /dev/null
+++ b/integrations/chat-models/watsonx/pom.xml
@@ -0,0 +1,64 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<!--
+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.
+-->
+<project xmlns="http://maven.apache.org/POM/4.0.0";
+ xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance";
+ xsi:schemaLocation="http://maven.apache.org/POM/4.0.0
http://maven.apache.org/xsd/maven-4.0.0.xsd";>
+ <modelVersion>4.0.0</modelVersion>
+
+ <parent>
+ <groupId>org.apache.flink</groupId>
+ <artifactId>flink-agents-integrations-chat-models</artifactId>
+ <version>0.4-SNAPSHOT</version>
+ <relativePath>../pom.xml</relativePath>
+ </parent>
+
+ <artifactId>flink-agents-integrations-chat-models-watsonx</artifactId>
+ <name>Flink Agents : Integrations: Chat Models: IBM watsonx.ai</name>
+ <packaging>jar</packaging>
+
+ <!-- Talks to the watsonx.ai REST API with the JDK HttpClient. The official
+ IBM watsonx.ai Java SDK requires Java 17+, while this project targets
+ Java 11. -->
+ <dependencies>
+ <dependency>
+ <groupId>org.apache.flink</groupId>
+ <artifactId>flink-agents-api</artifactId>
+ <version>${project.version}</version>
+ </dependency>
+
+ <dependency>
+ <groupId>com.fasterxml.jackson.core</groupId>
+ <artifactId>jackson-databind</artifactId>
+ </dependency>
+
+ <dependency>
+ <groupId>org.apache.flink</groupId>
+ <artifactId>flink-annotations</artifactId>
+ <version>${flink.version}</version>
+ <scope>provided</scope>
+ </dependency>
+
+ <dependency>
+ <groupId>org.slf4j</groupId>
+ <artifactId>slf4j-api</artifactId>
+ <version>${slf4j.version}</version>
+ <scope>provided</scope>
+ </dependency>
+ </dependencies>
+
+</project>
diff --git
a/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
b/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
new file mode 100644
index 00000000..627ccd96
--- /dev/null
+++
b/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java
@@ -0,0 +1,711 @@
+/*
+ * 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> REQUEST_OWNED_PARAMS =
+ Set.of("model_id", "messages", "tools", "project_id", "space_id");
+ 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();
+ }
+ }
+ if (reasoningChunks.isEmpty()) {
+ return new String[] {content, null};
+ }
+ final String reasoning = 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));
+ }
+ });
+ }
+
+ final Set<String> collisions = new
java.util.HashSet<>(modelParams.keySet());
+ collisions.retainAll(REQUEST_OWNED_PARAMS);
+ if (!collisions.isEmpty()) {
+ throw new IllegalArgumentException(
+ "model parameters must not contain framework-owned fields:
" + collisions);
+ }
+ 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() {
+ if (staticToken != null && !staticToken.isEmpty()) {
+ return staticToken;
+ }
+ final long nowEpochSec = System.currentTimeMillis() / 1000;
+ // Refresh 60 seconds before the cached token expires.
+ if (cachedIamToken != null && nowEpochSec < iamTokenExpirationEpochSec
- 60) {
+ return cachedIamToken;
+ }
+
+ try {
+ final String form =
+
"grant_type=urn%3Aibm%3Aparams%3Aoauth%3Agrant-type%3Aapikey&apikey="
+ + URLEncoder.encode(apiKey,
StandardCharsets.UTF_8);
+ final HttpRequest request =
+ HttpRequest.newBuilder()
+ .uri(URI.create(iamUrl + "/identity/token"))
+ .timeout(Duration.ofSeconds(30))
+ .header("Content-Type",
"application/x-www-form-urlencoded")
+ .header("Accept", "application/json")
+ .POST(HttpRequest.BodyPublishers.ofString(form))
+ .build();
+ final HttpResponse<String> response = sendWithRetry(request);
+ if (response.statusCode() / 100 != 2) {
+ throw new RuntimeException(
+ String.format(
+ "IAM token request failed with status %d: %s",
+ response.statusCode(), response.body()));
+ }
+ final JsonNode tokenResponse = MAPPER.readTree(response.body());
+ cachedIamToken = tokenResponse.required("access_token").asText();
+ iamTokenExpirationEpochSec =
+ nowEpochSec +
tokenResponse.path("expires_in").asLong(3600);
+ return cachedIamToken;
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ throw new RuntimeException("Interrupted while obtaining a
watsonx.ai IAM token.", e);
+ } catch (RuntimeException e) {
+ throw e;
+ } catch (Exception e) {
+ throw new RuntimeException(e);
+ }
+ }
+
+ private synchronized void invalidateCachedIamToken(String rejectedToken) {
+ if (rejectedToken.equals(cachedIamToken)) {
+ cachedIamToken = null;
+ iamTokenExpirationEpochSec = 0;
+ }
+ }
+}
diff --git
a/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelSetup.java
b/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelSetup.java
new file mode 100644
index 00000000..01df5b09
--- /dev/null
+++
b/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelSetup.java
@@ -0,0 +1,75 @@
+/*
+ * 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 org.apache.flink.agents.api.chat.model.BaseChatModelSetup;
+import org.apache.flink.agents.api.resource.ResourceContext;
+import org.apache.flink.agents.api.resource.ResourceDescriptor;
+
+import java.util.HashMap;
+import java.util.Map;
+
+/** Chat model setup for IBM watsonx.ai. */
+public class WatsonxChatModelSetup extends BaseChatModelSetup {
+
+ public static final String DEFAULT_MODEL = "ibm/granite-4-h-small";
+ public static final double DEFAULT_TEMPERATURE = 0.1;
+
+ private final Double temperature;
+ private final Integer maxTokens;
+ private final boolean extractReasoning;
+ private final Map<String, Object> additionalKwargs;
+
+ public WatsonxChatModelSetup(ResourceDescriptor descriptor,
ResourceContext resourceContext) {
+ super(descriptor, resourceContext);
+ if (this.model == null || this.model.isEmpty()) {
+ this.model = DEFAULT_MODEL;
+ }
+ Number temperature = descriptor.getArgument("temperature");
+ this.temperature = temperature != null ? temperature.doubleValue() :
DEFAULT_TEMPERATURE;
+ if (!Double.isFinite(this.temperature)
+ || this.temperature < 0.0
+ || this.temperature > 2.0) {
+ throw new IllegalArgumentException(
+ "temperature must be a finite number between 0.0 and 2.0");
+ }
+ Number maxTokens = descriptor.getArgument("max_tokens");
+ this.maxTokens =
+ maxTokens != null
+ ? WatsonxChatModelConnection.requireInteger(maxTokens,
"max_tokens", 1)
+ : null;
+ this.extractReasoning =
Boolean.TRUE.equals(descriptor.getArgument("extract_reasoning"));
+ Map<String, Object> additionalKwargs =
descriptor.getArgument("additional_kwargs");
+ this.additionalKwargs = additionalKwargs != null ? additionalKwargs :
Map.of();
+ }
+
+ @Override
+ public Map<String, Object> getParameters() {
+ Map<String, Object> params = new HashMap<>();
+ params.put("model", model);
+ params.put("temperature", temperature);
+ params.put("extract_reasoning", extractReasoning);
+ if (maxTokens != null) {
+ params.put("max_tokens", maxTokens);
+ }
+ if (!additionalKwargs.isEmpty()) {
+ params.put("additional_kwargs", additionalKwargs);
+ }
+ return params;
+ }
+}
diff --git
a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
new file mode 100644
index 00000000..9b1cd5b6
--- /dev/null
+++
b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java
@@ -0,0 +1,617 @@
+/*
+ * 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 com.sun.net.httpserver.HttpExchange;
+import com.sun.net.httpserver.HttpServer;
+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 org.junit.jupiter.params.provider.ValueSource;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+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}. Request-level tests use
a local stub server,
+ * so the suite runs in CI without external network access or an 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 final String CHAT_RESPONSE =
+ "{\"choices\":[{\"index\":0,\"message\":{\"role\":\"assistant\","
+ + "\"content\":\"Hello!\"},\"finish_reason\":\"stop\"}]}";
+
+ 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();
+ }
+
+ private static ResourceDescriptor stubDescriptor(
+ String baseUrl, boolean useApiKey, int maxRetries) {
+ ResourceDescriptor.Builder builder =
+
ResourceDescriptor.Builder.newBuilder(WatsonxChatModelConnection.class.getName())
+ .addInitialArgument("url", baseUrl)
+ .addInitialArgument("project_id", "test-project")
+ .addInitialArgument("max_retries", maxRetries);
+ if (useApiKey) {
+ builder.addInitialArgument("api_key", "test-key");
+ builder.addInitialArgument("iam_url", baseUrl);
+ } else {
+ builder.addInitialArgument("token", "test-token");
+ }
+ return builder.build();
+ }
+
+ private static HttpServer startServer() throws IOException {
+ HttpServer server = HttpServer.create(new
InetSocketAddress("127.0.0.1", 0), 0);
+ server.setExecutor(null);
+ server.start();
+ return server;
+ }
+
+ private static String baseUrl(HttpServer server) {
+ return "http://127.0.0.1:"; + server.getAddress().getPort();
+ }
+
+ private static void sendJson(HttpExchange exchange, int status, String
body)
+ throws IOException {
+ byte[] bytes = body.getBytes(StandardCharsets.UTF_8);
+ exchange.getResponseHeaders().add("Content-Type", "application/json");
+ exchange.sendResponseHeaders(status, bytes.length);
+ exchange.getResponseBody().write(bytes);
+ exchange.close();
+ }
+
+ private static ChatMessage chat(WatsonxChatModelConnection connection) {
+ return connection.chat(
+ List.of(new ChatMessage(MessageRole.USER, "Hello!")),
+ List.of(),
+ Map.of("model", "ibm/granite-3-3-8b-instruct"));
+ }
+
+ @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,
+ "top_p",
+ 0.5,
+ "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.5);
+ 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");
+
+ for (String requestOwnedField :
+ List.of("model_id", "messages", "tools", "project_id",
"space_id")) {
+ assertThatThrownBy(
+ () ->
+ WatsonxChatModelConnection.buildPayload(
+ List.of(new
ChatMessage(MessageRole.USER, "Hello!")),
+ List.of(),
+ Map.of(
+ "model",
+
"ibm/granite-3-3-8b-instruct",
+ requestOwnedField,
+ "override")))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining(requestOwnedField);
+ }
+ }
+
+ @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("IAM token is cached and reused between chat requests")
+ void testIamTokenIsCached() throws Exception {
+ HttpServer server = startServer();
+ AtomicInteger iamRequests = new AtomicInteger();
+ AtomicInteger chatRequests = new AtomicInteger();
+ server.createContext(
+ "/identity/token",
+ exchange -> {
+ iamRequests.incrementAndGet();
+ sendJson(exchange, 200,
"{\"access_token\":\"token-1\",\"expires_in\":3600}");
+ });
+ server.createContext(
+ "/ml/v1/text/chat",
+ exchange -> {
+ chatRequests.incrementAndGet();
+ sendJson(exchange, 200, CHAT_RESPONSE);
+ });
+
+ try {
+ WatsonxChatModelConnection connection =
+ new WatsonxChatModelConnection(
+ stubDescriptor(baseUrl(server), true, 0), NOOP,
NO_ENVIRONMENT);
+ assertThat(chat(connection).getContent()).isEqualTo("Hello!");
+ assertThat(chat(connection).getContent()).isEqualTo("Hello!");
+ assertThat(iamRequests).hasValue(1);
+ assertThat(chatRequests).hasValue(2);
+ } finally {
+ server.stop(0);
+ }
+ }
+
+ @Test
+ @DisplayName("IAM token is refreshed inside the 60-second expiry margin")
+ void testIamTokenRefreshMargin() throws Exception {
+ HttpServer server = startServer();
+ AtomicInteger iamRequests = new AtomicInteger();
+ server.createContext(
+ "/identity/token",
+ exchange -> {
+ int tokenNumber = iamRequests.incrementAndGet();
+ sendJson(
+ exchange,
+ 200,
+ "{\"access_token\":\"token-" + tokenNumber +
"\",\"expires_in\":60}");
+ });
+ server.createContext(
+ "/ml/v1/text/chat", exchange -> sendJson(exchange, 200,
CHAT_RESPONSE));
+
+ try {
+ WatsonxChatModelConnection connection =
+ new WatsonxChatModelConnection(
+ stubDescriptor(baseUrl(server), true, 0), NOOP,
NO_ENVIRONMENT);
+ chat(connection);
+ chat(connection);
+ assertThat(iamRequests).hasValue(2);
+ } finally {
+ server.stop(0);
+ }
+ }
+
+ @ParameterizedTest
+ @ValueSource(ints = {401, 403})
+ @DisplayName("Rejected IAM token is refreshed and retried once")
+ void testRejectedIamTokenIsRefreshed(int rejectedStatus) throws Exception {
+ HttpServer server = startServer();
+ AtomicInteger iamRequests = new AtomicInteger();
+ AtomicInteger chatRequests = new AtomicInteger();
+ server.createContext(
+ "/identity/token",
+ exchange -> {
+ int tokenNumber = iamRequests.incrementAndGet();
+ sendJson(
+ exchange,
+ 200,
+ "{\"access_token\":\"token-" + tokenNumber +
"\",\"expires_in\":3600}");
+ });
+ server.createContext(
+ "/ml/v1/text/chat",
+ exchange -> {
+ int requestNumber = chatRequests.incrementAndGet();
+ sendJson(
+ exchange,
+ requestNumber == 1 ? rejectedStatus : 200,
+ requestNumber == 1 ? "{}" : CHAT_RESPONSE);
+ });
+
+ try {
+ WatsonxChatModelConnection connection =
+ new WatsonxChatModelConnection(
+ stubDescriptor(baseUrl(server), true, 0), NOOP,
NO_ENVIRONMENT);
+ assertThat(chat(connection).getContent()).isEqualTo("Hello!");
+ assertThat(iamRequests).hasValue(2);
+ assertThat(chatRequests).hasValue(2);
+ } finally {
+ server.stop(0);
+ }
+ }
+
+ @Test
+ @DisplayName("Retryable response is retried while a client error is not")
+ void testRetryLoop() throws Exception {
+ HttpServer server = startServer();
+ AtomicInteger chatRequests = new AtomicInteger();
+ server.createContext(
+ "/ml/v1/text/chat",
+ exchange -> {
+ int requestNumber = chatRequests.incrementAndGet();
+ sendJson(
+ exchange,
+ requestNumber == 1 ? 503 : 200,
+ requestNumber == 1 ? "{}" : CHAT_RESPONSE);
+ });
+
+ try {
+ WatsonxChatModelConnection retryingConnection =
+ new WatsonxChatModelConnection(
+ stubDescriptor(baseUrl(server), false, 1), NOOP,
NO_ENVIRONMENT);
+
assertThat(chat(retryingConnection).getContent()).isEqualTo("Hello!");
+ assertThat(chatRequests).hasValue(2);
+
+ server.removeContext("/ml/v1/text/chat");
+ chatRequests.set(0);
+ server.createContext(
+ "/ml/v1/text/chat",
+ exchange -> {
+ chatRequests.incrementAndGet();
+ sendJson(exchange, 400, "{\"error\":\"bad request\"}");
+ });
+ WatsonxChatModelConnection nonRetryingConnection =
+ new WatsonxChatModelConnection(
+ stubDescriptor(baseUrl(server), false, 3), NOOP,
NO_ENVIRONMENT);
+ assertThatThrownBy(() -> chat(nonRetryingConnection))
+ .isInstanceOf(RuntimeException.class)
+ .hasMessageContaining("status 400");
+ assertThat(chatRequests).hasValue(1);
+ } finally {
+ server.stop(0);
+ }
+ }
+
+ @Test
+ @DisplayName("Reasoning blocks are extracted without changing plain
content")
+ void testExtractReasoning() {
+
assertThat(WatsonxChatModelConnection.extractReasoning("<think>Plan</think>\nAnswer"))
+ .containsExactly("Answer", "Plan");
+ String plainContent = "| 1 | 2 |\n\n indented";
+ assertThat(WatsonxChatModelConnection.extractReasoning(plainContent))
+ .containsExactly(plainContent, null);
+ }
+
+ @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);
+ }
+
+ @Test
+ @DisplayName("Tool arguments in messy model-emitted formats are parsed
into a map")
+ void testParseToolArguments() throws Exception {
+ // clean JSON object string
+ assertThat(
+ WatsonxChatModelConnection.parseToolArguments(
+ MAPPER.readTree("\"{\\\"a\\\": 1, \\\"b\\\":
2}\"")))
+ .isEqualTo(Map.of("a", 1, "b", 2));
+ // double-encoded JSON string
+ assertThat(
+ WatsonxChatModelConnection.parseToolArguments(
+ MAPPER.readTree("\"\\\"{\\\\\\\"a\\\\\\\":
1}\\\"\"")))
+ .isEqualTo(Map.of("a", 1));
+ // single-quoted pseudo-JSON
+ assertThat(
+ WatsonxChatModelConnection.parseToolArguments(
+ MAPPER.readTree("\"{'a': 17, 'b': 25}\"")))
+ .isEqualTo(Map.of("a", 17, "b", 25));
+ // already an object node
+
assertThat(WatsonxChatModelConnection.parseToolArguments(MAPPER.readTree("{\"a\":
1}")))
+ .isEqualTo(Map.of("a", 1));
+ // missing or empty -> empty map
+
assertThat(WatsonxChatModelConnection.parseToolArguments(null)).isEmpty();
+
assertThat(WatsonxChatModelConnection.parseToolArguments(MAPPER.readTree("\"\"")))
+ .isEmpty();
+ // garbage -> descriptive error carrying the raw value
+ assertThatThrownBy(
+ () ->
+ WatsonxChatModelConnection.parseToolArguments(
+ MAPPER.readTree("\"not json at
all\"")))
+ .isInstanceOf(RuntimeException.class)
+ .hasMessageContaining("not json at all");
+ }
+
+ @Test
+ @DisplayName("Chat response with tool calls parses arguments and preserves
the original id")
+ void testParseResponseWithToolCalls() throws Exception {
+ JsonNode response =
+ MAPPER.readTree(
+ "{\"choices\": [{\"index\": 0, \"message\": {\"role\":
\"assistant\","
+ + " \"tool_calls\": [{\"id\": \"call_abc123\",
\"type\":"
+ + " \"function\", \"function\": {\"name\":
\"add\", \"arguments\":"
+ + " \"{\\\"a\\\": 1, \\\"b\\\": 2}\"}}]},
\"finish_reason\":"
+ + " \"tool_calls\"}]}");
+
+ ChatMessage message =
WatsonxChatModelConnection.parseResponse(response, null);
+
+ assertThat(message.getToolCalls()).hasSize(1);
+ Map<String, Object> toolCall = message.getToolCalls().get(0);
+ assertThat(toolCall.get("id")).isEqualTo("call_abc123");
+ assertThat(toolCall.get("original_id")).isEqualTo("call_abc123");
+ @SuppressWarnings("unchecked")
+ Map<String, Object> function = (Map<String, Object>)
toolCall.get("function");
+ assertThat(function.get("name")).isEqualTo("add");
+ assertThat(function.get("arguments")).isEqualTo(Map.of("a", 1, "b",
2));
+ }
+}
diff --git
a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelLiveTest.java
b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelLiveTest.java
new file mode 100644
index 00000000..50197e18
--- /dev/null
+++
b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelLiveTest.java
@@ -0,0 +1,86 @@
+/*
+ * 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 org.apache.flink.agents.api.chat.messages.ChatMessage;
+import org.apache.flink.agents.api.chat.messages.MessageRole;
+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.api.condition.EnabledIf;
+
+import java.util.List;
+import java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Live tests for {@link WatsonxChatModelConnection} against the real
watsonx.ai service.
+ *
+ * <p>The test requires {@code WATSONX_URL}, either {@code WATSONX_API_KEY} or
{@code
+ * WATSONX_TOKEN}, and either {@code WATSONX_PROJECT_ID} or {@code
WATSONX_SPACE_ID}. Override the
+ * model with {@code WATSONX_CHAT_MODEL} if the default is not available in
your region.
+ */
+@EnabledIf("credentialsAvailable")
+class WatsonxChatModelLiveTest {
+
+ private static final ResourceContext NOOP =
ResourceContext.fromGetResource((a, b) -> null);
+
+ static boolean credentialsAvailable() {
+ return isSet("WATSONX_URL")
+ && (isSet("WATSONX_API_KEY") || isSet("WATSONX_TOKEN"))
+ && (isSet("WATSONX_PROJECT_ID") || isSet("WATSONX_SPACE_ID"));
+ }
+
+ private static boolean isSet(String name) {
+ String value = System.getenv(name);
+ return value != null && !value.isBlank();
+ }
+
+ private static String model() {
+ String model = System.getenv("WATSONX_CHAT_MODEL");
+ return model != null ? model : WatsonxChatModelSetup.DEFAULT_MODEL;
+ }
+
+ private static WatsonxChatModelConnection connection() {
+ return new WatsonxChatModelConnection(
+
ResourceDescriptor.Builder.newBuilder(WatsonxChatModelConnection.class.getName())
+ .build(),
+ NOOP);
+ }
+
+ @Test
+ @DisplayName("Basic chat returns a non-empty assistant message")
+ void testBasicChat() {
+ ChatMessage response =
+ connection()
+ .chat(
+ List.of(
+ new ChatMessage(
+ MessageRole.USER,
+ "Say hello in one short
sentence.")),
+ List.of(),
+ Map.of("model", model(), "max_tokens", 100));
+
+ assertThat(response.getRole()).isEqualTo(MessageRole.ASSISTANT);
+ assertThat(response.getContent()).isNotBlank();
+ assertThat(response.getExtraArgs().get("promptTokens")).isNotNull();
+
assertThat(response.getExtraArgs().get("completionTokens")).isNotNull();
+ }
+}
diff --git
a/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelSetupTest.java
b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelSetupTest.java
new file mode 100644
index 00000000..e9cafeb3
--- /dev/null
+++
b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelSetupTest.java
@@ -0,0 +1,114 @@
+/*
+ * 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 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 java.util.Map;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
+
+/** Unit tests for {@link WatsonxChatModelSetup}. */
+class WatsonxChatModelSetupTest {
+
+ private static final ResourceContext NOOP =
ResourceContext.fromGetResource((a, b) -> null);
+
+ @Test
+ @DisplayName("Default model applies when the model argument is omitted")
+ void testDefaultModel() {
+ WatsonxChatModelSetup setup =
+ new WatsonxChatModelSetup(
+
ResourceDescriptor.Builder.newBuilder(WatsonxChatModelSetup.class.getName())
+ .addInitialArgument("connection", "watsonx")
+ .build(),
+ NOOP);
+
+ Map<String, Object> params = setup.getParameters();
+
assertThat(params.get("model")).isEqualTo(WatsonxChatModelSetup.DEFAULT_MODEL);
+
assertThat(params.get("temperature")).isEqualTo(WatsonxChatModelSetup.DEFAULT_TEMPERATURE);
+ assertThat(params).doesNotContainKey("max_tokens");
+ }
+
+ @Test
+ @DisplayName("Configured model, temperature and max_tokens are exposed as
parameters")
+ void testConfiguredParameters() {
+ WatsonxChatModelSetup setup =
+ new WatsonxChatModelSetup(
+
ResourceDescriptor.Builder.newBuilder(WatsonxChatModelSetup.class.getName())
+ .addInitialArgument("connection", "watsonx")
+ .addInitialArgument("model",
"ibm/granite-3-3-8b-instruct")
+ .addInitialArgument("temperature", 0.2)
+ .addInitialArgument("max_tokens", 512)
+ .addInitialArgument("extract_reasoning", true)
+ .addInitialArgument(
+ "additional_kwargs", Map.of("top_p",
0.9, "seed", 7))
+ .build(),
+ NOOP);
+
+ Map<String, Object> params = setup.getParameters();
+
assertThat(params.get("model")).isEqualTo("ibm/granite-3-3-8b-instruct");
+ assertThat(params.get("temperature")).isEqualTo(0.2);
+ assertThat(params.get("max_tokens")).isEqualTo(512);
+ assertThat(params.get("extract_reasoning")).isEqualTo(true);
+ assertThat(params.get("additional_kwargs")).isEqualTo(Map.of("top_p",
0.9, "seed", 7));
+ }
+
+ @Test
+ @DisplayName("Out-of-range and non-finite temperatures are rejected")
+ void testTemperatureValidation() {
+ for (double invalidTemperature :
+ new double[] {-0.1, 2.5, Double.NaN,
Double.POSITIVE_INFINITY}) {
+ assertThatThrownBy(
+ () ->
+ new WatsonxChatModelSetup(
+
ResourceDescriptor.Builder.newBuilder(
+
WatsonxChatModelSetup.class.getName())
+
.addInitialArgument("connection", "watsonx")
+ .addInitialArgument(
+ "temperature",
invalidTemperature)
+ .build(),
+ NOOP))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("temperature");
+ }
+ }
+
+ @Test
+ @DisplayName("max_tokens must be a positive integer")
+ void testMaxTokensValidation() {
+ for (Number invalidMaxTokens :
+ new Number[] {0, -1, 1.5, Double.NaN,
Double.POSITIVE_INFINITY}) {
+ assertThatThrownBy(
+ () ->
+ new WatsonxChatModelSetup(
+
ResourceDescriptor.Builder.newBuilder(
+
WatsonxChatModelSetup.class.getName())
+
.addInitialArgument("connection", "watsonx")
+ .addInitialArgument(
+ "max_tokens",
invalidMaxTokens)
+ .build(),
+ NOOP))
+ .isInstanceOf(IllegalArgumentException.class)
+ .hasMessageContaining("max_tokens");
+ }
+ }
+}
diff --git a/python/flink_agents/api/chat_models/chat_model.py
b/python/flink_agents/api/chat_models/chat_model.py
index c3a9d786..cb1c663b 100644
--- a/python/flink_agents/api/chat_models/chat_model.py
+++ b/python/flink_agents/api/chat_models/chat_model.py
@@ -217,7 +217,10 @@ class BaseChatModelConnection(Resource, ABC):
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 cleaned, None
+
+ reasoning = "\n\n".join(reasoning_chunks)
cleaned = re.sub(r"\n{3,}", "\n\n", cleaned)
cleaned = re.sub(r" {2,}", " ", cleaned)
cleaned = cleaned.strip()
diff --git a/python/flink_agents/api/resource.py
b/python/flink_agents/api/resource.py
index 11542c5d..4dcc904f 100644
--- a/python/flink_agents/api/resource.py
+++ b/python/flink_agents/api/resource.py
@@ -266,6 +266,10 @@ class ResourceName:
VLLM_CONNECTION =
"flink_agents.integrations.chat_models.vllm.vllm_chat_model.VLLMChatModelConnection"
VLLM_SETUP =
"flink_agents.integrations.chat_models.vllm.vllm_chat_model.VLLMChatModelSetup"
+ # Watsonx
+ WATSONX_CONNECTION =
"flink_agents.integrations.chat_models.watsonx.watsonx_chat_model.WatsonxChatModelConnection"
+ WATSONX_SETUP =
"flink_agents.integrations.chat_models.watsonx.watsonx_chat_model.WatsonxChatModelSetup"
+
# Java Wrapper
JAVA_WRAPPER_CONNECTION = (
"flink_agents.api.chat_models.java_chat_model.JavaChatModelConnection"
@@ -308,6 +312,10 @@ class ResourceName:
VLLM_CONNECTION =
"org.apache.flink.agents.integrations.chatmodels.openai.VLLMChatModelConnection"
VLLM_SETUP =
"org.apache.flink.agents.integrations.chatmodels.openai.VLLMChatModelSetup"
+ # IBM watsonx.ai
+ WATSONX_CONNECTION =
"org.apache.flink.agents.integrations.chatmodels.watsonx.WatsonxChatModelConnection"
+ WATSONX_SETUP =
"org.apache.flink.agents.integrations.chatmodels.watsonx.WatsonxChatModelSetup"
+
class EmbeddingModel:
"""EmbeddingModel resource names."""
diff --git a/python/flink_agents/api/yaml/aliases.py
b/python/flink_agents/api/yaml/aliases.py
index 65dd5fe1..584ba2fd 100644
--- a/python/flink_agents/api/yaml/aliases.py
+++ b/python/flink_agents/api/yaml/aliases.py
@@ -67,6 +67,7 @@ CLAZZ_ALIASES: Dict[ResourceType, Dict[str, Dict[str, str]]]
= {
"tongyi": ResourceName.ChatModel.TONGYI_CONNECTION,
"azure_openai": ResourceName.ChatModel.AZURE_OPENAI_CONNECTION,
"vllm": ResourceName.ChatModel.VLLM_CONNECTION,
+ "watsonx": ResourceName.ChatModel.WATSONX_CONNECTION,
},
"java": {
"ollama": ResourceName.ChatModel.Java.OLLAMA_CONNECTION,
@@ -77,6 +78,7 @@ CLAZZ_ALIASES: Dict[ResourceType, Dict[str, Dict[str, str]]]
= {
"azure_openai":
ResourceName.ChatModel.Java.AZURE_OPENAI_CONNECTION,
"bedrock": ResourceName.ChatModel.Java.BEDROCK_CONNECTION,
"vllm": ResourceName.ChatModel.Java.VLLM_CONNECTION,
+ "watsonx": ResourceName.ChatModel.Java.WATSONX_CONNECTION,
},
},
ResourceType.CHAT_MODEL: {
@@ -87,6 +89,7 @@ CLAZZ_ALIASES: Dict[ResourceType, Dict[str, Dict[str, str]]]
= {
"tongyi": ResourceName.ChatModel.TONGYI_SETUP,
"azure_openai": ResourceName.ChatModel.AZURE_OPENAI_SETUP,
"vllm": ResourceName.ChatModel.VLLM_SETUP,
+ "watsonx": ResourceName.ChatModel.WATSONX_SETUP,
},
"java": {
"ollama": ResourceName.ChatModel.Java.OLLAMA_SETUP,
@@ -97,6 +100,7 @@ CLAZZ_ALIASES: Dict[ResourceType, Dict[str, Dict[str, str]]]
= {
"azure_openai": ResourceName.ChatModel.Java.AZURE_OPENAI_SETUP,
"bedrock": ResourceName.ChatModel.Java.BEDROCK_SETUP,
"vllm": ResourceName.ChatModel.Java.VLLM_SETUP,
+ "watsonx": ResourceName.ChatModel.Java.WATSONX_SETUP,
},
},
ResourceType.EMBEDDING_MODEL_CONNECTION: {
diff --git a/python/flink_agents/api/yaml/tests/test_aliases.py
b/python/flink_agents/api/yaml/tests/test_aliases.py
index dfb3b296..d7ef0bc0 100644
--- a/python/flink_agents/api/yaml/tests/test_aliases.py
+++ b/python/flink_agents/api/yaml/tests/test_aliases.py
@@ -28,7 +28,7 @@ from flink_agents.api.events.tool_event import (
ToolRequestEvent,
ToolResponseEvent,
)
-from flink_agents.api.resource import ResourceType
+from flink_agents.api.resource import ResourceName, ResourceType
from flink_agents.api.yaml.aliases import (
CLAZZ_ALIASES,
EVENT_ALIASES,
@@ -121,6 +121,25 @@ def test_resolve_clazz_dispatches_on_language() -> None:
assert py.startswith("flink_agents")
+def test_watsonx_aliases_resolve_for_java_and_python() -> None:
+ assert (
+ resolve_clazz("watsonx", ResourceType.CHAT_MODEL_CONNECTION, "python")
+ == ResourceName.ChatModel.WATSONX_CONNECTION
+ )
+ assert (
+ resolve_clazz("watsonx", ResourceType.CHAT_MODEL, "python")
+ == ResourceName.ChatModel.WATSONX_SETUP
+ )
+ assert (
+ resolve_clazz("watsonx", ResourceType.CHAT_MODEL_CONNECTION, "java")
+ == ResourceName.ChatModel.Java.WATSONX_CONNECTION
+ )
+ assert (
+ resolve_clazz("watsonx", ResourceType.CHAT_MODEL, "java")
+ == ResourceName.ChatModel.Java.WATSONX_SETUP
+ )
+
+
def test_resolve_clazz_default_language_is_python() -> None:
default = resolve_clazz("ollama", ResourceType.CHAT_MODEL_CONNECTION)
explicit = resolve_clazz("ollama", ResourceType.CHAT_MODEL_CONNECTION,
"python")
diff --git
a/python/flink_agents/integrations/chat_models/tests/test_ollama_chat_model.py
b/python/flink_agents/integrations/chat_models/tests/test_ollama_chat_model.py
index 1f3fd261..e9782649 100644
---
a/python/flink_agents/integrations/chat_models/tests/test_ollama_chat_model.py
+++
b/python/flink_agents/integrations/chat_models/tests/test_ollama_chat_model.py
@@ -136,7 +136,7 @@ def test_extract_think_tags() -> None:
assert reasoning == "This is just my thought process."
# Test with no think tags
- content = "This is a regular response without any thinking tags."
+ content = "| 1 | 2 |\n\n This is a regular response without thinking
tags."
cleaned, reasoning = OllamaChatModelConnection._extract_reasoning(content)
assert cleaned == content
assert reasoning is None
diff --git
a/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py
b/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py
index 9b814c3c..61da5d6f 100644
---
a/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py
+++
b/python/flink_agents/integrations/chat_models/tests/test_output_schema_param_declared.py
@@ -33,6 +33,7 @@ from flink_agents.integrations.chat_models import
ollama_chat_model, tongyi_chat
from flink_agents.integrations.chat_models.anthropic import
anthropic_chat_model
from flink_agents.integrations.chat_models.azure import azure_openai_chat_model
from flink_agents.integrations.chat_models.openai import openai_chat_model
+from flink_agents.integrations.chat_models.watsonx import watsonx_chat_model
from flink_agents.runtime.java import java_chat_model as
runtime_java_chat_model
# A class is only discoverable through __subclasses__() once it has been
imported.
@@ -48,6 +49,7 @@ _MODULES_DEFINING_CONNECTIONS = (
runtime_java_chat_model,
tongyi_chat_model,
tool_parameter_injection_agent,
+ watsonx_chat_model,
)
diff --git a/python/flink_agents/integrations/chat_models/watsonx/__init__.py
b/python/flink_agents/integrations/chat_models/watsonx/__init__.py
new file mode 100644
index 00000000..e154fadd
--- /dev/null
+++ b/python/flink_agents/integrations/chat_models/watsonx/__init__.py
@@ -0,0 +1,17 @@
+################################################################################
+# 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.
+#################################################################################
diff --git
a/python/flink_agents/integrations/chat_models/watsonx/tests/__init__.py
b/python/flink_agents/integrations/chat_models/watsonx/tests/__init__.py
new file mode 100644
index 00000000..e154fadd
--- /dev/null
+++ b/python/flink_agents/integrations/chat_models/watsonx/tests/__init__.py
@@ -0,0 +1,17 @@
+################################################################################
+# 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.
+#################################################################################
diff --git
a/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_chat_model.py
b/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_chat_model.py
new file mode 100644
index 00000000..57063f60
--- /dev/null
+++
b/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_chat_model.py
@@ -0,0 +1,393 @@
+################################################################################
+# 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 os
+from typing import Any, Dict
+from unittest.mock import MagicMock
+
+import pytest
+
+from flink_agents.api.chat_message import ChatMessage, MessageRole
+from flink_agents.api.resource import Resource, ResourceType
+from flink_agents.api.resource_context import ResourceContext
+from flink_agents.integrations.chat_models.watsonx.watsonx_chat_model import (
+ DEFAULT_MODEL,
+ WatsonxChatModelConnection,
+ WatsonxChatModelSetup,
+ convert_to_watsonx_messages,
+)
+
+test_model = os.environ.get("WATSONX_CHAT_MODEL", DEFAULT_MODEL)
+credentials_available = (
+ "WATSONX_URL" in os.environ
+ and ("WATSONX_API_KEY" in os.environ or "WATSONX_TOKEN" in os.environ)
+) and ("WATSONX_PROJECT_ID" in os.environ or "WATSONX_SPACE_ID" in os.environ)
+
+
+def _fake_connection(**kwargs: Any) -> WatsonxChatModelConnection:
+ """Create a connection with fake credentials for offline tests."""
+ return WatsonxChatModelConnection(
+ name="watsonx",
+ url=kwargs.pop("url", "https://us-south.ml.cloud.ibm.com";),
+ api_key=kwargs.pop("api_key", "fake-key"),
+ project_id=kwargs.pop("project_id", "fake-project"),
+ **kwargs,
+ )
+
+
[email protected]
[email protected](
+ not credentials_available, reason="watsonx.ai credentials are not set"
+)
+def test_watsonx_chat() -> None:
+ """Test basic chat functionality of WatsonxChatModelConnection."""
+ connection = WatsonxChatModelConnection(name="watsonx")
+ response = connection.chat(
+ [ChatMessage(role=MessageRole.USER, content="Hello!")],
model=test_model
+ )
+ assert response is not None
+ assert response.content is not None
+ assert response.content.strip() != ""
+ assert response.role == MessageRole.ASSISTANT
+
+
+def _mock_chat_response(
+ message: Dict[str, Any], finish_reason: str = "stop"
+) -> Dict[str, Any]:
+ return {
+ "id": "chatcmpl-1",
+ "model_id": test_model,
+ "choices": [{"index": 0, "message": message, "finish_reason":
finish_reason}],
+ "usage": {
+ "prompt_tokens": 100,
+ "completion_tokens": 50,
+ "total_tokens": 150,
+ },
+ }
+
+
+def test_watsonx_chat_mocked(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Test chat response handling and params passing (mock watsonx client)."""
+ mock_model = MagicMock()
+ mock_model.chat.return_value = _mock_chat_response(
+ {"role": "assistant", "content": "Hello there!"}
+ )
+ model_inference = MagicMock(return_value=mock_model)
+ monkeypatch.setattr(
+
"flink_agents.integrations.chat_models.watsonx.watsonx_chat_model.ModelInference",
+ model_inference,
+ )
+
+ connection = _fake_connection()
+ api_client = MagicMock()
+ connection._client = api_client
+
+ def get_resource(name: str, type: ResourceType) -> Resource:
+ return connection
+
+ mock_ctx = MagicMock(spec=ResourceContext)
+ mock_ctx.get_resource = get_resource
+
+ llm = WatsonxChatModelSetup(
+ name="watsonx",
+ model=test_model,
+ connection="watsonx",
+ temperature=0.5,
+ max_tokens=256,
+ additional_kwargs={"top_p": 0.9},
+ resource_context=mock_ctx,
+ )
+
+ llm.open()
+
+ response = llm.chat(
+ [ChatMessage(role=MessageRole.USER, content="Hello!")], top_p=0.5
+ )
+
+ mock_model.chat.assert_called_once()
+ call_kwargs = mock_model.chat.call_args.kwargs
+ assert call_kwargs["messages"] == [{"role": "user", "content": "Hello!"}]
+ assert call_kwargs["params"] == {
+ "temperature": 0.5,
+ "max_tokens": 256,
+ "top_p": 0.5,
+ }
+
+ assert response.role == MessageRole.ASSISTANT
+ assert response.content == "Hello there!"
+ assert response.extra_args["model_name"] == test_model
+ assert response.extra_args["promptTokens"] == 100
+ assert response.extra_args["completionTokens"] == 50
+ model_inference.assert_called_once_with(
+ model_id=test_model,
+ api_client=api_client,
+ project_id="fake-project",
+ space_id=None,
+ max_retries=0,
+ )
+
+
+def test_watsonx_tool_call_response_mocked(monkeypatch: pytest.MonkeyPatch) ->
None:
+ """Test that tool call responses are converted to the framework format."""
+ mock_model = MagicMock()
+ mock_model.chat.return_value = _mock_chat_response(
+ {
+ "role": "assistant",
+ "tool_calls": [
+ {
+ "id": "call_abc123",
+ "type": "function",
+ "function": {"name": "add", "arguments": '{"a": 1, "b":
2}'},
+ }
+ ],
+ }
+ )
+ monkeypatch.setattr(
+ WatsonxChatModelConnection, "_get_model", lambda self, model:
mock_model
+ )
+
+ connection = _fake_connection()
+ response = connection.chat(
+ [ChatMessage(role=MessageRole.USER, content="What is 1 + 2?")],
+ model=test_model,
+ )
+
+ assert len(response.tool_calls) == 1
+ tool_call = response.tool_calls[0]
+ assert tool_call["function"]["name"] == "add"
+ assert tool_call["function"]["arguments"] == {"a": 1, "b": 2}
+ assert tool_call["original_id"] == "call_abc123"
+
+
+def test_chat_retries_transient_failures(monkeypatch: pytest.MonkeyPatch) ->
None:
+ """Transient HTTP failures are retried up to max_retries, then succeed."""
+ import httpx
+ from ibm_watsonx_ai.wml_client_error import ApiRequestFailure
+
+ request = httpx.Request("POST", "https://test.invalid/ml/v1/text/chat";)
+ rate_limited = ApiRequestFailure(
+ "rate limited",
+ httpx.Response(
+ 429,
+ text="too many requests",
+ headers={"Retry-After": "5"},
+ request=request,
+ ),
+ )
+ mock_model = MagicMock()
+ mock_model.chat.side_effect = [
+ rate_limited,
+ rate_limited,
+ _mock_chat_response({"role": "assistant", "content": "Recovered!"}),
+ ]
+ monkeypatch.setattr(
+ WatsonxChatModelConnection, "_get_model", lambda self, model:
mock_model
+ )
+ sleep = MagicMock()
+ monkeypatch.setattr(
+
"flink_agents.integrations.chat_models.watsonx.watsonx_chat_model.time.sleep",
+ sleep,
+ )
+
+ connection = _fake_connection(max_retries=3)
+ response = connection.chat(
+ [ChatMessage(role=MessageRole.USER, content="Hello!")],
model=test_model
+ )
+
+ assert response.content == "Recovered!"
+ assert mock_model.chat.call_count == 3
+ assert [call.args[0] for call in sleep.call_args_list] == [5, 5]
+
+ # Non-retryable failures propagate immediately.
+ unauthorized = ApiRequestFailure(
+ "bad key", httpx.Response(401, text="unauthorized", request=request)
+ )
+ mock_model.chat.side_effect = [unauthorized]
+ mock_model.chat.reset_mock()
+ with pytest.raises(ApiRequestFailure):
+ connection.chat(
+ [ChatMessage(role=MessageRole.USER, content="Hello!")],
model=test_model
+ )
+ assert mock_model.chat.call_count == 1
+
+
+def test_connection_close() -> None:
+ """close() releases the HTTP client and cached models without errors."""
+ connection = _fake_connection()
+ connection.close() # closing before any request is a no-op
+
+ connection._http_client = MagicMock()
+ http_client = connection._http_client
+ connection.close()
+ http_client.close.assert_called_once()
+ assert connection._http_client is None
+ assert connection._models == {}
+
+
+def test_parse_tool_arguments_messy_formats() -> None:
+ """Tool arguments in messy model-emitted formats are parsed into a dict."""
+ from flink_agents.integrations.chat_models.watsonx.watsonx_chat_model
import (
+ _parse_tool_arguments,
+ )
+
+ assert _parse_tool_arguments('{"a": 1, "b": 2}') == {"a": 1, "b": 2}
+ # double-encoded JSON string
+ assert _parse_tool_arguments('"{\\"a\\": 1}"') == {"a": 1}
+ # single-quoted / Python-literal style dict
+ assert _parse_tool_arguments("{'a': 17, 'b': 25}") == {"a": 17, "b": 25}
+ # already a dict
+ assert _parse_tool_arguments({"a": 1}) == {"a": 1}
+ # missing or empty -> empty dict
+ assert _parse_tool_arguments(None) == {}
+ assert _parse_tool_arguments("") == {}
+ # garbage -> descriptive error carrying the raw value
+ with pytest.raises(TypeError, match="not json at all"):
+ _parse_tool_arguments("not json at all")
+
+
+def test_convert_to_watsonx_messages_round_trip() -> None:
+ """Test conversion of assistant tool calls and tool results to watsonx
format."""
+ messages = [
+ ChatMessage(role=MessageRole.SYSTEM, content="You are helpful."),
+ ChatMessage(role=MessageRole.USER, content="What is 1 + 2?"),
+ ChatMessage(
+ role=MessageRole.ASSISTANT,
+ tool_calls=[
+ {
+ "id": "internal-id",
+ "type": "function",
+ "function": {"name": "add", "arguments": {"a": 1, "b": 2}},
+ "original_id": "call_abc123",
+ }
+ ],
+ ),
+ ChatMessage(
+ role=MessageRole.TOOL,
+ content="3",
+ extra_args={"external_id": "call_abc123"},
+ ),
+ ]
+
+ converted = convert_to_watsonx_messages(messages)
+
+ assert converted[0] == {"role": "system", "content": "You are helpful."}
+ assert converted[1] == {"role": "user", "content": "What is 1 + 2?"}
+ assert converted[2] == {
+ "role": "assistant",
+ "tool_calls": [
+ {
+ "id": "call_abc123",
+ "type": "function",
+ "function": {"name": "add", "arguments": '{"a": 1, "b": 2}'},
+ }
+ ],
+ }
+ assert converted[3] == {
+ "role": "tool",
+ "content": "3",
+ "tool_call_id": "call_abc123",
+ }
+
+
+def test_configuration_contract(monkeypatch: pytest.MonkeyPatch) -> None:
+ """Validate required connection fields, scoping, timeout, and defaults."""
+ for var in (
+ "WATSONX_URL",
+ "WATSONX_API_KEY",
+ "WATSONX_TOKEN",
+ "WATSONX_PROJECT_ID",
+ "WATSONX_SPACE_ID",
+ ):
+ monkeypatch.delenv(var, raising=False)
+
+ with pytest.raises(ValueError, match="url"):
+ WatsonxChatModelConnection(name="watsonx")
+
+ with pytest.raises(ValueError, match="credentials"):
+ WatsonxChatModelConnection(
+ name="watsonx", url="https://us-south.ml.cloud.ibm.com";
+ )
+
+ with pytest.raises(ValueError, match="project or space"):
+ WatsonxChatModelConnection(
+ name="watsonx",
+ url="https://us-south.ml.cloud.ibm.com";,
+ api_key="fake-key",
+ )
+
+ connection = WatsonxChatModelConnection(
+ name="watsonx",
+ url=" https://us-south.ml.cloud.ibm.com ",
+ api_key=" fake-key ",
+ space_id=" fake-space ",
+ )
+
+ assert connection.url == "https://us-south.ml.cloud.ibm.com";
+ assert connection.api_key == "fake-key"
+ assert connection.project_id is None
+ assert connection.space_id == "fake-space"
+
+ with pytest.raises(ValueError, match=r"cannot both be provided.*exactly
one"):
+ WatsonxChatModelConnection(
+ name="watsonx",
+ url="https://us-south.ml.cloud.ibm.com";,
+ api_key="fake-key",
+ project_id="fake-project",
+ space_id="fake-space",
+ )
+
+ with pytest.raises(ValueError, match=r"api_key and token.*exactly one"):
+ WatsonxChatModelConnection(
+ name="watsonx",
+ url=" https://us-south.ml.cloud.ibm.com ",
+ api_key=" fake-key ",
+ token=" fake-token ",
+ project_id=" fake-project ",
+ )
+
+ for request_timeout in (0, -1, float("nan"), float("inf")):
+ with pytest.raises(ValueError, match="request_timeout"):
+ _fake_connection(request_timeout=request_timeout)
+
+ assert WatsonxChatModelSetup(connection="conn").model == DEFAULT_MODEL
+
+ with pytest.raises(ValueError, match="additional_kwargs"):
+ _fake_connection().chat(
+ [ChatMessage(role=MessageRole.USER, content="Hello!")],
+ model=test_model,
+ additional_kwargs={"temperature": 5.0},
+ )
+
+
[email protected](
+ "reserved_key", ["model_id", "messages", "tools", "project_id", "space_id"]
+)
+def test_additional_kwargs_reject_request_owned_fields(reserved_key: str) ->
None:
+ """Framework-owned request fields cannot be replaced by static
configuration."""
+ with pytest.raises(ValueError, match=reserved_key):
+ _fake_connection().chat(
+ [ChatMessage(role=MessageRole.USER, content="Hello!")],
+ model=test_model,
+ additional_kwargs={reserved_key: "override"},
+ )
+ if reserved_key not in {"messages", "tools"}:
+ with pytest.raises(ValueError, match=reserved_key):
+ _fake_connection().chat(
+ [ChatMessage(role=MessageRole.USER, content="Hello!")],
+ model=test_model,
+ **{reserved_key: "override"},
+ )
diff --git
a/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
b/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
new file mode 100644
index 00000000..6544d221
--- /dev/null
+++ b/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py
@@ -0,0 +1,498 @@
+################################################################################
+# 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})
+REQUEST_OWNED_PARAMS = frozenset(
+ {"model_id", "messages", "tools", "project_id", "space_id"}
+)
+RESERVED_ADDITIONAL_KWARGS = frozenset(
+ {
+ "model",
+ "temperature",
+ "max_tokens",
+ "extract_reasoning",
+ "tool_choice",
+ "tool_choice_option",
+ }
+) | REQUEST_OWNED_PARAMS
+
+
+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)
+
+ request_params = {**additional_kwargs, **kwargs}
+ collisions = REQUEST_OWNED_PARAMS & request_params.keys()
+ if collisions:
+ msg = (
+ "request parameters must not contain framework-owned fields: "
+ f"{sorted(collisions)}."
+ )
+ 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=request_params or None,
+ )
+
+ extra_args: Dict[str, Any] = {}
+
+ usage = response.get("usage")
+ if model_name and usage:
+ extra_args["model_name"] = model_name
+ extra_args["promptTokens"] = usage.get("prompt_tokens", 0)
+ extra_args["completionTokens"] = usage.get("completion_tokens", 0)
+
+ choice: Dict[str, Any] = response["choices"][0]
+ finish_reason = choice.get("finish_reason")
+ if finish_reason not in (None, "stop", "tool_calls"):
+ logger.warning(
+ "watsonx.ai chat for model %s finished with reason '%s'; "
+ "the response may be truncated or incomplete",
+ model_name,
+ finish_reason,
+ )
+
+ response_message: Dict[str, Any] = choice["message"]
+
+ tool_calls: List[Dict[str, Any]] = []
+ for tc in response_message.get("tool_calls") or []:
+ fn = tc.get("function", {}) or {}
+ args = _parse_tool_arguments(fn.get("arguments"))
+ tool_calls.append(
+ {
+ "id": uuid.uuid4(),
+ "type": tc.get("type", "function"),
+ "function": {
+ "name": fn.get("name"),
+ "arguments": args,
+ },
+ "original_id": tc.get("id"),
+ }
+ )
+
+ content = response_message.get("content") or ""
+
+ if extract_reasoning and content:
+ content, reasoning = self._extract_reasoning(content)
+ if reasoning:
+ extra_args["reasoning"] = reasoning
+
+ return ChatMessage(
+ role=MessageRole(response_message.get("role", "assistant")),
+ content=content,
+ tool_calls=tool_calls,
+ extra_args=extra_args,
+ )
+
+
+DEFAULT_TEMPERATURE = 0.1
+
+
+class WatsonxChatModelSetup(BaseChatModelSetup):
+ """Chat model configuration for IBM watsonx.ai."""
+
+ temperature: float = Field(
+ default=DEFAULT_TEMPERATURE,
+ description="The temperature to use for sampling.",
+ ge=0.0,
+ le=2.0,
+ )
+ max_tokens: int | None = Field(
+ default=None,
+ description="The maximum number of tokens to generate.",
+ gt=0,
+ )
+ additional_kwargs: Dict[str, Any] = Field(
+ default_factory=dict,
+ description="Additional chat parameters for the watsonx.ai API.",
+ )
+ extract_reasoning: bool = Field(
+ default=False,
+ description="If True, extracts reasoning content from the response and
stores it.",
+ )
+
+ def __init__(
+ self,
+ *,
+ model: str = DEFAULT_MODEL,
+ temperature: float = DEFAULT_TEMPERATURE,
+ max_tokens: int | None = None,
+ additional_kwargs: Dict[str, Any] | None = None,
+ extract_reasoning: bool = False,
+ **kwargs: Any,
+ ) -> None:
+ """Initialize the model configuration."""
+ if additional_kwargs is None:
+ additional_kwargs = {}
+ super().__init__(
+ model=model,
+ temperature=temperature,
+ max_tokens=max_tokens,
+ additional_kwargs=additional_kwargs,
+ extract_reasoning=extract_reasoning,
+ **kwargs,
+ )
+
+ @property
+ def model_kwargs(self) -> Dict[str, Any]:
+ """Return watsonx.ai model configuration."""
+ base_kwargs: Dict[str, Any] = {
+ "model": self.model,
+ "temperature": self.temperature,
+ "extract_reasoning": self.extract_reasoning,
+ }
+ if self.max_tokens is not None:
+ base_kwargs["max_tokens"] = self.max_tokens
+ if self.additional_kwargs:
+ base_kwargs["additional_kwargs"] = self.additional_kwargs
+ return base_kwargs
diff --git a/python/pyproject.toml b/python/pyproject.toml
index a9660d43..d284a6fd 100644
--- a/python/pyproject.toml
+++ b/python/pyproject.toml
@@ -56,6 +56,8 @@ dependencies = [
"dashscope~=1.24.2",
"openai>=1.66.3",
"anthropic>=0.77.0",
+ "ibm-watsonx-ai>=1.3.42,<1.6.0;python_version<'3.11'",
+ "ibm-watsonx-ai>=1.6.0,<2;python_version>='3.11'",
"chromadb==1.0.21",
"mem0ai>=0.1.43,<2.0.0",
"onnxruntime<1.24.1;python_version<'3.11'",
