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