diff --git a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py index 074b9d199b2..74b027d77ff 100644 --- a/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py +++ b/litellm/a2a_protocol/providers/watsonx_orchestrate/handler.py @@ -325,41 +325,10 @@ class WatsonxOrchestrateHandler: stream=True, ) response.raise_for_status() - 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 - except httpx.TransportError as exc: verbose_logger.warning( - f"WXO: Streaming transport failed ({exc!r}), " - "falling back to non-streaming + fake streaming", + 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( @@ -381,3 +350,30 @@ class WatsonxOrchestrateHandler: 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/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 index a13120104ba..df2876900b6 100644 --- 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 @@ -3,6 +3,7 @@ import os import sys from pathlib import Path +import httpx import pytest sys.path.insert(0, os.path.abspath("../../../../..")) @@ -370,6 +371,108 @@ async def test_handle_streaming_does_not_fallback_on_invalid_json(monkeypatch): 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"