cansuk99 commented on code in PR #922: URL: https://github.com/apache/flink-agents/pull/922#discussion_r3700357873
########## python/flink_agents/integrations/chat_models/watsonx/watsonx_chat_model.py: ########## @@ -0,0 +1,476 @@ +################################################################################ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +################################################################################# +import ast +import contextlib +import json +import logging +import os +import time +import uuid +from typing import Any, Dict, List, Sequence + +import httpx +from ibm_watsonx_ai import APIClient, Credentials +from ibm_watsonx_ai.foundation_models import ModelInference +from ibm_watsonx_ai.wml_client_error import ApiRequestFailure +from pydantic import Field, PrivateAttr +from typing_extensions import override + +from flink_agents.api.chat_message import ChatMessage, MessageRole +from flink_agents.api.chat_models.chat_model import ( + BaseChatModelConnection, + BaseChatModelSetup, +) +from flink_agents.api.tools.tool import Tool +from flink_agents.integrations.chat_models.chat_model_utils import to_openai_tool + +logger = logging.getLogger(__name__) + +DEFAULT_MODEL = "ibm/granite-4-h-small" +DEFAULT_REQUEST_TIMEOUT = 120.0 +DEFAULT_MAX_RETRIES = 3 +RETRYABLE_STATUS_CODES = frozenset({408, 429, 500, 502, 503, 504}) +RESERVED_ADDITIONAL_KWARGS = frozenset( + { + "model", + "temperature", + "max_tokens", + "extract_reasoning", + "tool_choice", + "tool_choice_option", + } +) + + +def _normalize(value: str | None) -> str | None: + if value is None or not value.strip(): + return None + return value.strip() + + +def _retry_delay_seconds(attempt: int, response: httpx.Response | None) -> float: + """Return capped exponential backoff, honoring a numeric Retry-After header.""" + backoff = min(2**attempt, 10) + if response is not None: + retry_after = response.headers.get("Retry-After") + if retry_after is not None: + with contextlib.suppress(ValueError): + return max(backoff, min(float(retry_after.strip()), 30)) + return backoff + + +def convert_to_watsonx_messages( + messages: Sequence[ChatMessage], +) -> List[Dict[str, Any]]: + """Convert framework messages to the watsonx.ai chat format.""" + watsonx_messages: List[Dict[str, Any]] = [] + for message in messages: + role = message.role + + if role == MessageRole.ASSISTANT: + assistant_message: Dict[str, Any] = {"role": "assistant"} + if message.content: + assistant_message["content"] = message.content + if message.tool_calls: + assistant_message["tool_calls"] = [ + _convert_to_watsonx_tool_call(tool_call) + for tool_call in message.tool_calls + ] + watsonx_messages.append(assistant_message) + elif role == MessageRole.TOOL: + tool_call_id = message.extra_args.get("external_id") + if not tool_call_id or not isinstance(tool_call_id, str): + msg = "Tool message must have 'external_id' as a string in extra_args" + raise ValueError(msg) + watsonx_messages.append( + { + "role": "tool", + "content": message.content, + "tool_call_id": tool_call_id, + } + ) + else: + watsonx_messages.append({"role": role.value, "content": message.content}) + return watsonx_messages + + +def _parse_tool_arguments(args: Any) -> Dict[str, Any]: + """Parse model-emitted tool arguments, including common malformed variants.""" + if args is None or args == "": + return {} + raw = args + for _ in range(3): + if not isinstance(args, str): + break + try: + args = json.loads(args) + except ValueError: + with contextlib.suppress(Exception): + literal = ast.literal_eval(args) + if isinstance(literal, dict): + args = literal + break + if not isinstance(args, dict): + msg = ( + "Failed to parse tool call arguments returned by watsonx.ai " + f"as a JSON object: {raw!r}" + ) + raise TypeError(msg) + return args + + +def _convert_to_watsonx_tool_call(tool_call: Dict[str, Any]) -> Dict[str, Any]: + """Convert a framework tool call to watsonx.ai format.""" + watsonx_tool_call_id = tool_call.get("original_id") + if watsonx_tool_call_id is None: + tool_call_id = tool_call.get("id") + if tool_call_id is None: + msg = "Tool call must have either 'original_id' or 'id' field" + raise ValueError(msg) + watsonx_tool_call_id = str(tool_call_id) + + arguments = tool_call["function"]["arguments"] + return { + "id": watsonx_tool_call_id, + "type": "function", + "function": { + "name": tool_call["function"]["name"], + "arguments": json.dumps(arguments) + if isinstance(arguments, dict) + else arguments, + }, + } + + +class WatsonxChatModelConnection(BaseChatModelConnection): + """Connection to the IBM watsonx.ai chat API.""" + + url: str = Field(description="The watsonx.ai service endpoint.") + api_key: str | None = Field(default=None, description="The IBM Cloud API key.") + token: str | None = Field( + default=None, description="A bearer token, as an alternative to api_key." + ) + project_id: str | None = Field( + default=None, description="The watsonx.ai project id." + ) + space_id: str | None = Field( + default=None, description="The watsonx.ai deployment space id." + ) + request_timeout: float = Field( + default=DEFAULT_REQUEST_TIMEOUT, + description="The timeout, in seconds, for chat requests to watsonx.ai.", + gt=0, + allow_inf_nan=False, + ) + max_retries: int = Field( + default=DEFAULT_MAX_RETRIES, + description="Maximum number of retries for transient failures.", + ge=0, + ) + + _client: APIClient | None = PrivateAttr(default=None) + _http_client: httpx.Client | None = PrivateAttr(default=None) + _models: Dict[str, ModelInference] = PrivateAttr(default_factory=dict) + + def __init__( + self, + *, + url: str | None = None, + api_key: str | None = None, + token: str | None = None, + project_id: str | None = None, + space_id: str | None = None, + request_timeout: float = DEFAULT_REQUEST_TIMEOUT, + max_retries: int = DEFAULT_MAX_RETRIES, + **kwargs: Any, + ) -> None: + """Initialize the connection.""" + resolved_url = _normalize(url) or _normalize(os.environ.get("WATSONX_URL")) + resolved_api_key = _normalize(api_key) or _normalize( + os.environ.get("WATSONX_API_KEY") + ) + resolved_token = _normalize(token) or _normalize( + os.environ.get("WATSONX_TOKEN") + ) + resolved_project_id = _normalize(project_id) or _normalize( + os.environ.get("WATSONX_PROJECT_ID") + ) + resolved_space_id = _normalize(space_id) or _normalize( + os.environ.get("WATSONX_SPACE_ID") + ) + + if not resolved_url: + msg = ( + "watsonx.ai url is not provided. Please pass it as an argument " + "or set the 'WATSONX_URL' environment variable." + ) + raise ValueError(msg) + if not resolved_api_key and not resolved_token: + msg = ( + "watsonx.ai credentials are not provided. Please pass 'api_key' " + "or 'token' as an argument, or set the 'WATSONX_API_KEY' or " + "'WATSONX_TOKEN' environment variable." + ) + raise ValueError(msg) + if resolved_api_key and resolved_token: + msg = ( + "watsonx.ai api_key and token cannot both be provided. Please configure " + "exactly one credential source." + ) + raise ValueError(msg) + if not resolved_project_id and not resolved_space_id: + msg = ( + "watsonx.ai project or space is not provided. Please pass " + "'project_id' or 'space_id' as an argument, or set the " + "'WATSONX_PROJECT_ID' or 'WATSONX_SPACE_ID' environment variable." + ) + raise ValueError(msg) + if resolved_project_id and resolved_space_id: + msg = ( + "watsonx.ai project and space cannot both be provided. Please configure " + "exactly one of 'project_id' or 'space_id'." + ) + raise ValueError(msg) + + super().__init__( + url=resolved_url, + api_key=resolved_api_key, + token=resolved_token, + project_id=resolved_project_id, + space_id=resolved_space_id, + request_timeout=request_timeout, + max_retries=max_retries, + **kwargs, + ) + + @property + def client(self) -> APIClient: + """Return the lazily initialized API client.""" + if self._client is None: + credential_kwargs: Dict[str, Any] = {"url": self.url} + if self.api_key: + credential_kwargs["api_key"] = self.api_key + if self.token: + credential_kwargs["token"] = self.token + self._http_client = httpx.Client(timeout=self.request_timeout) + self._client = APIClient( + credentials=Credentials(**credential_kwargs), + project_id=self.project_id, + space_id=self.space_id, + httpx_client=self._http_client, + ) + return self._client + + @override + def close(self) -> None: + """Close the underlying HTTP client.""" + self._models = {} + self._client = None + if self._http_client is not None: + with contextlib.suppress(Exception): + self._http_client.close() + self._http_client = None + + def _get_model(self, model: str) -> ModelInference: + if model not in self._models: + self._models[model] = ModelInference( + model_id=model, + api_client=self.client, + project_id=self.project_id, + space_id=self.space_id, + ) + return self._models[model] + + def _chat_with_retry(self, model_name: str, **chat_kwargs: Any) -> Dict[str, Any]: Review Comment: Hey @weiqingy he SDK retry layer was not accounted for in the initial implementation - thank you for catching this! I have disabled SDK retries by passing` max_retries=0` to `ModelInference` and kept `_chat_with_retry` as the single retry owner. Here I optimized for behavioural parity between the Java and Python implementation, as otherwise if we prioritize the SDK built-in retry which does not exist in Java, the connectors would differ. Does this approach make sense to you, or would you prefer delegating retries to the native SDK in Python while retaining the connector-level implementation in Java? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
