From 091e8ec7387b2b0c18e57e8ca966c3b1e2bdc330 Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Tue, 24 Feb 2026 09:01:37 -0800 Subject: [PATCH] refactor: extract large functions to fix PLR0915 ruff lint errors --- .../litellm_core_utils/realtime_streaming.py | 248 +++++++++--------- litellm/proxy/health_check.py | 122 ++++----- litellm/proxy/proxy_server.py | 127 +++++---- 3 files changed, 267 insertions(+), 230 deletions(-) diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index b42bf9e5711..759eaf60035 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -203,6 +203,129 @@ class RealTimeStreaming: return True return False + async def _handle_provider_config_message(self, raw_response) -> None: + """Process a backend message when a provider_config is set (transformed path).""" + returned_object = self.provider_config.transform_realtime_response( # type: ignore[union-attr] + raw_response, + self.model, + self.logging_obj, + realtime_response_transform_input={ + "session_configuration_request": self.session_configuration_request, + "current_output_item_id": self.current_output_item_id, + "current_response_id": self.current_response_id, + "current_delta_chunks": self.current_delta_chunks, + "current_conversation_id": self.current_conversation_id, + "current_item_chunks": self.current_item_chunks, + "current_delta_type": self.current_delta_type, + }, + ) + + transformed_response = returned_object["response"] + self.current_output_item_id = returned_object["current_output_item_id"] + self.current_response_id = returned_object["current_response_id"] + self.current_delta_chunks = returned_object["current_delta_chunks"] + self.current_conversation_id = returned_object["current_conversation_id"] + self.current_item_chunks = returned_object["current_item_chunks"] + self.current_delta_type = returned_object["current_delta_type"] + self.session_configuration_request = returned_object["session_configuration_request"] + events = ( + transformed_response + if isinstance(transformed_response, list) + else [transformed_response] + ) + for event in events: + ## GUARDRAIL: inject create_response=false on session.created + if isinstance(event, dict) and event.get("type") == "session.created": + if self._has_realtime_guardrails(): + await self.backend_ws.send( + json.dumps( + { + "type": "session.update", + "session": { + "turn_detection": { + "type": "server_vad", + "create_response": False, + } + }, + } + ) + ) + for event in events: + event_str = json.dumps(event) + ## GUARDRAIL: run on transcription events in provider_config path too + if ( + isinstance(event, dict) + and event.get("type") + == "conversation.item.input_audio_transcription.completed" + ): + transcript = event.get("transcript", "") + self.store_message(event_str) + await self.websocket.send_text(event_str) + blocked = await self.run_realtime_guardrails( + transcript, item_id=event.get("item_id") + ) + if not blocked: + await self.backend_ws.send( + json.dumps({"type": "response.create"}) + ) + continue + ## LOGGING + self.store_message(event_str) + await self.websocket.send_text(event_str) + + async def _handle_raw_backend_message(self, raw_response) -> bool: + """Process a backend message without provider_config (raw path). + + Returns True if the caller should skip the default store+forward (i.e. continue the loop). + """ + try: + event_obj = json.loads(raw_response) + + if event_obj.get("type") == "session.created": + # If any realtime guardrails are registered, proactively + # set create_response=false so the LLM never auto-responds + # before our guardrail has a chance to run. + if self._has_realtime_guardrails(): + await self.backend_ws.send( + json.dumps( + { + "type": "session.update", + "session": { + "turn_detection": { + "type": "server_vad", + "create_response": False, + } + }, + } + ) + ) + verbose_logger.debug( + "[realtime guardrail] injected create_response=false into session" + ) + + if ( + event_obj.get("type") + == "conversation.item.input_audio_transcription.completed" + ): + transcript = event_obj.get("transcript", "") + ## LOGGING — must happen before continue below + self.store_message(raw_response) + # Forward transcript to client so user sees what they said + await self.websocket.send_text(raw_response) + blocked = await self.run_realtime_guardrails( + transcript, + item_id=event_obj.get("item_id"), + ) + if not blocked: + # Clean — trigger LLM response + await self.backend_ws.send( + json.dumps({"type": "response.create"}) + ) + return True + except (json.JSONDecodeError, AttributeError): + pass + return False + async def backend_to_client_send_messages(self): import websockets @@ -216,128 +339,11 @@ class RealTimeStreaming: raw_response = await self.backend_ws.recv() # type: ignore[assignment] if self.provider_config: - returned_object = self.provider_config.transform_realtime_response( - raw_response, - self.model, - self.logging_obj, - realtime_response_transform_input={ - "session_configuration_request": self.session_configuration_request, - "current_output_item_id": self.current_output_item_id, - "current_response_id": self.current_response_id, - "current_delta_chunks": self.current_delta_chunks, - "current_conversation_id": self.current_conversation_id, - "current_item_chunks": self.current_item_chunks, - "current_delta_type": self.current_delta_type, - }, - ) - - transformed_response = returned_object["response"] - self.current_output_item_id = returned_object[ - "current_output_item_id" - ] - self.current_response_id = returned_object["current_response_id"] - self.current_delta_chunks = returned_object["current_delta_chunks"] - self.current_conversation_id = returned_object[ - "current_conversation_id" - ] - self.current_item_chunks = returned_object["current_item_chunks"] - self.current_delta_type = returned_object["current_delta_type"] - self.session_configuration_request = returned_object[ - "session_configuration_request" - ] - events = ( - transformed_response - if isinstance(transformed_response, list) - else [transformed_response] - ) - for event in events: - ## GUARDRAIL: inject create_response=false on session.created - if isinstance(event, dict) and event.get("type") == "session.created": - if self._has_realtime_guardrails(): - await self.backend_ws.send( - json.dumps( - { - "type": "session.update", - "session": { - "turn_detection": { - "type": "server_vad", - "create_response": False, - } - }, - } - ) - ) - for event in events: - event_str = json.dumps(event) - ## GUARDRAIL: run on transcription events in provider_config path too - if ( - isinstance(event, dict) - and event.get("type") - == "conversation.item.input_audio_transcription.completed" - ): - transcript = event.get("transcript", "") - self.store_message(event_str) - await self.websocket.send_text(event_str) - blocked = await self.run_realtime_guardrails( - transcript, item_id=event.get("item_id") - ) - if not blocked: - await self.backend_ws.send( - json.dumps({"type": "response.create"}) - ) - continue - ## LOGGING - self.store_message(event_str) - await self.websocket.send_text(event_str) - + await self._handle_provider_config_message(raw_response) else: - ## GUARDRAIL: intercept transcription events before triggering LLM - try: - event_obj = json.loads(raw_response) - - if event_obj.get("type") == "session.created": - # If any realtime guardrails are registered, proactively - # set create_response=false so the LLM never auto-responds - # before our guardrail has a chance to run. - if self._has_realtime_guardrails(): - await self.backend_ws.send( - json.dumps( - { - "type": "session.update", - "session": { - "turn_detection": { - "type": "server_vad", - "create_response": False, - } - }, - } - ) - ) - verbose_logger.debug( - "[realtime guardrail] injected create_response=false into session" - ) - - if ( - event_obj.get("type") - == "conversation.item.input_audio_transcription.completed" - ): - transcript = event_obj.get("transcript", "") - ## LOGGING — must happen before continue below - self.store_message(raw_response) - # Forward transcript to client so user sees what they said - await self.websocket.send_text(raw_response) - blocked = await self.run_realtime_guardrails( - transcript, - item_id=event_obj.get("item_id"), - ) - if not blocked: - # Clean — trigger LLM response - await self.backend_ws.send( - json.dumps({"type": "response.create"}) - ) - continue - except (json.JSONDecodeError, AttributeError): - pass + handled = await self._handle_raw_backend_message(raw_response) + if handled: + continue ## LOGGING self.store_message(raw_response) await self.websocket.send_text(raw_response) diff --git a/litellm/proxy/health_check.py b/litellm/proxy/health_check.py index 4777f64405d..de4517efc88 100644 --- a/litellm/proxy/health_check.py +++ b/litellm/proxy/health_check.py @@ -11,7 +11,7 @@ from typing import List, Optional import litellm logger = logging.getLogger(__name__) -from litellm.constants import HEALTH_CHECK_TIMEOUT_SECONDS, DEFAULT_HEALTH_CHECK_PROMPT +from litellm.constants import DEFAULT_HEALTH_CHECK_PROMPT, HEALTH_CHECK_TIMEOUT_SECONDS ILLEGAL_DISPLAY_PARAMS = [ "messages", @@ -98,6 +98,66 @@ async def run_with_timeout(task, timeout): return {"error": "Timeout exceeded"} +async def _run_model_health_check(model: dict): + litellm_params = model["litellm_params"] + model_info = model.get("model_info", {}) + mode = model_info.get("mode", None) + litellm_params = _update_litellm_params_for_health_check(model_info, litellm_params) + timeout = model_info.get("health_check_timeout") or HEALTH_CHECK_TIMEOUT_SECONDS + + return await run_with_timeout( + litellm.ahealth_check( + litellm_params, + mode=mode, + prompt=DEFAULT_HEALTH_CHECK_PROMPT, + input=["test from litellm"], + ), + timeout, + ) + + +async def _run_health_checks_with_bounded_concurrency( + models: list, concurrency_limit: int +) -> tuple[list, int]: + """ + Run health checks with at most `concurrency_limit` active tasks. + Preserves result ordering to match `models`. + """ + results: list = [None] * len(models) + tasks_to_index: dict[asyncio.Task, int] = {} + model_iter = iter(enumerate(models)) + peak_in_flight = 0 + + def _schedule_next() -> bool: + nonlocal peak_in_flight + try: + idx, next_model = next(model_iter) + except StopIteration: + return False + task = asyncio.create_task(_run_model_health_check(next_model)) + tasks_to_index[task] = idx + peak_in_flight = max(peak_in_flight, len(tasks_to_index)) + return True + + for _ in range(min(concurrency_limit, len(models))): + _schedule_next() + + while tasks_to_index: + done, _ = await asyncio.wait( + set(tasks_to_index.keys()), + return_when=asyncio.FIRST_COMPLETED, + ) + for task in done: + idx = tasks_to_index.pop(task) + try: + results[idx] = task.result() + except Exception as e: + results[idx] = e + _schedule_next() + + return results, peak_in_flight + + async def _perform_health_check( model_list: list, details: Optional[bool] = True, @@ -115,66 +175,6 @@ async def _perform_health_check( cycle_id = instrumentation_context.get("cycle_id", "unknown") source = instrumentation_context.get("source", "unknown") - async def _run_model_health_check(model: dict): - litellm_params = model["litellm_params"] - model_info = model.get("model_info", {}) - mode = model_info.get("mode", None) - litellm_params = _update_litellm_params_for_health_check( - model_info, litellm_params - ) - timeout = model_info.get("health_check_timeout") or HEALTH_CHECK_TIMEOUT_SECONDS - - return await run_with_timeout( - litellm.ahealth_check( - litellm_params, - mode=mode, - prompt=DEFAULT_HEALTH_CHECK_PROMPT, - input=["test from litellm"], - ), - timeout, - ) - - async def _run_health_checks_with_bounded_concurrency( - models: list, concurrency_limit: int - ) -> tuple[list, int]: - """ - Run health checks with at most `concurrency_limit` active tasks. - Preserves result ordering to match `models`. - """ - results: list = [None] * len(models) - tasks_to_index: dict[asyncio.Task, int] = {} - model_iter = iter(enumerate(models)) - peak_in_flight = 0 - - def _schedule_next() -> bool: - nonlocal peak_in_flight - try: - idx, next_model = next(model_iter) - except StopIteration: - return False - task = asyncio.create_task(_run_model_health_check(next_model)) - tasks_to_index[task] = idx - peak_in_flight = max(peak_in_flight, len(tasks_to_index)) - return True - - for _ in range(min(concurrency_limit, len(models))): - _schedule_next() - - while tasks_to_index: - done, _ = await asyncio.wait( - set(tasks_to_index.keys()), - return_when=asyncio.FIRST_COMPLETED, - ) - for task in done: - idx = tasks_to_index.pop(task) - try: - results[idx] = task.result() - except Exception as e: - results[idx] = e - _schedule_next() - - return results, peak_in_flight - dispatch_mode = "unbounded" peak_in_flight = 0 if isinstance(max_concurrency, int) and max_concurrency > 0: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ddab560aaa3..b3b6bf0ccf7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -713,7 +713,7 @@ async def _initialize_shared_aiohttp_session(): try: from aiohttp import ClientSession, TCPConnector - connector_kwargs = { + connector_kwargs: Dict[str, Any] = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, } @@ -1959,6 +1959,65 @@ def _rss_mb_for_log() -> str: return f"{rss_mb:.2f}" +async def _run_direct_health_check_with_instrumentation( + model_list: list, + details: Optional[bool], + max_concurrency: Optional[int], + instrumentation_context: dict, +): + try: + return await perform_health_check( + model_list=model_list, + details=details, + max_concurrency=max_concurrency, + instrumentation_context=instrumentation_context, + ) + except TypeError as e: + if "instrumentation_context" not in str(e): + raise + # Backward compatibility for monkeypatched or wrapped callables + # that do not accept instrumentation_context. + return await perform_health_check( + model_list=model_list, + details=details, + max_concurrency=max_concurrency, + ) + + +def _schedule_background_health_check_db_save( + prisma_client, + shared_health_manager, + model_list: list, + healthy_endpoints: list, + unhealthy_endpoints: list, +): + """Fire-and-forget: persist health check results to DB if prisma is available.""" + if prisma_client is None: + return + import time as time_module + + from litellm.proxy.health_endpoints._health_endpoints import ( + _save_background_health_checks_to_db, + ) + + checked_by = ( + shared_health_manager.pod_id + if shared_health_manager is not None + else "background_health_check" + ) + start_time = time_module.time() + asyncio.create_task( + _save_background_health_checks_to_db( + prisma_client, + model_list, + healthy_endpoints, + unhealthy_endpoints, + start_time, + checked_by=checked_by, + ) + ) + + async def _run_background_health_check(): """ Periodically run health checks in the background on the endpoints. @@ -2053,25 +2112,6 @@ async def _run_background_health_check(): "cycle_id": cycle_id, } - async def _run_direct_health_check_with_instrumentation(): - try: - return await perform_health_check( - model_list=_llm_model_list, - details=health_check_details, - max_concurrency=health_check_concurrency, - instrumentation_context=instrumentation_context, - ) - except TypeError as e: - if "instrumentation_context" not in str(e): - raise - # Backward compatibility for monkeypatched or wrapped callables - # that do not accept instrumentation_context. - return await perform_health_check( - model_list=_llm_model_list, - details=health_check_details, - max_concurrency=health_check_concurrency, - ) - # Use shared health check if available, otherwise fall back to direct health check # Convert health_check_details to bool for perform_shared_health_check (defaults to True if None) details_bool = ( @@ -2094,11 +2134,21 @@ async def _run_background_health_check(): str(e), ) healthy_endpoints, unhealthy_endpoints = ( - await _run_direct_health_check_with_instrumentation() + await _run_direct_health_check_with_instrumentation( + _llm_model_list, + health_check_details, + health_check_concurrency, + instrumentation_context, + ) ) else: healthy_endpoints, unhealthy_endpoints = ( - await _run_direct_health_check_with_instrumentation() + await _run_direct_health_check_with_instrumentation( + _llm_model_list, + health_check_details, + health_check_concurrency, + instrumentation_context, + ) ) # Update the global variable with the health check results @@ -2127,32 +2177,13 @@ async def _run_background_health_check(): ) # Save background health checks to database (non-blocking) - if prisma_client is not None: - import time as time_module - - from litellm.proxy.health_endpoints._health_endpoints import ( - _save_background_health_checks_to_db, - ) - - # Use pod_id or a system identifier for checked_by if shared health check is enabled - checked_by = None - if shared_health_manager is not None: - checked_by = shared_health_manager.pod_id - else: - # Use a system identifier for background health checks - checked_by = "background_health_check" - - start_time = time_module.time() - asyncio.create_task( - _save_background_health_checks_to_db( - prisma_client, - _llm_model_list, - healthy_endpoints, - unhealthy_endpoints, - start_time, - checked_by=checked_by, - ) - ) + _schedule_background_health_check_db_save( + prisma_client, + shared_health_manager, + _llm_model_list, + healthy_endpoints, + unhealthy_endpoints, + ) await asyncio.sleep(health_check_interval)