diff --git a/litellm/a2a_protocol/providers/config_manager.py b/litellm/a2a_protocol/providers/config_manager.py index d684efd475..ecb8f66bde 100644 --- a/litellm/a2a_protocol/providers/config_manager.py +++ b/litellm/a2a_protocol/providers/config_manager.py @@ -48,4 +48,11 @@ class A2AProviderConfigManager: return BedrockAgentCoreA2AConfig() + if custom_llm_provider == "watsonx_orchestrate": + from litellm.a2a_protocol.providers.watsonx_orchestrate.config import ( + WatsonxOrchestrateA2AConfig, + ) + + return WatsonxOrchestrateA2AConfig() + return None diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/__init__.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/__init__.py new file mode 100644 index 0000000000..096bcc0121 --- /dev/null +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/__init__.py @@ -0,0 +1,3 @@ +""" +IBM watsonx Orchestrate (WXO) A2A provider. +""" diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py new file mode 100644 index 0000000000..dbd4a0558f --- /dev/null +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/config.py @@ -0,0 +1,55 @@ +""" +A2A provider configuration for IBM watsonx Orchestrate (WXO). +""" + +from typing import Any, AsyncIterator, Dict, Optional + +from litellm.a2a_protocol.providers.base import BaseA2AProviderConfig +from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import ( + WatsonxOrchestrateHandler, +) + + +class WatsonxOrchestrateA2AConfig(BaseA2AProviderConfig): + """A2A bridge for IBM watsonx Orchestrate (REST runs API + poll/SSE).""" + + async def handle_non_streaming( + self, + request_id: str, + params: Dict[str, Any], + api_base: Optional[str] = None, + **kwargs: Any, + ) -> Dict[str, Any]: + """Handle a non-streaming A2A request via WXO runs API.""" + litellm_params = kwargs.get("litellm_params") + if not litellm_params: + raise ValueError( + "litellm_params is required for WatsonxOrchestrateA2AConfig " + "(must contain cp4d_host, instance_id, wxo_agent_id, api_key)" + ) + return await WatsonxOrchestrateHandler.handle_non_streaming( + request_id=request_id, + params=params, + litellm_params=litellm_params, + ) + + async def handle_streaming( + self, + request_id: str, + params: Dict[str, Any], + api_base: Optional[str] = None, + **kwargs: Any, + ) -> AsyncIterator[Dict[str, Any]]: + """Handle a streaming A2A request via WXO streaming runs API.""" + litellm_params = kwargs.get("litellm_params") + if not litellm_params: + raise ValueError( + "litellm_params is required for WatsonxOrchestrateA2AConfig " + "(must contain cp4d_host, instance_id, wxo_agent_id, api_key)" + ) + async for chunk in WatsonxOrchestrateHandler.handle_streaming( + request_id=request_id, + params=params, + litellm_params=litellm_params, + ): + yield chunk diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py new file mode 100644 index 0000000000..dbc0247618 --- /dev/null +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -0,0 +1,373 @@ +""" +Handler for IBM watsonx Orchestrate (WXO) agent provider. +""" + +import asyncio +import hashlib +import json +import time +from typing import Any, AsyncIterator, Dict, NamedTuple, Optional, Tuple, cast + +import httpx + +from litellm._logging import verbose_logger +from litellm.a2a_protocol.providers.watsonx_orchestrate.transformation import ( + WatsonxOrchestrateTransformation, +) +from litellm.llms.custom_httpx.http_handler import ( + AsyncHTTPHandler, + get_async_httpx_client, +) +from litellm.types.llms.custom_http import httpxSpecialProvider + +_IBM_CLOUD_IAM_URL = "https://iam.cloud.ibm.com/identity/token" +_POLL_INTERVAL_S = 2.0 +_MAX_POLL_ATTEMPTS = 90 +_TOKEN_CACHE_TTL_BUFFER_S = 60 +_token_cache: Dict[str, Tuple[str, float]] = {} + + +class WXORequestParams(NamedTuple): + cp4d_host: str + instance_id: str + wxo_agent_id: str + api_key: str + username: Optional[str] + auth_mode: str + thread_id: Optional[str] + + +class WatsonxOrchestrateHandler: + @staticmethod + def _http_client(timeout: float = 90.0) -> AsyncHTTPHandler: + return get_async_httpx_client( + llm_provider=cast(Any, httpxSpecialProvider.A2AProvider), + params={"timeout": timeout}, + ) + + @staticmethod + def _token_cache_key( + auth_mode: str, + cp4d_host: str, + api_key: str, + username: Optional[str], + ) -> str: + material = f"{auth_mode}:{cp4d_host}:{username or ''}:{api_key}" + return hashlib.sha256(material.encode()).hexdigest() + + @staticmethod + def _cp4d_token_ttl_seconds( + expiration: Any, now_wall: Optional[float] = None + ) -> int: + # CP4D returns expiration as absolute Unix epoch seconds, not a duration. + expires_at = int(expiration) + wall = now_wall if now_wall is not None else time.time() + return max(expires_at - int(wall), 0) + + @staticmethod + async def _get_bearer_token( + cp4d_host: str, + auth_mode: str, + api_key: str, + username: Optional[str] = None, + client: Optional[AsyncHTTPHandler] = None, + ) -> str: + cache_key = WatsonxOrchestrateHandler._token_cache_key( + auth_mode, cp4d_host, api_key, username + ) + now = time.monotonic() + cached = _token_cache.get(cache_key) + if cached and cached[1] > now: + return cached[0] + + if client is None: + client = WatsonxOrchestrateHandler._http_client(timeout=30.0) + + if auth_mode == "ibm_cloud": + response = await client.post( + _IBM_CLOUD_IAM_URL, + data={ + "grant_type": "urn:ibm:params:oauth:grant-type:apikey", + "apikey": api_key, + }, + headers={"Content-Type": "application/x-www-form-urlencoded"}, + ) + response.raise_for_status() + payload = response.json() + token = str(payload["access_token"]) + ttl_s = int(payload.get("expires_in", 3600)) + else: + if not username: + raise ValueError( + "'username' is required in litellm_params when auth_mode='cp4d'" + ) + token_url = f"{cp4d_host.rstrip('/')}/icp4d-api/v1/authorize" + response = await client.post( + token_url, + json={"username": username, "api_key": api_key}, + headers={"Content-Type": "application/json"}, + ) + response.raise_for_status() + payload = response.json() + token = str(payload["token"]) + expiration = payload.get("expiration") + if expiration is None: + ttl_s = 3600 + else: + ttl_s = WatsonxOrchestrateHandler._cp4d_token_ttl_seconds(expiration) + + expires_at = now + max(ttl_s - _TOKEN_CACHE_TTL_BUFFER_S, 0) + _token_cache[cache_key] = (token, expires_at) + for stale_key, (_, stale_expires_at) in list(_token_cache.items()): + if stale_expires_at <= now: + del _token_cache[stale_key] + return token + + @staticmethod + async def _poll_run( + base_url: str, + run_id: str, + auth_headers: Dict[str, str], + client: AsyncHTTPHandler, + max_attempts: int = _MAX_POLL_ATTEMPTS, + interval_s: float = _POLL_INTERVAL_S, + ) -> Dict[str, Any]: + url = f"{base_url}/v1/orchestrate/runs/{run_id}" + + for attempt in range(max_attempts): + await asyncio.sleep(interval_s) + response = await client.get(url, headers=auth_headers) + response.raise_for_status() + result: Dict[str, Any] = response.json() + status = result.get("status", "") + verbose_logger.debug( + f"WXO: Poll {attempt + 1}/{max_attempts} run='{run_id}' status='{status}'" + ) + if status in WatsonxOrchestrateTransformation.TERMINAL_STATES: + return result + + raise asyncio.TimeoutError( + f"WXO run '{run_id}' did not reach a terminal state after " + f"{max_attempts * interval_s:.0f}s" + ) + + @staticmethod + async def _get_successful_run_data( + run_data: Dict[str, Any], + base_url: str, + auth_headers: Dict[str, str], + client: AsyncHTTPHandler, + ) -> Dict[str, Any]: + status = run_data.get("status", "") + if status not in WatsonxOrchestrateTransformation.TERMINAL_STATES: + run_id = run_data.get("run_id") or run_data.get("id") or "" + if not run_id: + raise ValueError(f"WXO: No run_id in response: {run_data}") + run_data = await WatsonxOrchestrateHandler._poll_run( + base_url=base_url, + run_id=run_id, + auth_headers=auth_headers, + client=client, + ) + status = run_data.get("status", "") + + if status not in WatsonxOrchestrateTransformation.SUCCESS_STATES: + raise RuntimeError( + f"WXO run ended with non-success status '{status}': {run_data}" + ) + + return run_data + + @staticmethod + async def _accumulate_wxo_sse_text(response: Any) -> str: + accumulated_text = "" + async for line in response.aiter_lines(): + if not line.startswith("data:"): + continue + data_str = line[5:].strip() + if not data_str or data_str == "[DONE]": + continue + try: + event = json.loads(data_str) + except json.JSONDecodeError: + continue + chunk_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result( + event + ) + if chunk_text: + accumulated_text += chunk_text + return accumulated_text + + @staticmethod + def _extract_litellm_params(litellm_params: Dict[str, Any]) -> WXORequestParams: + cp4d_host = litellm_params.get("cp4d_host") or "" + instance_id = litellm_params.get("instance_id") or "" + wxo_agent_id = litellm_params.get("wxo_agent_id") or "" + api_key = litellm_params.get("api_key") or "" + + if not cp4d_host: + raise ValueError("'cp4d_host' is required in litellm_params for WXO agents") + if not instance_id: + raise ValueError( + "'instance_id' is required in litellm_params for WXO agents" + ) + if not wxo_agent_id: + raise ValueError( + "'wxo_agent_id' is required in litellm_params for WXO agents" + ) + if not api_key: + raise ValueError("'api_key' is required in litellm_params for WXO agents") + + return WXORequestParams( + cp4d_host=cp4d_host, + instance_id=instance_id, + wxo_agent_id=wxo_agent_id, + api_key=api_key, + username=litellm_params.get("username") or None, + auth_mode=litellm_params.get("auth_mode") or "cp4d", + thread_id=litellm_params.get("thread_id") or None, + ) + + @staticmethod + async def handle_non_streaming( + request_id: str, + params: Dict[str, Any], + litellm_params: Dict[str, Any], + ) -> Dict[str, Any]: + wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params) + + client = WatsonxOrchestrateHandler._http_client(timeout=90.0) + token = await WatsonxOrchestrateHandler._get_bearer_token( + cp4d_host=wxo.cp4d_host, + auth_mode=wxo.auth_mode, + api_key=wxo.api_key, + username=wxo.username, + client=client, + ) + base_url = WatsonxOrchestrateTransformation.get_api_base_url( + wxo.cp4d_host, wxo.instance_id + ) + auth_headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "Accept": "application/json", + } + + text = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params) + body = WatsonxOrchestrateTransformation.build_wxo_run_body( + wxo_agent_id=wxo.wxo_agent_id, text=text, thread_id=wxo.thread_id + ) + + run_response = await client.post( + f"{base_url}/v1/orchestrate/runs", + json=body, + headers=auth_headers, + ) + run_response.raise_for_status() + run_data: Dict[str, Any] = run_response.json() + + run_data = await WatsonxOrchestrateHandler._get_successful_run_data( + run_data=run_data, + base_url=base_url, + auth_headers=auth_headers, + client=client, + ) + + response_text = WatsonxOrchestrateTransformation.extract_text_from_wxo_result( + run_data + ) + return WatsonxOrchestrateTransformation.build_a2a_message_response( + request_id=request_id, text=response_text + ) + + @staticmethod + async def handle_streaming( + request_id: str, + params: Dict[str, Any], + litellm_params: Dict[str, Any], + chunk_size: int = 50, + delay_ms: int = 10, + ) -> AsyncIterator[Dict[str, Any]]: + wxo = WatsonxOrchestrateHandler._extract_litellm_params(litellm_params) + + client = WatsonxOrchestrateHandler._http_client(timeout=120.0) + token = await WatsonxOrchestrateHandler._get_bearer_token( + cp4d_host=wxo.cp4d_host, + auth_mode=wxo.auth_mode, + api_key=wxo.api_key, + username=wxo.username, + client=client, + ) + base_url = WatsonxOrchestrateTransformation.get_api_base_url( + wxo.cp4d_host, wxo.instance_id + ) + auth_headers = { + "Authorization": f"Bearer {token}", + "Content-Type": "application/json", + "Accept": "text/event-stream, application/json", + } + text = WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params) + body = WatsonxOrchestrateTransformation.build_wxo_run_body( + wxo_agent_id=wxo.wxo_agent_id, text=text, thread_id=wxo.thread_id + ) + + try: + response = await client.post( + f"{base_url}/v1/orchestrate/runs/stream", + json=body, + headers=auth_headers, + stream=True, + ) + response.raise_for_status() + except httpx.TransportError as exc: + verbose_logger.warning( + f"WXO: Streaming request failed before a run was submitted " + f"({exc!r}), falling back to non-streaming + fake streaming", + exc_info=True, + ) + result = await WatsonxOrchestrateHandler.handle_non_streaming( + request_id=request_id, + params=params, + litellm_params=litellm_params, + ) + response_text = ( + WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response( + result + ) + ) + async for ( + chunk + ) in WatsonxOrchestrateTransformation.fake_streaming_from_text( + text=response_text, + request_id=request_id, + chunk_size=chunk_size, + delay_ms=delay_ms, + ): + yield chunk + return + + content_type = response.headers.get("content-type", "").lower() + if "text/event-stream" not in content_type: + response_body = await response.aread() + result = json.loads(response_body) + result = await WatsonxOrchestrateHandler._get_successful_run_data( + run_data=result, + base_url=base_url, + auth_headers=auth_headers, + client=client, + ) + accumulated_text = ( + WatsonxOrchestrateTransformation.extract_text_from_wxo_result(result) + ) + else: + accumulated_text = await WatsonxOrchestrateHandler._accumulate_wxo_sse_text( + response + ) + + async for chunk in WatsonxOrchestrateTransformation.fake_streaming_from_text( + text=accumulated_text, + request_id=request_id, + chunk_size=chunk_size, + delay_ms=delay_ms, + ): + yield chunk diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py new file mode 100644 index 0000000000..824e9dbcdd --- /dev/null +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/transformation.py @@ -0,0 +1,224 @@ +""" +Transformation layer for IBM watsonx Orchestrate (WXO) agent provider. + +WXO uses a REST API (not A2A/JSON-RPC) with an async-poll execution model: + POST /v1/orchestrate/runs → submit run, get run_id + GET /v1/orchestrate/runs/{id} → poll until terminal state + POST /v1/orchestrate/runs/stream → native SSE streaming +""" + +import asyncio +from typing import Any, AsyncIterator, Dict, Optional +from uuid import uuid4 + +from litellm._logging import verbose_logger + + +class WatsonxOrchestrateTransformation: + """ + Handles request/response transformation between A2A and the WXO REST API. + """ + + TERMINAL_STATES = frozenset( + {"completed", "succeeded", "failed", "error", "cancelled"} + ) + SUCCESS_STATES = frozenset({"completed", "succeeded"}) + + @staticmethod + def get_api_base_url(cp4d_host: str, instance_id: str) -> str: + """Build the WXO API base URL from host and instance ID.""" + return f"{cp4d_host.rstrip('/')}/orchestrate/cpd/instances/{instance_id}" + + @staticmethod + def extract_text_from_a2a_params(params: Dict[str, Any]) -> str: + """ + Extract user message text from A2A MessageSendParams. + + A2A format: params.message.parts[*] where part.kind == "text" + """ + message = params.get("message", {}) + parts = message.get("parts", []) + texts = [] + for part in parts: + if not isinstance(part, dict): + continue + kind = part.get("kind") + if kind in (None, "", "text") and part.get("text"): + texts.append(part["text"]) + return " ".join(texts) or "" + + @staticmethod + def build_wxo_run_body( + wxo_agent_id: str, + text: str, + thread_id: Optional[str] = None, + ) -> Dict[str, Any]: + """Build the WXO POST /v1/orchestrate/runs request body.""" + body: Dict[str, Any] = { + "agent_id": wxo_agent_id, + "message": { + "role": "user", + "content": [ + { + "response_type": "text", + "text": text, + } + ], + }, + } + if thread_id: + body["thread_id"] = thread_id + return body + + @staticmethod + def extract_text_from_wxo_result(result: Any) -> str: + """ + Extract response text from a WXO run result. + + WXO can return text in several locations; checks in priority order per the API spec. + """ + if not isinstance(result, dict): + return "" + + # Primary: last_message.content[0].text + try: + text = result["last_message"]["content"][0]["text"] + if text: + return str(text) + except (KeyError, IndexError, TypeError): + pass + + # Secondary: result.data.message.content[0].text + try: + text = result["result"]["data"]["message"]["content"][0]["text"] + if text: + return str(text) + except (KeyError, IndexError, TypeError): + pass + + # Tertiary: results as a raw string + results = result.get("results") + if results and isinstance(results, str): + return results + + return "" + + @staticmethod + def extract_text_from_a2a_message_response(a2a_response: Dict[str, Any]) -> str: + result = a2a_response.get("result") + if not isinstance(result, dict): + verbose_logger.warning("WXO: A2A response missing result object") + return "" + parts = result.get("parts") + if not isinstance(parts, list): + verbose_logger.warning("WXO: A2A result has no parts list") + return "" + for part in parts: + if ( + isinstance(part, dict) + and part.get("kind") == "text" + and part.get("text") + ): + return str(part["text"]) + verbose_logger.warning("WXO: A2A result parts contained no text") + return "" + + @staticmethod + def build_a2a_message_response(request_id: str, text: str) -> Dict[str, Any]: + """ + Build a standard A2A non-streaming SendMessageResponse (kind=message). + """ + return { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "kind": "message", + "role": "agent", + "parts": [{"kind": "text", "text": text}], + "messageId": str(uuid4()), + }, + } + + @staticmethod + async def fake_streaming_from_text( + text: str, + request_id: str, + chunk_size: int = 50, + delay_ms: int = 10, + ) -> AsyncIterator[Dict[str, Any]]: + """ + Emit standard A2A streaming events from a completed text response. + + Event sequence: + 1. task (kind="task", state="submitted") + 2. status-update (kind="status-update", state="working") + 3. artifact-update chunks + 4. status-update (kind="status-update", state="completed", final=True) + """ + task_id = str(uuid4()) + context_id = str(uuid4()) + artifact_id = str(uuid4()) + + # 1. Task submitted + yield { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "id": task_id, + "kind": "task", + "status": {"state": "submitted"}, + }, + } + + # 2. Working + yield { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "final": False, + "kind": "status-update", + "status": {"state": "working"}, + "taskId": task_id, + }, + } + await asyncio.sleep(delay_ms / 1000.0) + + # 3. Artifact chunks (always emit at least one chunk, even for empty text) + text_to_chunk = text or "" + for i in range(0, max(len(text_to_chunk), 1), chunk_size): + chunk_text = text_to_chunk[i : i + chunk_size] + is_last = (i + chunk_size) >= max(len(text_to_chunk), 1) + yield { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "kind": "artifact-update", + "taskId": task_id, + "artifact": { + "artifactId": artifact_id, + "parts": [{"kind": "text", "text": chunk_text}], + }, + }, + } + if not is_last: + await asyncio.sleep(delay_ms / 1000.0) + + # 4. Completed + yield { + "jsonrpc": "2.0", + "id": request_id, + "result": { + "contextId": context_id, + "final": True, + "kind": "status-update", + "status": {"state": "completed"}, + "taskId": task_id, + }, + } + + verbose_logger.debug( + f"WXO: Fake streaming completed for request_id={request_id}" + ) diff --git a/litellm/proxy/public_endpoints/agent_create_fields.json b/litellm/proxy/public_endpoints/agent_create_fields.json index 931c9a4349..e58bd97cce 100644 --- a/litellm/proxy/public_endpoints/agent_create_fields.json +++ b/litellm/proxy/public_endpoints/agent_create_fields.json @@ -189,6 +189,78 @@ "litellm_params_template": { "custom_llm_provider": "vertex_ai" } + }, + { + "agent_type": "watsonx_orchestrate", + "agent_type_display_name": "watsonx Orchestrate", + "description": "Connect to IBM watsonx Orchestrate agents via CP4D or IBM Cloud IAM", + "logo_url": "/ui/assets/logos/watsonx.svg", + "credential_fields": [ + { + "key": "cp4d_host", + "label": "CP4D Host URL", + "placeholder": "https://cpd-cpd.apps.example.com", + "tooltip": "Your CP4D cluster base URL (e.g. https://cpd-cpd.apps.example.com). For IBM Cloud WXO, use the service endpoint.", + "required": true, + "field_type": "text", + "default_value": null, + "include_in_litellm_params": true + }, + { + "key": "instance_id", + "label": "WXO Instance ID", + "placeholder": "1769134113217795", + "tooltip": "The numeric watsonx Orchestrate instance ID. Find it in the WXO service URL: /orchestrate/cpd/instances/", + "required": true, + "field_type": "text", + "default_value": null, + "include_in_litellm_params": true + }, + { + "key": "wxo_agent_id", + "label": "WXO Agent ID", + "placeholder": "588c8cdf-60f4-454b-8468-8702b19dca46", + "tooltip": "UUID of the agent in watsonx Orchestrate. Find it via the WXO console or GET /v1/orchestrate/agents.", + "required": true, + "field_type": "text", + "default_value": null, + "include_in_litellm_params": true + }, + { + "key": "auth_mode", + "label": "Authentication Mode", + "placeholder": null, + "tooltip": "cp4d: on-prem / CloudPak for Data (requires username). ibm_cloud: IBM Cloud IAM (api_key only).", + "required": false, + "field_type": "select", + "options": ["cp4d", "ibm_cloud"], + "default_value": "cp4d", + "include_in_litellm_params": true + }, + { + "key": "username", + "label": "Username (CP4D only)", + "placeholder": "admin", + "tooltip": "Your CP4D username. Required when auth_mode is 'cp4d'.", + "required": false, + "field_type": "text", + "default_value": null, + "include_in_litellm_params": true + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": "CP4D API key (auth_mode=cp4d) or IBM Cloud API key (auth_mode=ibm_cloud).", + "required": true, + "field_type": "password", + "default_value": null, + "include_in_litellm_params": true + } + ], + "litellm_params_template": { + "custom_llm_provider": "watsonx_orchestrate" + } } ] diff --git a/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py b/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py new file mode 100644 index 0000000000..7968eed414 --- /dev/null +++ b/tests/test_litellm/a2a_protocol/providers/watsonx_orchestrate/test_watsonx_orchestrate_transformation.py @@ -0,0 +1,593 @@ +import asyncio +import json +import os +import sys +import time +from pathlib import Path + +import httpx +import pytest + +sys.path.insert(0, os.path.abspath("../../../../..")) + +from litellm.a2a_protocol.providers.config_manager import A2AProviderConfigManager +from litellm.a2a_protocol.providers.watsonx_orchestrate import handler as wxo_handler +from litellm.a2a_protocol.providers.watsonx_orchestrate.handler import ( + WatsonxOrchestrateHandler, +) +from litellm.a2a_protocol.providers.watsonx_orchestrate.transformation import ( + WatsonxOrchestrateTransformation, +) + + +class _JsonResponse: + def __init__(self, payload): + self.payload = payload + + def raise_for_status(self): + pass + + def json(self): + return self.payload + + +class _ShortTtlTokenClient: + def __init__(self): + self.calls = 0 + + async def post(self, *args, **kwargs): + self.calls += 1 + return _JsonResponse({"access_token": f"token-{self.calls}", "expires_in": 30}) + + +class _SSELines: + def __init__(self, lines): + self.lines = lines + + async def aiter_lines(self): + for line in self.lines: + yield line + + +class _InvalidJsonStreamResponse: + headers = {"content-type": "application/json"} + + def raise_for_status(self): + pass + + async def aread(self): + return b"not-json" + + +class _InvalidJsonStreamClient: + def __init__(self): + self.post_urls = [] + + async def post(self, url, **kwargs): + self.post_urls.append(url) + if "identity/token" in url: + return _JsonResponse({"access_token": "token", "expires_in": 3600}) + if url.endswith("/runs/stream"): + return _InvalidJsonStreamResponse() + if url.endswith("/runs"): + return _JsonResponse({"status": "completed", "results": "fallback text"}) + raise AssertionError(url) + + +class _JsonStreamResponse: + headers = {"content-type": "application/json"} + + def __init__(self, payload): + self.payload = payload + + def raise_for_status(self): + pass + + async def aread(self): + return json.dumps(self.payload).encode() + + +class _JsonStreamClient: + def __init__(self, stream_payload): + self.stream_payload = stream_payload + self.post_urls = [] + + async def post(self, url, **kwargs): + self.post_urls.append(url) + if "identity/token" in url: + return _JsonResponse({"access_token": "token", "expires_in": 3600}) + if url.endswith("/runs/stream"): + return _JsonStreamResponse(self.stream_payload) + raise AssertionError(url) + + +class TestWatsonxOrchestrateTransformation: + def test_get_api_base_url(self): + url = WatsonxOrchestrateTransformation.get_api_base_url( + "https://cpd.example.com/", + "1769134113217795", + ) + assert ( + url == "https://cpd.example.com/orchestrate/cpd/instances/1769134113217795" + ) + + def test_extract_text_from_a2a_params(self): + params = { + "message": { + "role": "user", + "parts": [ + {"kind": "text", "text": "Hello"}, + {"kind": "text", "text": "world"}, + ], + } + } + assert ( + WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params) + == "Hello world" + ) + + def test_extract_text_from_a2a_params_ignores_non_text_parts_with_text(self): + params = { + "message": { + "role": "user", + "parts": [ + {"kind": "data", "text": "metadata label", "data": {}}, + {"kind": "file", "text": "file label", "file": {}}, + {"kind": "text", "text": "Hello"}, + {"text": "legacy"}, + {"kind": "", "text": "empty-kind"}, + ], + } + } + assert ( + WatsonxOrchestrateTransformation.extract_text_from_a2a_params(params) + == "Hello legacy empty-kind" + ) + + def test_build_wxo_run_body_with_thread(self): + body = WatsonxOrchestrateTransformation.build_wxo_run_body( + wxo_agent_id="agent-uuid", + text="Hi", + thread_id="thread-1", + ) + assert body["agent_id"] == "agent-uuid" + assert body["thread_id"] == "thread-1" + assert body["message"]["content"][0]["response_type"] == "text" + assert body["message"]["content"][0]["text"] == "Hi" + + @pytest.mark.parametrize( + "result,expected", + [ + ( + { + "last_message": { + "content": [{"type": "text", "text": "from last_message"}] + } + }, + "from last_message", + ), + ( + { + "result": { + "data": { + "message": {"content": [{"text": "from nested result"}]} + } + } + }, + "from nested result", + ), + ({"results": "raw string"}, "raw string"), + ], + ) + def test_extract_text_from_wxo_result(self, result, expected): + assert ( + WatsonxOrchestrateTransformation.extract_text_from_wxo_result(result) + == expected + ) + + def test_build_a2a_message_response(self): + out = WatsonxOrchestrateTransformation.build_a2a_message_response( + "req-1", "answer" + ) + assert out["jsonrpc"] == "2.0" + assert out["id"] == "req-1" + assert out["result"]["kind"] == "message" + assert out["result"]["parts"][0]["text"] == "answer" + + def test_extract_text_from_a2a_message_response(self): + envelope = WatsonxOrchestrateTransformation.build_a2a_message_response( + "req-1", "answer" + ) + assert ( + WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response( + envelope + ) + == "answer" + ) + assert ( + WatsonxOrchestrateTransformation.extract_text_from_a2a_message_response( + {"result": {}} + ) + == "" + ) + + +def test_cp4d_token_ttl_from_absolute_expiration(): + wall = 1_750_000_000.0 + assert ( + WatsonxOrchestrateHandler._cp4d_token_ttl_seconds(1_750_003_600, wall) == 3600 + ) + assert WatsonxOrchestrateHandler._cp4d_token_ttl_seconds(1_749_999_000, wall) == 0 + + +@pytest.mark.asyncio +async def test_accumulate_wxo_sse_text_ignores_non_dict_json_events(): + response = _SSELines( + [ + "data: null", + "data: true", + 'data: {"results": "streamed text"}', + ] + ) + assert await WatsonxOrchestrateHandler._accumulate_wxo_sse_text(response) == ( + "streamed text" + ) + + +@pytest.mark.asyncio +async def test_short_lived_tokens_are_not_served_from_cache(): + client = _ShortTtlTokenClient() + token_1 = await WatsonxOrchestrateHandler._get_bearer_token( + cp4d_host="https://cpd.example.com", + auth_mode="ibm_cloud", + api_key="short-ttl-cache-key", + client=client, + ) + token_2 = await WatsonxOrchestrateHandler._get_bearer_token( + cp4d_host="https://cpd.example.com", + auth_mode="ibm_cloud", + api_key="short-ttl-cache-key", + client=client, + ) + assert token_1 == "token-1" + assert token_2 == "token-2" + assert client.calls == 2 + + +class _CP4DAuthClient: + def __init__(self, expiration): + self.expiration = expiration + self.calls = [] + + async def post(self, url, **kwargs): + self.calls.append((url, kwargs)) + return _JsonResponse({"token": "cp4d-token", "expiration": self.expiration}) + + +@pytest.mark.asyncio +async def test_cp4d_auth_posts_to_authorize_and_caches_token(): + client = _CP4DAuthClient(int(time.time()) + 3600) + token_1 = await WatsonxOrchestrateHandler._get_bearer_token( + cp4d_host="https://cpd.example.com/", + auth_mode="cp4d", + api_key="cp4d-e2e-cache-key", + username="cp4d-user", + client=client, + ) + token_2 = await WatsonxOrchestrateHandler._get_bearer_token( + cp4d_host="https://cpd.example.com/", + auth_mode="cp4d", + api_key="cp4d-e2e-cache-key", + username="cp4d-user", + client=client, + ) + + assert token_1 == "cp4d-token" + assert token_2 == "cp4d-token" + assert len(client.calls) == 1 + url, kwargs = client.calls[0] + assert url == "https://cpd.example.com/icp4d-api/v1/authorize" + assert kwargs["json"] == {"username": "cp4d-user", "api_key": "cp4d-e2e-cache-key"} + + +@pytest.mark.asyncio +async def test_cp4d_auth_requires_username(): + client = _CP4DAuthClient(int(time.time()) + 3600) + with pytest.raises(ValueError, match="username"): + await WatsonxOrchestrateHandler._get_bearer_token( + cp4d_host="https://cpd.example.com", + auth_mode="cp4d", + api_key="cp4d-missing-username-key", + username=None, + client=client, + ) + assert client.calls == [] + + +@pytest.mark.asyncio +async def test_expired_token_cache_entries_are_evicted(): + stale_key = "wxo-stale-cache-entry" + wxo_handler._token_cache[stale_key] = ("stale-token", time.monotonic() - 1) + + class _FreshTokenClient: + async def post(self, *args, **kwargs): + return _JsonResponse({"access_token": "fresh", "expires_in": 3600}) + + await WatsonxOrchestrateHandler._get_bearer_token( + cp4d_host="https://cpd.example.com", + auth_mode="ibm_cloud", + api_key="wxo-eviction-trigger-key", + client=_FreshTokenClient(), + ) + + assert stale_key not in wxo_handler._token_cache + + +@pytest.mark.asyncio +async def test_poll_run_raises_asyncio_timeout_when_never_terminal(): + class _NeverTerminalClient: + def __init__(self): + self.get_calls = 0 + + async def get(self, url, headers=None): + self.get_calls += 1 + return _JsonResponse({"status": "running"}) + + client = _NeverTerminalClient() + with pytest.raises(asyncio.TimeoutError): + await WatsonxOrchestrateHandler._poll_run( + base_url="https://cpd.example.com/orchestrate/cpd/instances/i", + run_id="run-1", + auth_headers={}, + client=client, + max_attempts=2, + interval_s=0, + ) + assert client.get_calls == 2 + + +@pytest.mark.asyncio +async def test_handle_streaming_polls_non_sse_json_until_complete(monkeypatch): + client = _JsonStreamClient({"status": "running", "run_id": "run-1"}) + poll_calls = [] + + async def poll_run(base_url, run_id, auth_headers, client, **kwargs): + poll_calls.append((base_url, run_id, auth_headers, client)) + return {"status": "completed", "results": "polled text"} + + monkeypatch.setattr( + WatsonxOrchestrateHandler, + "_http_client", + lambda timeout=90.0: client, + ) + monkeypatch.setattr(WatsonxOrchestrateHandler, "_poll_run", poll_run) + + params = { + "message": { + "parts": [ + {"kind": "text", "text": "Hello"}, + ], + } + } + litellm_params = { + "cp4d_host": "https://cpd.example.com", + "instance_id": "instance-id", + "wxo_agent_id": "agent-id", + "api_key": "pending-json-stream-cache-key", + "auth_mode": "ibm_cloud", + } + + events = [ + event + async for event in WatsonxOrchestrateHandler.handle_streaming( + request_id="req-1", + params=params, + litellm_params=litellm_params, + delay_ms=0, + ) + ] + artifact_text = "".join( + event["result"]["artifact"]["parts"][0]["text"] + for event in events + if event["result"].get("kind") == "artifact-update" + ) + + assert len(poll_calls) == 1 + assert poll_calls[0][1] == "run-1" + assert artifact_text == "polled text" + + +@pytest.mark.asyncio +async def test_handle_streaming_raises_for_non_sse_json_failure(monkeypatch): + client = _JsonStreamClient({"status": "failed", "run_id": "run-1"}) + monkeypatch.setattr( + WatsonxOrchestrateHandler, + "_http_client", + lambda timeout=90.0: client, + ) + + params = { + "message": { + "parts": [ + {"kind": "text", "text": "Hello"}, + ], + } + } + litellm_params = { + "cp4d_host": "https://cpd.example.com", + "instance_id": "instance-id", + "wxo_agent_id": "agent-id", + "api_key": "failed-json-stream-cache-key", + "auth_mode": "ibm_cloud", + } + + with pytest.raises(RuntimeError, match="non-success status 'failed'"): + async for _ in WatsonxOrchestrateHandler.handle_streaming( + request_id="req-1", + params=params, + litellm_params=litellm_params, + delay_ms=0, + ): + pass + + +@pytest.mark.asyncio +async def test_handle_streaming_does_not_fallback_on_invalid_json(monkeypatch): + client = _InvalidJsonStreamClient() + monkeypatch.setattr( + WatsonxOrchestrateHandler, + "_http_client", + lambda timeout=90.0: client, + ) + + params = { + "message": { + "parts": [ + {"kind": "text", "text": "Hello"}, + ], + } + } + litellm_params = { + "cp4d_host": "https://cpd.example.com", + "instance_id": "instance-id", + "wxo_agent_id": "agent-id", + "api_key": "invalid-json-stream-cache-key", + "auth_mode": "ibm_cloud", + } + + with pytest.raises(json.JSONDecodeError): + async for _ in WatsonxOrchestrateHandler.handle_streaming( + request_id="req-1", + params=params, + litellm_params=litellm_params, + ): + pass + + assert not any(url.endswith("/runs") for url in client.post_urls) + + +@pytest.mark.asyncio +async def test_handle_streaming_does_not_resubmit_run_on_poll_transport_error( + monkeypatch, +): + class _RunSubmissionClient: + def __init__(self): + self.post_urls = [] + + async def post(self, url, **kwargs): + self.post_urls.append(url) + if "identity/token" in url: + return _JsonResponse({"access_token": "token", "expires_in": 3600}) + if url.endswith("/runs/stream"): + return _JsonStreamResponse({"status": "running", "run_id": "run-1"}) + if url.endswith("/runs"): + return _JsonResponse({"status": "completed", "results": "duplicate"}) + raise AssertionError(url) + + client = _RunSubmissionClient() + + async def poll_run(base_url, run_id, auth_headers, client, **kwargs): + raise httpx.ConnectError("connection reset during poll") + + monkeypatch.setattr( + WatsonxOrchestrateHandler, + "_http_client", + lambda timeout=90.0: client, + ) + monkeypatch.setattr(WatsonxOrchestrateHandler, "_poll_run", poll_run) + + params = {"message": {"parts": [{"kind": "text", "text": "Hello"}]}} + litellm_params = { + "cp4d_host": "https://cpd.example.com", + "instance_id": "instance-id", + "wxo_agent_id": "agent-id", + "api_key": "poll-transport-error-cache-key", + "auth_mode": "ibm_cloud", + } + + with pytest.raises(httpx.TransportError): + async for _ in WatsonxOrchestrateHandler.handle_streaming( + request_id="req-1", + params=params, + litellm_params=litellm_params, + delay_ms=0, + ): + pass + + assert not any(url.endswith("/runs") for url in client.post_urls) + + +@pytest.mark.asyncio +async def test_handle_streaming_falls_back_when_initial_post_fails(monkeypatch): + class _StreamPostFailsClient: + def __init__(self): + self.post_urls = [] + + async def post(self, url, **kwargs): + self.post_urls.append(url) + if "identity/token" in url: + return _JsonResponse({"access_token": "token", "expires_in": 3600}) + if url.endswith("/runs/stream"): + raise httpx.ConnectError("cannot reach stream endpoint") + if url.endswith("/runs"): + return _JsonResponse({"status": "completed", "results": "fallback"}) + raise AssertionError(url) + + client = _StreamPostFailsClient() + monkeypatch.setattr( + WatsonxOrchestrateHandler, + "_http_client", + lambda timeout=90.0: client, + ) + + params = {"message": {"parts": [{"kind": "text", "text": "Hello"}]}} + litellm_params = { + "cp4d_host": "https://cpd.example.com", + "instance_id": "instance-id", + "wxo_agent_id": "agent-id", + "api_key": "stream-post-fails-cache-key", + "auth_mode": "ibm_cloud", + } + + events = [ + event + async for event in WatsonxOrchestrateHandler.handle_streaming( + request_id="req-1", + params=params, + litellm_params=litellm_params, + delay_ms=0, + ) + ] + artifact_text = "".join( + event["result"]["artifact"]["parts"][0]["text"] + for event in events + if event["result"].get("kind") == "artifact-update" + ) + + assert artifact_text == "fallback" + assert sum(url.endswith("/runs") for url in client.post_urls) == 1 + + +def test_config_manager_returns_wxo_provider(): + config = A2AProviderConfigManager.get_provider_config( + custom_llm_provider="watsonx_orchestrate" + ) + assert config is not None + assert config.__class__.__name__ == "WatsonxOrchestrateA2AConfig" + + +def test_wxo_dashboard_auth_fields(): + fields_path = ( + Path(__file__).resolve().parents[5] + / "litellm/proxy/public_endpoints/agent_create_fields.json" + ) + agent_fields = json.loads(fields_path.read_text()) + wxo_agent = next( + agent for agent in agent_fields if agent["agent_type"] == "watsonx_orchestrate" + ) + fields_by_key = {field["key"]: field for field in wxo_agent["credential_fields"]} + + assert fields_by_key["auth_mode"]["default_value"] == "cp4d" + # Username is CP4D-only; UI does not require it so ibm_cloud users are not blocked. + assert fields_by_key["username"]["required"] is False + assert "cp4d" in fields_by_key["username"]["tooltip"].lower()