From 39748fd4a3327f6e7e0796fe53a49d9e14ab62b9 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 2 Mar 2026 17:21:47 +0530 Subject: [PATCH] feat(responses): add _aresponses_websocket function and HTTP handler support for WebSocket mode Also fix pyrightconfig.json to use the conda venv for type checking, and remove redundant inline import of ResponsesAPIRequestUtils that was confusing pyright. Co-Authored-By: Claude Sonnet 4.6 --- litellm/llms/custom_httpx/llm_http_handler.py | 93 ++++++++++++++++ litellm/responses/main.py | 101 +++++++++++++++++- pyrightconfig.json | 4 +- 3 files changed, 195 insertions(+), 3 deletions(-) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index d6fdc58099f..29d494dbb50 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -69,6 +69,7 @@ from litellm.responses.streaming_iterator import ( BaseResponsesAPIStreamingIterator, MockResponsesAPIStreamingIterator, ResponsesAPIStreamingIterator, + ResponsesWebSocketStreaming, SyncResponsesAPIStreamingIterator, ) from litellm.types.containers.main import ( @@ -4731,6 +4732,98 @@ class BaseLLMHTTPHandler: f"Unexpected error while closing WebSocket: {close_error}" ) + async def async_responses_websocket( + self, + model: str, + websocket: Any, + logging_obj: LiteLLMLoggingObj, + responses_api_provider_config: BaseResponsesAPIConfig, + api_base: Optional[str] = None, + api_key: Optional[str] = None, + timeout: Optional[float] = None, + user_api_key_dict: Optional[Any] = None, + litellm_metadata: Optional[Dict[str, Any]] = None, + ): + """ + Handles Responses API WebSocket mode. + + Opens a persistent WebSocket to the provider's /v1/responses endpoint + and proxies response.create events bidirectionally for lower-latency + agentic workflows. + """ + import websockets + from websockets.asyncio.client import ClientConnection + + litellm_params = GenericLiteLLMParams() + headers = responses_api_provider_config.validate_environment( + headers={}, + model=model, + litellm_params=litellm_params, + ) + if api_key: + headers["Authorization"] = f"Bearer {api_key}" + + http_url = responses_api_provider_config.get_complete_url( + api_base=api_base, + litellm_params={}, + ) + # /responses -> wss:// URL + ws_url = http_url.replace("https://", "wss://").replace("http://", "ws://") + + try: + ssl_context = get_shared_realtime_ssl_context() + if ws_url.startswith("wss://") and ssl_context is False: + ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) + ssl_context.check_hostname = False + ssl_context.verify_mode = ssl.CERT_NONE + + logging_obj.pre_call( + input=None, + api_key=api_key or "", + additional_args={ + "api_base": ws_url, + "headers": headers, + "complete_input_dict": {"mode": "responses_websocket"}, + }, + ) + + async with websockets.connect( # type: ignore + ws_url, + additional_headers=headers, + max_size=REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES, + ssl=ssl_context, + ) as backend_ws: + _request_data: Dict[str, Any] = {} + if litellm_metadata: + _request_data["litellm_metadata"] = litellm_metadata + streaming = ResponsesWebSocketStreaming( + websocket=websocket, + backend_ws=cast(ClientConnection, backend_ws), + logging_obj=logging_obj, + user_api_key_dict=user_api_key_dict, + request_data=_request_data, + ) + await streaming.bidirectional_forward() + + except websockets.exceptions.InvalidStatusCode as e: # type: ignore + verbose_logger.exception(f"Error connecting to responses WS backend: {e}") + await websocket.close(code=e.status_code, reason=str(e)) + except Exception as e: + verbose_logger.exception(f"Error in responses WS: {e}") + try: + await websocket.close( + code=1011, reason=f"Internal server error: {str(e)}" + ) + except RuntimeError as close_error: + if "already completed" in str(close_error) or "websocket.close" in str( + close_error + ): + pass + else: + raise Exception( + f"Unexpected error while closing WebSocket: {close_error}" + ) + def image_edit_handler( self, model: str, diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 05fd6026af2..6bdeaf66e62 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -51,6 +51,8 @@ if TYPE_CHECKING: from litellm.types.llms.openai import ResponseText # type: ignore else: ResponseText = str # Fallback for ResponseText import +from litellm.litellm_core_utils.get_litellm_params import get_litellm_params +from litellm.secret_managers.main import get_secret_str from litellm.types.responses.main import * from litellm.types.router import GenericLiteLLMParams from litellm.utils import ProviderConfigManager, client @@ -182,8 +184,6 @@ async def aresponses_api_with_mcp( mcp_server_auth_headers: Optional[Dict[str, Dict[str, str]]] = None secret_fields = kwargs.get("secret_fields") if secret_fields and isinstance(secret_fields, dict): - from litellm.responses.utils import ResponsesAPIRequestUtils - mcp_auth_header, mcp_server_auth_headers, _, _ = ( ResponsesAPIRequestUtils.extract_mcp_headers_from_request( secret_fields=secret_fields, tools=tools @@ -1651,3 +1651,100 @@ def compact_responses( completion_kwargs=local_vars, extra_kwargs=kwargs, ) + + +# --------------------------------------------------------------------------- +# Responses API WebSocket mode +# --------------------------------------------------------------------------- + + +def _build_litellm_metadata_for_ws(kwargs: dict) -> dict: + metadata: dict = {**(kwargs.get("litellm_metadata") or {})} + guardrails = ( + (kwargs.get("metadata") or {}).get("guardrails") + or kwargs.get("guardrails") + or [] + ) + if guardrails: + metadata["guardrails"] = guardrails + return metadata + + +@client +async def _aresponses_websocket( + model: str, + websocket: Any, + api_base: Optional[str] = None, + api_key: Optional[str] = None, + timeout: Optional[float] = None, + **kwargs, +): + """ + Private function to handle the Responses API WebSocket mode. + + For PROXY use only. + + Resolves the LLM provider from ``model``, looks up the matching + ``BaseResponsesAPIConfig``, and hands off to + ``BaseLLMHTTPHandler.async_responses_websocket``. + """ + litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore + user = kwargs.get("user", None) + litellm_params = GenericLiteLLMParams(**kwargs) + litellm_params_dict = get_litellm_params(**kwargs) + + model, _custom_llm_provider, dynamic_api_key, dynamic_api_base = ( + litellm.get_llm_provider( + model=model, + api_base=api_base, + api_key=api_key, + ) + ) + + litellm_logging_obj.update_environment_variables( + model=model, + user=user, + optional_params={}, + litellm_params=litellm_params_dict, + custom_llm_provider=_custom_llm_provider, + ) + + responses_api_provider_config: Optional[BaseResponsesAPIConfig] = None + if _custom_llm_provider is not None: + responses_api_provider_config = ( + ProviderConfigManager.get_provider_responses_api_config( + model=model, + provider=litellm.LlmProviders(_custom_llm_provider), + ) + ) + + if responses_api_provider_config is None: + raise ValueError( + f"Responses API WebSocket mode is not supported for provider: {_custom_llm_provider}" + ) + + resolved_api_base = ( + dynamic_api_base + or litellm_params.api_base + or litellm.api_base + or None + ) + resolved_api_key = ( + dynamic_api_key + or litellm_params.api_key + or litellm.api_key + or litellm.openai_key + or get_secret_str("OPENAI_API_KEY") + ) + + await base_llm_http_handler.async_responses_websocket( + model=model, + websocket=websocket, + logging_obj=litellm_logging_obj, + responses_api_provider_config=responses_api_provider_config, + api_base=resolved_api_base, + api_key=resolved_api_key, + timeout=timeout, + user_api_key_dict=kwargs.get("user_api_key_dict"), + litellm_metadata=_build_litellm_metadata_for_ws(kwargs), + ) diff --git a/pyrightconfig.json b/pyrightconfig.json index f930e44d305..ec0a1823038 100644 --- a/pyrightconfig.json +++ b/pyrightconfig.json @@ -2,6 +2,8 @@ "ignore": [], "exclude": ["**/node_modules", "**/__pycache__", "litellm/types/utils.py", "litellm/proxy/_types.py"], "reportMissingImports": false, - "reportPrivateImportUsage": false + "reportPrivateImportUsage": false, + "venvPath": "/Users/sameerkankute/miniconda3/envs", + "venv": "litellm-dev" } \ No newline at end of file