mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
refactor: extract large functions to fix PLR0915 ruff lint errors
This commit is contained in:
Ishaan Jaffer
parent
4711cbb85e
commit
091e8ec738
3 changed files with 267 additions and 230 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue