scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-05 08:07:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 136

fix: fix websocket import

Browse source
This commit is contained in:
Krrish Dholakia 2026-02-02 09:01:06 -08:00
parent faff9d1dc5
commit 975bea3933
1 changed files with 63 additions and 52 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

115
litellm/proxy/proxy_server.py
View file

@ -346,9 +346,7 @@ from litellm.proxy.management_endpoints.fallback_management_endpoints import (
from litellm.proxy.management_endpoints.internal_user_endpoints import (
router as internal_user_router,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import (
user_update,
)
from litellm.proxy.management_endpoints.internal_user_endpoints import user_update
from litellm.proxy.management_endpoints.key_management_endpoints import (
delete_verification_tokens,
duration_in_seconds,
@ -403,9 +401,7 @@ from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router
from litellm.proxy.openai_files_endpoints.files_endpoints import (
router as openai_files_router,
)
from litellm.proxy.openai_files_endpoints.files_endpoints import (
set_files_config,
)
from litellm.proxy.openai_files_endpoints.files_endpoints import set_files_config
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
passthrough_endpoint_router,
)
@ -501,9 +497,7 @@ from litellm.types.proxy.management_endpoints.ui_sso import (
LiteLLM_UpperboundKeyGenerateParams,
)
from litellm.types.realtime import RealtimeQueryParams
from litellm.types.router import (
DeploymentTypedDict,
)
from litellm.types.router import DeploymentTypedDict
from litellm.types.router import ModelInfo as RouterModelInfo
from litellm.types.router import (
RouterGeneralSettings,
@ -590,9 +584,9 @@ except ImportError:
server_root_path = get_server_root_path()
_license_check = LicenseCheck()
premium_user: bool = _license_check.is_premium()
premium_user_data: Optional[
"EnterpriseLicenseData"
] = _license_check.airgapped_license_data
premium_user_data: Optional["EnterpriseLicenseData"] = (
_license_check.airgapped_license_data
)
global_max_parallel_request_retries_env: Optional[str] = os.getenv(
"LITELLM_GLOBAL_MAX_PARALLEL_REQUEST_RETRIES"
)
@ -1257,9 +1251,9 @@ master_key: Optional[str] = None
config_agents: Optional[List[AgentConfig]] = None
otel_logging = False
prisma_client: Optional[PrismaClient] = None
shared_aiohttp_session: Optional[
"ClientSession"
] = None # Global shared session for connection reuse
shared_aiohttp_session: Optional["ClientSession"] = (
None # Global shared session for connection reuse
)
user_api_key_cache = DualCache(
default_in_memory_ttl=UserAPIKeyCacheTTLEnum.in_memory_cache_ttl.value
)
@ -1267,11 +1261,13 @@ model_max_budget_limiter = _PROXY_VirtualKeyModelMaxBudgetLimiter(
dual_cache=user_api_key_cache
)
litellm.logging_callback_manager.add_litellm_callback(model_max_budget_limiter)
redis_usage_cache: Optional[
RedisCache
] = None # redis cache used for tracking spend, tpm/rpm limits
redis_usage_cache: Optional[RedisCache] = (
None # redis cache used for tracking spend, tpm/rpm limits
)
polling_via_cache_enabled: Union[Literal["all"], List[str], bool] = False
native_background_mode: List[str] = [] # Models that should use native provider background mode instead of polling
native_background_mode: List[str] = (
[]
) # Models that should use native provider background mode instead of polling
polling_cache_ttl: int = 3600 # Default 1 hour TTL for polling cache
user_custom_auth = None
user_custom_key_generate = None
@ -1609,9 +1605,9 @@ async def update_cache( # noqa: PLR0915
_id = "team_id:{}".format(team_id)
try:
# Fetch the existing cost for the given user
existing_spend_obj: Optional[
LiteLLM_TeamTable
] = await user_api_key_cache.async_get_cache(key=_id)
existing_spend_obj: Optional[LiteLLM_TeamTable] = (
await user_api_key_cache.async_get_cache(key=_id)
)
if existing_spend_obj is None:
# do nothing if team not in api key cache
return
@ -4071,10 +4067,10 @@ class ProxyConfig:
)
try:
guardrails_in_db: List[
Guardrail
] = await GuardrailRegistry.get_all_guardrails_from_db(
prisma_client=prisma_client
guardrails_in_db: List[Guardrail] = (
await GuardrailRegistry.get_all_guardrails_from_db(
prisma_client=prisma_client
)
)
verbose_proxy_logger.debug(
"guardrails from the DB %s", str(guardrails_in_db)
@ -4433,9 +4429,9 @@ async def initialize( # noqa: PLR0915
user_api_base = api_base
dynamic_config[user_model]["api_base"] = api_base
if api_version:
os.environ[
"AZURE_API_VERSION"
] = api_version # set this for azure - litellm can read this from the env
os.environ["AZURE_API_VERSION"] = (
api_version # set this for azure - litellm can read this from the env
)
if max_tokens: # model-specific param
dynamic_config[user_model]["max_tokens"] = max_tokens
if temperature: # model-specific param
@ -6512,6 +6508,8 @@ async def realtime_websocket_endpoint(
),
user_api_key_dict=Depends(user_api_key_auth_websocket),
):
import websockets
await websocket.accept()
# Only use explicit parameters, not all query params
@ -8070,7 +8068,8 @@ async def _apply_search_filter_to_models(
# Fetch database models if we need more for the current page
if router_models_count < models_needed_for_page:
models_to_fetch = min(
models_needed_for_page - router_models_count, db_models_total_count
models_needed_for_page - router_models_count,
db_models_total_count,
)
if models_to_fetch > 0:
@ -8106,21 +8105,21 @@ async def _apply_search_filter_to_models(
def _normalize_datetime_for_sorting(dt: Any) -> Optional[datetime]:
"""
Normalize a datetime value to a timezone-aware UTC datetime for sorting.
This function handles:
- None values: returns None
- String values: parses ISO format strings and converts to UTC-aware datetime
- Datetime objects: converts naive datetimes to UTC-aware, and aware datetimes to UTC
Args:
dt: Datetime value (None, str, or datetime object)
Returns:
UTC-aware datetime object, or None if input is None or cannot be parsed
"""
if dt is None:
return None
if isinstance(dt, str):
try:
# Handle ISO format strings, including 'Z' suffix
@ -8134,14 +8133,14 @@ def _normalize_datetime_for_sorting(dt: Any) -> Optional[datetime]:
return parsed_dt
except (ValueError, AttributeError):
return None
if isinstance(dt, datetime):
# If naive, assume UTC and make it aware
if dt.tzinfo is None:
return dt.replace(tzinfo=timezone.utc)
# If aware, convert to UTC
return dt.astimezone(timezone.utc)
return None
@ -8161,46 +8160,60 @@ def _sort_models(
Returns:
Sorted list of models
"""
if not sort_by or sort_by not in ["model_name", "created_at", "updated_at", "costs", "status"]:
if not sort_by or sort_by not in [
"model_name",
"created_at",
"updated_at",
"costs",
"status",
]:
return all_models
reverse = sort_order.lower() == "desc"
def get_sort_key(model: Dict[str, Any]) -> Any:
model_info = model.get("model_info", {})
if sort_by == "model_name":
return model.get("model_name", "").lower()
elif sort_by == "created_at":
created_at = model_info.get("created_at")
normalized_dt = _normalize_datetime_for_sorting(created_at)
if normalized_dt is None:
# Put None values at the end for asc, at the start for desc
return (datetime.max.replace(tzinfo=timezone.utc) if not reverse else datetime.min.replace(tzinfo=timezone.utc))
return (
datetime.max.replace(tzinfo=timezone.utc)
if not reverse
else datetime.min.replace(tzinfo=timezone.utc)
)
return normalized_dt
elif sort_by == "updated_at":
updated_at = model_info.get("updated_at")
normalized_dt = _normalize_datetime_for_sorting(updated_at)
if normalized_dt is None:
return (datetime.max.replace(tzinfo=timezone.utc) if not reverse else datetime.min.replace(tzinfo=timezone.utc))
return (
datetime.max.replace(tzinfo=timezone.utc)
if not reverse
else datetime.min.replace(tzinfo=timezone.utc)
)
return normalized_dt
elif sort_by == "costs":
input_cost = model_info.get("input_cost_per_token", 0) or 0
output_cost = model_info.get("output_cost_per_token", 0) or 0
total_cost = input_cost + output_cost
# Put 0 or None costs at the end for asc, at the start for desc
if total_cost == 0:
return (float("inf") if not reverse else float("-inf"))
return float("inf") if not reverse else float("-inf")
return total_cost
elif sort_by == "status":
# False (config) comes before True (db) for asc
db_model = model_info.get("db_model", False)
return db_model
return None
try:
@ -8396,9 +8409,7 @@ async def _find_model_by_id(
)
if db_model:
# Convert database model to router format
decrypted_models = proxy_config.decrypt_model_list_from_db(
[db_model]
)
decrypted_models = proxy_config.decrypt_model_list_from_db([db_model])
if decrypted_models:
found_model = decrypted_models[0]
except Exception as e:
@ -10708,9 +10719,9 @@ async def get_config_list(
hasattr(sub_field_info, "description")
and sub_field_info.description is not None
):
nested_fields[
idx
].field_description = sub_field_info.description
nested_fields[idx].field_description = (
sub_field_info.description
)
idx += 1
_stored_in_db = None