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 b1fe9df0b..d43d95f3e 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 static final class ChatModel { public static final String AZURE_OPENAI_SETUP = "org.apache.flink.agents.integrations.chatmodels.openai.AzureOpenAIChatModelSetup"; + // 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"; @@ -134,6 +140,12 @@ public static final class Python { public static final String TONGYI_SETUP = "flink_agents.integrations.chat_models.tongyi_chat_model.TongyiChatModelSetup"; + // 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 9a60b1adb..5cdcf2c32 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 @@ -88,12 +88,14 @@ public final class Aliases { chatConnJava.put("anthropic", ResourceName.ChatModel.ANTHROPIC_CONNECTION); chatConnJava.put("azure", ResourceName.ChatModel.AZURE_CONNECTION); chatConnJava.put("bedrock", ResourceName.ChatModel.BEDROCK_CONNECTION); + chatConnJava.put("watsonx", ResourceName.ChatModel.WATSONX_CONNECTION); Map chatConnPython = new HashMap<>(); chatConnPython.put("ollama", ResourceName.ChatModel.Python.OLLAMA_CONNECTION); chatConnPython.put("openai", ResourceName.ChatModel.Python.OPENAI_COMPLETIONS_CONNECTION); chatConnPython.put("anthropic", ResourceName.ChatModel.Python.ANTHROPIC_CONNECTION); chatConnPython.put("tongyi", ResourceName.ChatModel.Python.TONGYI_CONNECTION); chatConnPython.put("azure_openai", ResourceName.ChatModel.Python.AZURE_OPENAI_CONNECTION); + chatConnPython.put("watsonx", ResourceName.ChatModel.Python.WATSONX_CONNECTION); ca.put(ResourceType.CHAT_MODEL_CONNECTION, buildLangBuckets(chatConnJava, chatConnPython)); // CHAT_MODEL @@ -104,12 +106,14 @@ public final class Aliases { chatJava.put("anthropic", ResourceName.ChatModel.ANTHROPIC_SETUP); chatJava.put("azure", ResourceName.ChatModel.AZURE_SETUP); chatJava.put("bedrock", ResourceName.ChatModel.BEDROCK_SETUP); + chatJava.put("watsonx", ResourceName.ChatModel.WATSONX_SETUP); Map chatPython = new HashMap<>(); chatPython.put("ollama", ResourceName.ChatModel.Python.OLLAMA_SETUP); chatPython.put("openai", ResourceName.ChatModel.Python.OPENAI_COMPLETIONS_SETUP); chatPython.put("anthropic", ResourceName.ChatModel.Python.ANTHROPIC_SETUP); chatPython.put("tongyi", ResourceName.ChatModel.Python.TONGYI_SETUP); chatPython.put("azure_openai", ResourceName.ChatModel.Python.AZURE_OPENAI_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 38211cb0f..ac481de18 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 @@ -59,6 +59,22 @@ void clazzAliasPythonBucket() { assertThat(fqn).isEqualTo(ResourceName.ChatModel.Python.OLLAMA_CONNECTION); } + @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 9ee655bb2..ae917be67 100644 --- a/dist/pom.xml +++ b/dist/pom.xml @@ -92,6 +92,11 @@ under the License. flink-agents-integrations-chat-models-gemini ${project.version} + + org.apache.flink + flink-agents-integrations-chat-models-watsonx + ${project.version} + org.apache.flink flink-agents-integrations-embedding-models-ollama diff --git a/docs/content/docs/development/chat_models.md b/docs/content/docs/development/chat_models.md index c7b40ae05..0ae5d4de1 100644 --- a/docs/content/docs/development/chat_models.md +++ b/docs/content/docs/development/chat_models.md @@ -1111,6 +1111,163 @@ Some popular options include: Model availability and specifications may change. Always check the official DashScope documentation for the latest information before implementing in production. {{< /hint >}} +### 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` | Bearer token; 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` | Bearer token; 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. `` 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. `` 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). @@ -1370,4 +1527,4 @@ public class MyChatModelSetup extends BaseChatModelSetup { The built-in `chat_model_action` listens to `ChatRequestEvent` and `ToolResponseEvent`. To request a chat completion, send a `ChatRequestEvent`. If the model returns a final answer, the action sends a `ChatResponseEvent`. -If the model asks to call tools, `chat_model_action` sends a `ToolRequestEvent` instead of a final `ChatResponseEvent`. After the tools finish, it receives the matching `ToolResponseEvent`, appends the tool results to the chat history, and calls the model again. This loop continues until the model returns a final response. For details on how tools are executed, see [Built-in Events and Actions in Tool Use]({{< ref "docs/development/tool_use#built-in-events-and-actions" >}}). \ No newline at end of file +If the model asks to call tools, `chat_model_action` sends a `ToolRequestEvent` instead of a final `ChatResponseEvent`. After the tools finish, it receives the matching `ToolResponseEvent`, appends the tool results to the chat history, and calls the model again. This loop continues until the model returns a final response. For details on how tools are executed, see [Built-in Events and Actions in Tool Use]({{< ref "docs/development/tool_use#built-in-events-and-actions" >}}). diff --git a/docs/content/docs/development/yaml.md b/docs/content/docs/development/yaml.md index 87a9d7695..276647017 100644 --- a/docs/content/docs/development/yaml.md +++ b/docs/content/docs/development/yaml.md @@ -532,6 +532,7 @@ Common chat-model aliases: | `azure` | — | Azure OpenAI (Java) | | `bedrock` | — | Bedrock (Java) | | `tongyi` | Tongyi (Python) | — | +| `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/integrations/chat-models/pom.xml b/integrations/chat-models/pom.xml index ef8064528..faf2850fa 100644 --- a/integrations/chat-models/pom.xml +++ b/integrations/chat-models/pom.xml @@ -37,6 +37,7 @@ under the License. gemini ollama openai + watsonx diff --git a/integrations/chat-models/watsonx/pom.xml b/integrations/chat-models/watsonx/pom.xml new file mode 100644 index 000000000..1e2988351 --- /dev/null +++ b/integrations/chat-models/watsonx/pom.xml @@ -0,0 +1,64 @@ + + + + 4.0.0 + + + org.apache.flink + flink-agents-integrations-chat-models + 0.3-SNAPSHOT + ../pom.xml + + + flink-agents-integrations-chat-models-watsonx + Flink Agents : Integrations: Chat Models: IBM watsonx.ai + jar + + + + + org.apache.flink + flink-agents-api + ${project.version} + + + + com.fasterxml.jackson.core + jackson-databind + + + + org.apache.flink + flink-annotations + ${flink.version} + provided + + + + org.slf4j + slf4j-api + ${slf4j.version} + provided + + + + 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 000000000..0c0300235 --- /dev/null +++ b/integrations/chat-models/watsonx/src/main/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnection.java @@ -0,0 +1,678 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.integrations.chatmodels.watsonx; + +import com.fasterxml.jackson.core.json.JsonReadFeature; +import com.fasterxml.jackson.core.type.TypeReference; +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.chat.model.BaseChatModelConnection; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.apache.flink.agents.api.tools.Tool; +import org.apache.flink.annotation.VisibleForTesting; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.IOException; +import java.net.URI; +import java.net.URLEncoder; +import java.net.http.HttpClient; +import java.net.http.HttpRequest; +import java.net.http.HttpResponse; +import java.nio.charset.StandardCharsets; +import java.time.Duration; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Function; +import java.util.regex.Matcher; +import java.util.regex.Pattern; + +/** Chat model connection for the IBM watsonx.ai text chat REST API. */ +public class WatsonxChatModelConnection extends BaseChatModelConnection { + + private static final Logger LOG = LoggerFactory.getLogger(WatsonxChatModelConnection.class); + + static final String DEFAULT_IAM_URL = "https://iam.cloud.ibm.com"; + static final String DEFAULT_API_VERSION = "2025-04-23"; + static final long DEFAULT_REQUEST_TIMEOUT_SEC = 120; + static final int DEFAULT_MAX_RETRIES = 3; + private static final Set RETRYABLE_STATUS_CODES = Set.of(408, 429, 500, 502, 503, 504); + + private static final Set CONTROL_PARAMS = + Set.of( + "model", + "tool_choice", + "tool_choice_option", + "extract_reasoning", + "additional_kwargs"); + private static final Set RESERVED_ADDITIONAL_KWARGS = + Set.of( + "model", + "temperature", + "max_tokens", + "extract_reasoning", + "tool_choice", + "tool_choice_option"); + + private static final Pattern[] REASONING_PATTERNS = { + Pattern.compile("(.*?)", Pattern.DOTALL | Pattern.CASE_INSENSITIVE), + Pattern.compile("(.*?)", Pattern.DOTALL | Pattern.CASE_INSENSITIVE), + Pattern.compile("(.*?)", 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 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 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 messages, List tools, Map modelParams) { + try { + final String modelName = (String) modelParams.get("model"); + final boolean extractReasoning = + Boolean.TRUE.equals(modelParams.get("extract_reasoning")); + final ObjectNode payload = buildPayload(messages, tools, modelParams); + if (projectId != null && !projectId.isEmpty()) { + payload.put("project_id", projectId); + } else { + payload.put("space_id", spaceId); + } + + final HttpRequest request = + HttpRequest.newBuilder() + .uri(URI.create(url + "/ml/v1/text/chat?version=" + apiVersion)) + .timeout(requestTimeout) + .header("Authorization", "Bearer " + getBearerToken()) + .header("Content-Type", "application/json") + .header("Accept", "application/json") + .POST( + HttpRequest.BodyPublishers.ofString( + MAPPER.writeValueAsString(payload))) + .build(); + + final HttpResponse response = sendWithRetry(request); + if (response.statusCode() / 100 != 2) { + throw new RuntimeException( + String.format( + "watsonx.ai chat request failed with status %d: %s", + response.statusCode(), response.body())); + } + + final ChatMessage chatMessage = + parseResponse(MAPPER.readTree(response.body()), modelName); + if (extractReasoning) { + final String[] parts = extractReasoning(chatMessage.getContent()); + chatMessage.setContent(parts[0]); + if (parts[1] != null) { + chatMessage.getExtraArgs().put("reasoning", parts[1]); + } + } + return chatMessage; + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new RuntimeException("Interrupted while calling watsonx.ai.", e); + } catch (RuntimeException e) { + throw e; + } catch (Exception e) { + throw new RuntimeException(e); + } + } + + /** + * Sends the request, retrying HTTP 408, 429, 500, 502, 503, and 504 responses and I/O errors up + * to {@code max_retries} times with capped exponential backoff, honoring {@code Retry-After}. + */ + private HttpResponse sendWithRetry(HttpRequest request) + throws IOException, InterruptedException { + for (int attempt = 0; ; attempt++) { + try { + final HttpResponse 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 reasoningChunks = new ArrayList<>(); + String cleaned = content; + for (Pattern pattern : REASONING_PATTERNS) { + final Matcher matcher = pattern.matcher(cleaned); + final StringBuilder rest = new StringBuilder(); + boolean found = false; + int position = 0; + while (matcher.find()) { + final String chunk = matcher.group(1).trim(); + if (!chunk.isEmpty()) { + reasoningChunks.add(chunk); + } + rest.append(cleaned, position, matcher.start()); + position = matcher.end(); + found = true; + } + if (found) { + rest.append(cleaned, position, cleaned.length()); + cleaned = rest.toString(); + } + } + final String reasoning = + reasoningChunks.isEmpty() ? null : String.join("\n\n", reasoningChunks); + cleaned = cleaned.replaceAll("\\n{3,}", "\n\n").replaceAll(" {2,}", " ").trim(); + return new String[] {cleaned, reasoning}; + } + + @VisibleForTesting + static ObjectNode buildPayload( + List messages, List tools, Map 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 additionalKwargs = + (Map) modelParams.get("additional_kwargs"); + if (additionalKwargs != null) { + final Set collisions = new java.util.HashSet<>(additionalKwargs.keySet()); + collisions.retainAll(RESERVED_ADDITIONAL_KWARGS); + if (!collisions.isEmpty()) { + throw new IllegalArgumentException( + "additional_kwargs must not contain reserved typed fields: " + + collisions + + ". Set these via the corresponding Setup field instead."); + } + additionalKwargs.forEach( + (key, value) -> { + if (value != null) { + payload.set(key, MAPPER.valueToTree(value)); + } + }); + } + + for (Map.Entry 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. + * + *
    + *
  • SYSTEM/USER messages carry {@code role} and {@code content}. + *
  • ASSISTANT messages may carry {@code tool_calls} with JSON string arguments. + *
  • TOOL messages carry {@code tool_call_id} referencing the original call, taken from the + * {@code externalId} entry of the message extra args. + *
+ */ + @VisibleForTesting + static ArrayNode convertMessages(List 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> 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> toolCalls) { + final ArrayNode result = MAPPER.createArrayNode(); + for (Map 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 function = (Map) 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 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> toolCalls = new java.util.ArrayList<>(); + for (JsonNode toolCallNode : toolCallsNode) { + final String id = toolCallNode.required("id").asText(); + final JsonNode functionNode = toolCallNode.required("function"); + final Map 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 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. + * + *

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 parseToolArguments(JsonNode argumentsNode) { + if (argumentsNode == null || argumentsNode.isNull()) { + return Map.of(); + } + if (argumentsNode.isObject()) { + return MAPPER.convertValue(argumentsNode, new TypeReference>() {}); + } + 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>() {}); + } + 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 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); + } + } +} 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 000000000..01df5b097 --- /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 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 additionalKwargs = descriptor.getArgument("additional_kwargs"); + this.additionalKwargs = additionalKwargs != null ? additionalKwargs : Map.of(); + } + + @Override + public Map getParameters() { + Map 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 000000000..8ef440b4e --- /dev/null +++ b/integrations/chat-models/watsonx/src/test/java/org/apache/flink/agents/integrations/chatmodels/watsonx/WatsonxChatModelConnectionTest.java @@ -0,0 +1,395 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.flink.agents.integrations.chatmodels.watsonx; + +import com.fasterxml.jackson.databind.JsonNode; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ArrayNode; +import com.fasterxml.jackson.databind.node.ObjectNode; +import org.apache.flink.agents.api.chat.messages.ChatMessage; +import org.apache.flink.agents.api.chat.messages.MessageRole; +import org.apache.flink.agents.api.chat.model.BaseChatModelConnection; +import org.apache.flink.agents.api.resource.ResourceContext; +import org.apache.flink.agents.api.resource.ResourceDescriptor; +import org.junit.jupiter.api.DisplayName; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +import java.util.List; +import java.util.Map; +import java.util.function.Function; +import java.util.stream.Stream; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** + * Unit tests for {@link WatsonxChatModelConnection}. These exercise the protocol-conversion logic + * with no network access, so they run in CI without any API key. + */ +class WatsonxChatModelConnectionTest { + + private static final ResourceContext NOOP = ResourceContext.fromGetResource((a, b) -> null); + private static final Function NO_ENVIRONMENT = ignored -> null; + private static final ObjectMapper MAPPER = new ObjectMapper(); + + private static ResourceDescriptor descriptor(String url, String apiKey, String projectId) { + ResourceDescriptor.Builder b = + ResourceDescriptor.Builder.newBuilder(WatsonxChatModelConnection.class.getName()); + if (url != null) { + b.addInitialArgument("url", url); + } + if (apiKey != null) { + b.addInitialArgument("api_key", apiKey); + } + if (projectId != null) { + b.addInitialArgument("project_id", projectId); + } + return b.build(); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("missingRequiredConfiguration") + void testConstructorRejectsMissingRequiredConfiguration( + String ignoredCaseName, + String url, + String apiKey, + String projectId, + String expectedMessage) { + assertThatThrownBy( + () -> + new WatsonxChatModelConnection( + descriptor(url, apiKey, projectId), NOOP, NO_ENVIRONMENT)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining(expectedMessage); + } + + private static Stream missingRequiredConfiguration() { + return Stream.of( + Arguments.of("missing url", null, "test-key", "test-project", "url"), + Arguments.of( + "missing credentials", + "https://us-south.ml.cloud.ibm.com", + null, + "test-project", + "credentials"), + Arguments.of( + "missing project or space", + "https://us-south.ml.cloud.ibm.com", + "test-key", + null, + "project or space")); + } + + @Test + @DisplayName("Constructor accepts space_id without project_id") + void testConstructorWithSpaceId() { + ResourceDescriptor descriptor = + ResourceDescriptor.Builder.newBuilder(WatsonxChatModelConnection.class.getName()) + .addInitialArgument("url", " https://us-south.ml.cloud.ibm.com ") + .addInitialArgument("api_key", " test-key ") + .addInitialArgument("space_id", " test-space ") + .build(); + + assertThat(new WatsonxChatModelConnection(descriptor, NOOP, NO_ENVIRONMENT)) + .isInstanceOf(BaseChatModelConnection.class); + } + + @Test + @DisplayName("Constructor rejects ambiguous scope and credentials") + void testConstructorRejectsAmbiguousConfiguration() { + ResourceDescriptor descriptor = + ResourceDescriptor.Builder.newBuilder(WatsonxChatModelConnection.class.getName()) + .addInitialArgument("url", "https://us-south.ml.cloud.ibm.com") + .addInitialArgument("api_key", "test-key") + .addInitialArgument("project_id", "test-project") + .addInitialArgument("space_id", "test-space") + .build(); + + assertThatThrownBy(() -> new WatsonxChatModelConnection(descriptor, NOOP, NO_ENVIRONMENT)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("cannot both be provided") + .hasMessageContaining("exactly one"); + + ResourceDescriptor credentials = + ResourceDescriptor.Builder.newBuilder(WatsonxChatModelConnection.class.getName()) + .addInitialArgument("url", " https://us-south.ml.cloud.ibm.com ") + .addInitialArgument("api_key", " test-key ") + .addInitialArgument("token", " test-token ") + .addInitialArgument("project_id", " test-project ") + .build(); + assertThatThrownBy(() -> new WatsonxChatModelConnection(credentials, NOOP, NO_ENVIRONMENT)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("api_key and token") + .hasMessageContaining("exactly one"); + } + + @Test + @DisplayName("Request timeout accepts positive fractional seconds and rejects invalid values") + void testRequestTimeoutValidation() { + ResourceDescriptor fractionalTimeout = + ResourceDescriptor.Builder.newBuilder(WatsonxChatModelConnection.class.getName()) + .addInitialArgument("url", "https://us-south.ml.cloud.ibm.com") + .addInitialArgument("api_key", "test-key") + .addInitialArgument("project_id", "test-project") + .addInitialArgument("request_timeout", 0.5) + .build(); + assertThat(new WatsonxChatModelConnection(fractionalTimeout, NOOP, NO_ENVIRONMENT)) + .isInstanceOf(BaseChatModelConnection.class); + + for (double invalidTimeout : List.of(0.0, -1.0, Double.NaN, Double.POSITIVE_INFINITY)) { + ResourceDescriptor invalid = + ResourceDescriptor.Builder.newBuilder( + WatsonxChatModelConnection.class.getName()) + .addInitialArgument("url", "https://us-south.ml.cloud.ibm.com") + .addInitialArgument("api_key", "test-key") + .addInitialArgument("project_id", "test-project") + .addInitialArgument("request_timeout", invalidTimeout) + .build(); + assertThatThrownBy(() -> new WatsonxChatModelConnection(invalid, NOOP, NO_ENVIRONMENT)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("request_timeout"); + } + } + + @Test + @DisplayName("max_retries must be a non-negative integer") + void testMaxRetriesValidation() { + for (Number invalidMaxRetries : + new Number[] {-1, 0.9, Double.NaN, Double.POSITIVE_INFINITY}) { + ResourceDescriptor invalid = + ResourceDescriptor.Builder.newBuilder( + WatsonxChatModelConnection.class.getName()) + .addInitialArgument("url", "https://us-south.ml.cloud.ibm.com") + .addInitialArgument("api_key", "test-key") + .addInitialArgument("project_id", "test-project") + .addInitialArgument("max_retries", invalidMaxRetries) + .build(); + assertThatThrownBy(() -> new WatsonxChatModelConnection(invalid, NOOP, NO_ENVIRONMENT)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("max_retries"); + } + } + + @Test + @DisplayName("System, user, assistant and tool messages convert to the watsonx format") + void testConvertMessages() { + ChatMessage assistant = new ChatMessage(MessageRole.ASSISTANT, ""); + assistant.setToolCalls( + List.of( + Map.of( + "id", "internal-uuid", + "original_id", "call_abc123", + "type", "function", + "function", + Map.of( + "name", + "add", + "arguments", + Map.of("a", 1, "b", 2))))); + ChatMessage toolResult = + new ChatMessage(MessageRole.TOOL, "3", Map.of("externalId", "call_abc123")); + + ArrayNode converted = + WatsonxChatModelConnection.convertMessages( + List.of( + new ChatMessage(MessageRole.SYSTEM, "You are helpful."), + new ChatMessage(MessageRole.USER, "What is 1 + 2?"), + assistant, + toolResult)); + + assertThat(converted).hasSize(4); + assertThat(converted.get(0).get("role").asText()).isEqualTo("system"); + assertThat(converted.get(0).get("content").asText()).isEqualTo("You are helpful."); + assertThat(converted.get(1).get("role").asText()).isEqualTo("user"); + + JsonNode assistantNode = converted.get(2); + assertThat(assistantNode.get("role").asText()).isEqualTo("assistant"); + assertThat(assistantNode.has("content")).isFalse(); + JsonNode toolCall = assistantNode.get("tool_calls").get(0); + assertThat(toolCall.get("id").asText()).isEqualTo("call_abc123"); + assertThat(toolCall.get("function").get("name").asText()).isEqualTo("add"); + // arguments must be serialized as a JSON string + assertThat(toolCall.get("function").get("arguments").isTextual()).isTrue(); + + JsonNode toolNode = converted.get(3); + assertThat(toolNode.get("role").asText()).isEqualTo("tool"); + assertThat(toolNode.get("tool_call_id").asText()).isEqualTo("call_abc123"); + assertThat(toolNode.get("content").asText()).isEqualTo("3"); + } + + @Test + @DisplayName("Tool message without externalId is rejected") + void testConvertToolMessageWithoutExternalId() { + assertThatThrownBy( + () -> + WatsonxChatModelConnection.convertMessages( + List.of(new ChatMessage(MessageRole.TOOL, "3")))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("externalId"); + } + + @Test + @DisplayName("Model params are copied top-level into the payload") + void testBuildPayload() { + ObjectNode payload = + WatsonxChatModelConnection.buildPayload( + List.of(new ChatMessage(MessageRole.USER, "Hello!")), + List.of(), + Map.of( + "model", + "ibm/granite-3-3-8b-instruct", + "temperature", + 0.5, + "max_tokens", + 256, + "extract_reasoning", + true, + "additional_kwargs", + Map.of("top_p", 0.9))); + + assertThat(payload.get("model_id").asText()).isEqualTo("ibm/granite-3-3-8b-instruct"); + assertThat(payload.get("temperature").asDouble()).isEqualTo(0.5); + assertThat(payload.get("max_tokens").asInt()).isEqualTo(256); + assertThat(payload.get("top_p").asDouble()).isEqualTo(0.9); + assertThat(payload.get("messages")).hasSize(1); + // framework control params must not leak into the request + assertThat(payload.has("model")).isFalse(); + assertThat(payload.has("extract_reasoning")).isFalse(); + assertThat(payload.has("tools")).isFalse(); + + assertThatThrownBy( + () -> + WatsonxChatModelConnection.buildPayload( + List.of(new ChatMessage(MessageRole.USER, "Hello!")), + List.of(), + Map.of( + "model", + "ibm/granite-3-3-8b-instruct", + "additional_kwargs", + Map.of("temperature", 5.0)))) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("additional_kwargs") + .hasMessageContaining("temperature"); + } + + @Test + @DisplayName("Chat response with content and usage parses into a ChatMessage") + void testParseResponse() throws Exception { + JsonNode response = + MAPPER.readTree( + "{\"choices\": [{\"index\": 0, \"message\": {\"role\": \"assistant\"," + + " \"content\": \"Hello there!\"}, \"finish_reason\": \"stop\"}]," + + " \"usage\": {\"prompt_tokens\": 100, \"completion_tokens\": 50," + + " \"total_tokens\": 150}}"); + + ChatMessage message = + WatsonxChatModelConnection.parseResponse(response, "ibm/granite-3-3-8b-instruct"); + + assertThat(message.getRole()).isEqualTo(MessageRole.ASSISTANT); + assertThat(message.getContent()).isEqualTo("Hello there!"); + assertThat(message.getExtraArgs().get("model_name")) + .isEqualTo("ibm/granite-3-3-8b-instruct"); + assertThat(message.getExtraArgs().get("promptTokens")).isEqualTo(100L); + assertThat(message.getExtraArgs().get("completionTokens")).isEqualTo(50L); + } + + @Test + @DisplayName("Transient HTTP statuses are retryable and backoff honors Retry-After") + void testRetryPolicy() { + assertThat(WatsonxChatModelConnection.isRetryableStatus(408)).isTrue(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(429)).isTrue(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(500)).isTrue(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(502)).isTrue(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(503)).isTrue(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(504)).isTrue(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(200)).isFalse(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(400)).isFalse(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(401)).isFalse(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(501)).isFalse(); + assertThat(WatsonxChatModelConnection.isRetryableStatus(505)).isFalse(); + + // exponential backoff, capped + assertThat(WatsonxChatModelConnection.retryDelayMillis(0, null)).isEqualTo(1000L); + assertThat(WatsonxChatModelConnection.retryDelayMillis(1, null)).isEqualTo(2000L); + assertThat(WatsonxChatModelConnection.retryDelayMillis(10, null)).isEqualTo(10_000L); + // Retry-After wins when larger, is capped, and non-numeric values are ignored + assertThat(WatsonxChatModelConnection.retryDelayMillis(0, "5")).isEqualTo(5000L); + assertThat(WatsonxChatModelConnection.retryDelayMillis(0, "600")).isEqualTo(30_000L); + assertThat(WatsonxChatModelConnection.retryDelayMillis(0, "not-a-number")).isEqualTo(1000L); + } + + @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 toolCall = message.getToolCalls().get(0); + assertThat(toolCall.get("id")).isEqualTo("call_abc123"); + assertThat(toolCall.get("original_id")).isEqualTo("call_abc123"); + @SuppressWarnings("unchecked") + Map function = (Map) 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 000000000..50197e18d --- /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. + * + *

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 000000000..e9cafeb36 --- /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 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 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/resource.py b/python/flink_agents/api/resource.py index 21d033429..b05fca974 100644 --- a/python/flink_agents/api/resource.py +++ b/python/flink_agents/api/resource.py @@ -257,6 +257,10 @@ class ChatModel: TONGYI_CONNECTION = "flink_agents.integrations.chat_models.tongyi_chat_model.TongyiChatModelConnection" TONGYI_SETUP = "flink_agents.integrations.chat_models.tongyi_chat_model.TongyiChatModelSetup" + # 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" @@ -299,6 +303,10 @@ class Java: AZURE_OPENAI_CONNECTION = "org.apache.flink.agents.integrations.chatmodels.openai.AzureOpenAIChatModelConnection" AZURE_OPENAI_SETUP = "org.apache.flink.agents.integrations.chatmodels.openai.AzureOpenAIChatModelSetup" + # 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 cef9b59a2..ad9c68ef3 100644 --- a/python/flink_agents/api/yaml/aliases.py +++ b/python/flink_agents/api/yaml/aliases.py @@ -66,6 +66,7 @@ "anthropic": ResourceName.ChatModel.ANTHROPIC_CONNECTION, "tongyi": ResourceName.ChatModel.TONGYI_CONNECTION, "azure_openai": ResourceName.ChatModel.AZURE_OPENAI_CONNECTION, + "watsonx": ResourceName.ChatModel.WATSONX_CONNECTION, }, "java": { "ollama": ResourceName.ChatModel.Java.OLLAMA_CONNECTION, @@ -74,6 +75,7 @@ "anthropic": ResourceName.ChatModel.Java.ANTHROPIC_CONNECTION, "azure": ResourceName.ChatModel.Java.AZURE_CONNECTION, "bedrock": ResourceName.ChatModel.Java.BEDROCK_CONNECTION, + "watsonx": ResourceName.ChatModel.Java.WATSONX_CONNECTION, }, }, ResourceType.CHAT_MODEL: { @@ -83,6 +85,7 @@ "anthropic": ResourceName.ChatModel.ANTHROPIC_SETUP, "tongyi": ResourceName.ChatModel.TONGYI_SETUP, "azure_openai": ResourceName.ChatModel.AZURE_OPENAI_SETUP, + "watsonx": ResourceName.ChatModel.WATSONX_SETUP, }, "java": { "ollama": ResourceName.ChatModel.Java.OLLAMA_SETUP, @@ -91,6 +94,7 @@ "anthropic": ResourceName.ChatModel.Java.ANTHROPIC_SETUP, "azure": ResourceName.ChatModel.Java.AZURE_SETUP, "bedrock": ResourceName.ChatModel.Java.BEDROCK_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 b43db0d40..d59bd24c8 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 @@ 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/watsonx/__init__.py b/python/flink_agents/integrations/chat_models/watsonx/__init__.py new file mode 100644 index 000000000..e154fadd3 --- /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 000000000..e154fadd3 --- /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 000000000..4a5dfc347 --- /dev/null +++ b/python/flink_agents/integrations/chat_models/watsonx/tests/test_watsonx_chat_model.py @@ -0,0 +1,360 @@ +################################################################################ +# 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, + ) + + +@pytest.mark.integration +@pytest.mark.skipif( + 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!"} + ) + monkeypatch.setattr( + WatsonxChatModelConnection, "_get_model", lambda self, model: mock_model + ) + + connection = _fake_connection() + + 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!")]) + + 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.9, + } + + 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 + + +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}, + ) 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 000000000..95300b2e4 --- /dev/null +++ b/python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py @@ -0,0 +1,476 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################# +import ast +import contextlib +import json +import logging +import os +import time +import uuid +from typing import Any, Dict, List, Sequence + +import httpx +from ibm_watsonx_ai import APIClient, Credentials +from ibm_watsonx_ai.foundation_models import ModelInference +from ibm_watsonx_ai.wml_client_error import ApiRequestFailure +from pydantic import Field, PrivateAttr +from typing_extensions import override + +from flink_agents.api.chat_message import ChatMessage, MessageRole +from flink_agents.api.chat_models.chat_model import ( + BaseChatModelConnection, + BaseChatModelSetup, +) +from flink_agents.api.tools.tool import Tool +from flink_agents.integrations.chat_models.chat_model_utils import to_openai_tool + +logger = logging.getLogger(__name__) + +DEFAULT_MODEL = "ibm/granite-4-h-small" +DEFAULT_REQUEST_TIMEOUT = 120.0 +DEFAULT_MAX_RETRIES = 3 +RETRYABLE_STATUS_CODES = frozenset({408, 429, 500, 502, 503, 504}) +RESERVED_ADDITIONAL_KWARGS = frozenset( + { + "model", + "temperature", + "max_tokens", + "extract_reasoning", + "tool_choice", + "tool_choice_option", + } +) + + +def _normalize(value: str | None) -> str | None: + if value is None or not value.strip(): + return None + return value.strip() + + +def _retry_delay_seconds(attempt: int, response: httpx.Response | None) -> float: + """Return capped exponential backoff, honoring a numeric Retry-After header.""" + backoff = min(2**attempt, 10) + if response is not None: + retry_after = response.headers.get("Retry-After") + if retry_after is not None: + with contextlib.suppress(ValueError): + return max(backoff, min(float(retry_after.strip()), 30)) + return backoff + + +def convert_to_watsonx_messages( + messages: Sequence[ChatMessage], +) -> List[Dict[str, Any]]: + """Convert framework messages to the watsonx.ai chat format.""" + watsonx_messages: List[Dict[str, Any]] = [] + for message in messages: + role = message.role + + if role == MessageRole.ASSISTANT: + assistant_message: Dict[str, Any] = {"role": "assistant"} + if message.content: + assistant_message["content"] = message.content + if message.tool_calls: + assistant_message["tool_calls"] = [ + _convert_to_watsonx_tool_call(tool_call) + for tool_call in message.tool_calls + ] + watsonx_messages.append(assistant_message) + elif role == MessageRole.TOOL: + tool_call_id = message.extra_args.get("external_id") + if not tool_call_id or not isinstance(tool_call_id, str): + msg = "Tool message must have 'external_id' as a string in extra_args" + raise ValueError(msg) + watsonx_messages.append( + { + "role": "tool", + "content": message.content, + "tool_call_id": tool_call_id, + } + ) + else: + watsonx_messages.append({"role": role.value, "content": message.content}) + return watsonx_messages + + +def _parse_tool_arguments(args: Any) -> Dict[str, Any]: + """Parse model-emitted tool arguments, including common malformed variants.""" + if args is None or args == "": + return {} + raw = args + for _ in range(3): + if not isinstance(args, str): + break + try: + args = json.loads(args) + except ValueError: + with contextlib.suppress(Exception): + literal = ast.literal_eval(args) + if isinstance(literal, dict): + args = literal + break + if not isinstance(args, dict): + msg = ( + "Failed to parse tool call arguments returned by watsonx.ai " + f"as a JSON object: {raw!r}" + ) + raise TypeError(msg) + return args + + +def _convert_to_watsonx_tool_call(tool_call: Dict[str, Any]) -> Dict[str, Any]: + """Convert a framework tool call to watsonx.ai format.""" + watsonx_tool_call_id = tool_call.get("original_id") + if watsonx_tool_call_id is None: + tool_call_id = tool_call.get("id") + if tool_call_id is None: + msg = "Tool call must have either 'original_id' or 'id' field" + raise ValueError(msg) + watsonx_tool_call_id = str(tool_call_id) + + arguments = tool_call["function"]["arguments"] + return { + "id": watsonx_tool_call_id, + "type": "function", + "function": { + "name": tool_call["function"]["name"], + "arguments": json.dumps(arguments) + if isinstance(arguments, dict) + else arguments, + }, + } + + +class WatsonxChatModelConnection(BaseChatModelConnection): + """Connection to the IBM watsonx.ai chat API.""" + + url: str = Field(description="The watsonx.ai service endpoint.") + api_key: str | None = Field(default=None, description="The IBM Cloud API key.") + token: str | None = Field( + default=None, description="A bearer token, as an alternative to api_key." + ) + project_id: str | None = Field( + default=None, description="The watsonx.ai project id." + ) + space_id: str | None = Field( + default=None, description="The watsonx.ai deployment space id." + ) + request_timeout: float = Field( + default=DEFAULT_REQUEST_TIMEOUT, + description="The timeout, in seconds, for chat requests to watsonx.ai.", + gt=0, + allow_inf_nan=False, + ) + max_retries: int = Field( + default=DEFAULT_MAX_RETRIES, + description="Maximum number of retries for transient failures.", + ge=0, + ) + + _client: APIClient | None = PrivateAttr(default=None) + _http_client: httpx.Client | None = PrivateAttr(default=None) + _models: Dict[str, ModelInference] = PrivateAttr(default_factory=dict) + + def __init__( + self, + *, + url: str | None = None, + api_key: str | None = None, + token: str | None = None, + project_id: str | None = None, + space_id: str | None = None, + request_timeout: float = DEFAULT_REQUEST_TIMEOUT, + max_retries: int = DEFAULT_MAX_RETRIES, + **kwargs: Any, + ) -> None: + """Initialize the connection.""" + resolved_url = _normalize(url) or _normalize(os.environ.get("WATSONX_URL")) + resolved_api_key = _normalize(api_key) or _normalize( + os.environ.get("WATSONX_API_KEY") + ) + resolved_token = _normalize(token) or _normalize( + os.environ.get("WATSONX_TOKEN") + ) + resolved_project_id = _normalize(project_id) or _normalize( + os.environ.get("WATSONX_PROJECT_ID") + ) + resolved_space_id = _normalize(space_id) or _normalize( + os.environ.get("WATSONX_SPACE_ID") + ) + + if not resolved_url: + msg = ( + "watsonx.ai url is not provided. Please pass it as an argument " + "or set the 'WATSONX_URL' environment variable." + ) + raise ValueError(msg) + if not resolved_api_key and not resolved_token: + msg = ( + "watsonx.ai credentials are not provided. Please pass 'api_key' " + "or 'token' as an argument, or set the 'WATSONX_API_KEY' or " + "'WATSONX_TOKEN' environment variable." + ) + raise ValueError(msg) + if resolved_api_key and resolved_token: + msg = ( + "watsonx.ai api_key and token cannot both be provided. Please configure " + "exactly one credential source." + ) + raise ValueError(msg) + if not resolved_project_id and not resolved_space_id: + msg = ( + "watsonx.ai project or space is not provided. Please pass " + "'project_id' or 'space_id' as an argument, or set the " + "'WATSONX_PROJECT_ID' or 'WATSONX_SPACE_ID' environment variable." + ) + raise ValueError(msg) + if resolved_project_id and resolved_space_id: + msg = ( + "watsonx.ai project and space cannot both be provided. Please configure " + "exactly one of 'project_id' or 'space_id'." + ) + raise ValueError(msg) + + super().__init__( + url=resolved_url, + api_key=resolved_api_key, + token=resolved_token, + project_id=resolved_project_id, + space_id=resolved_space_id, + request_timeout=request_timeout, + max_retries=max_retries, + **kwargs, + ) + + @property + def client(self) -> APIClient: + """Return the lazily initialized API client.""" + if self._client is None: + credential_kwargs: Dict[str, Any] = {"url": self.url} + if self.api_key: + credential_kwargs["api_key"] = self.api_key + if self.token: + credential_kwargs["token"] = self.token + self._http_client = httpx.Client(timeout=self.request_timeout) + self._client = APIClient( + credentials=Credentials(**credential_kwargs), + project_id=self.project_id, + space_id=self.space_id, + httpx_client=self._http_client, + ) + return self._client + + @override + def close(self) -> None: + """Close the underlying HTTP client.""" + self._models = {} + self._client = None + if self._http_client is not None: + with contextlib.suppress(Exception): + self._http_client.close() + self._http_client = None + + def _get_model(self, model: str) -> ModelInference: + if model not in self._models: + self._models[model] = ModelInference( + model_id=model, + api_client=self.client, + project_id=self.project_id, + space_id=self.space_id, + ) + return self._models[model] + + def _chat_with_retry(self, model_name: str, **chat_kwargs: Any) -> Dict[str, Any]: + """Call chat with retries for selected HTTP statuses and transport failures. + + Retries use capped exponential backoff and honor a numeric ``Retry-After`` + response header. + """ + attempt = 0 + while True: + try: + return self._get_model(model_name).chat(**chat_kwargs) + except (ApiRequestFailure, httpx.TransportError) as e: # noqa: PERF203 + response = getattr(e, "response", None) + status = getattr(response, "status_code", None) + retryable = isinstance(e, httpx.TransportError) or ( + status in RETRYABLE_STATUS_CODES + ) + if attempt >= self.max_retries or not retryable: + raise + delay = _retry_delay_seconds(attempt, response) + logger.warning( + "watsonx.ai chat request for model %s failed with %s; " + "retry %d/%d in %ds", + model_name, + status if status is not None else type(e).__name__, + attempt + 1, + self.max_retries, + delay, + ) + time.sleep(delay) + attempt += 1 + + def chat( + self, + messages: Sequence[ChatMessage], + tools: List[Tool] | None = None, + **kwargs: Any, + ) -> ChatMessage: + """Process a sequence of messages, and return a response.""" + model_name = kwargs.pop("model", DEFAULT_MODEL) + extract_reasoning = bool(kwargs.pop("extract_reasoning", False)) + tool_choice = kwargs.pop("tool_choice", None) + tool_choice_option = kwargs.pop("tool_choice_option", None) + additional_kwargs = kwargs.pop("additional_kwargs", None) or {} + collisions = RESERVED_ADDITIONAL_KWARGS & additional_kwargs.keys() + if collisions: + msg = ( + "additional_kwargs must not contain reserved typed fields: " + f"{sorted(collisions)}. Set these via the corresponding Setup field instead." + ) + raise ValueError(msg) + + tool_specs: List[Dict[str, Any]] | None = ( + [to_openai_tool(metadata=tool.metadata) for tool in tools] + if tools + else None + ) + + response = self._chat_with_retry( + model_name, + messages=convert_to_watsonx_messages(messages), + tools=tool_specs, + tool_choice=tool_choice, + tool_choice_option=tool_choice_option, + params={**kwargs, **additional_kwargs} or None, + ) + + 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 2d2d8d41c..5b1142a13 100644 --- a/python/pyproject.toml +++ b/python/pyproject.toml @@ -56,6 +56,8 @@ dependencies = [ "dashscope~=1.24.2", "openai>=1.66.3", "anthropic>=0.64.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'",