mirror of
https://github.com/tiennm99/litellm.git
synced 2026-08-03 10:23:28 +00:00
feat(a2a): watsonx Orchestrate agent provider (#29410)
* feat(a2a): add watsonx Orchestrate agent provider Bridge A2A message/send to WXO runs API (CP4D and IBM Cloud IAM auth), with dashboard agent type metadata and unit tests for transformations. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(a2a): use shared httpx client and cache WXO auth tokens Route WXO streaming through get_async_httpx_client (TLS verification enabled). Cache bearer tokens with TTL buffer. Extract A2A reply text via a dedicated helper instead of hard-coded JSON paths. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(a2a): treat CP4D token expiration as absolute Unix time CP4D /authorize returns expiration as epoch seconds, not TTL. Compute remaining lifetime against wall clock so cached tokens refresh before expiry. Co-authored-by: Cursor <cursoragent@cursor.com> * style(a2a): black-format watsonx orchestrate handler for CI py312 Co-authored-by: Cursor <cursoragent@cursor.com> * Fix watsonx orchestrate edge cases * Fix WXO streaming fallback error handling * Fix watsonx orchestrate run completion handling * fix(a2a): make WXO username optional in agent create UI Username is only required for cp4d auth; ibm_cloud uses api_key alone. Backend still validates username when auth_mode is cp4d. Co-authored-by: Cursor <cursoragent@cursor.com> * test(a2a): align WXO dashboard field test with optional username Username is not required in agent_create_fields.json; backend validates for cp4d auth_mode only. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(a2a): send Accept header on WXO streaming run request * fix(a2a/wxo): scope streaming transport fallback to initial POST only Narrow the httpx.TransportError fallback in handle_streaming so it only covers the initial POST /runs/stream. Errors during polling or SSE consumption now propagate instead of triggering handle_non_streaming, which would have submitted a duplicate WXO run for the same request. * refactor(a2a/wxo): type run-param extraction and use text response_type Return a typed WXORequestParams NamedTuple from _extract_litellm_params instead of a positional tuple so call sites read params by name, and send the user message with response_type 'text' so the run body is valid across all WXO agent configurations rather than the search-specific type. * fix(a2a/wxo): evict expired token cache entries and raise asyncio.TimeoutError on poll timeout --------- Co-authored-by: Cursor <cursoragent@cursor.com> Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Sameer Kankute
co-authored by
Cursor
mateo-berri
Cursor
mateo-berri
parent
f48a87ef12
commit
c8bcfbb20c
@@ -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
|
||||
|
||||
@@ -0,0 +1,3 @@
|
||||
"""
|
||||
IBM watsonx Orchestrate (WXO) A2A provider.
|
||||
"""
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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}"
|
||||
)
|
||||
@@ -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/<INSTANCE_ID>",
|
||||
"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"
|
||||
}
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
+593
@@ -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()
|
||||
Reference in New Issue
Block a user