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

Merge 527bb6f99d into ec94a1f82a

Browse source
This commit is contained in:
Lin Junrong 2026-08-27 19:08:04 -05:00 • committed by GitHub
commit 9508ab55b6
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 71 additions and 40 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

111
litellm/proxy/auth/auth_checks.py
View file

@ -279,6 +279,17 @@ db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s
all_routes: Final = LiteLLMRoutes.openai_routes.value + LiteLLMRoutes.management_routes.value
# The event loop only keeps weak references to tasks, so a background task whose
# only reference was the create_task() call can be collected before it runs. Hold
# each one until it completes.
_background_tasks: Final[set[asyncio.Task]] = set() # mutable-ok: task registry
def _hold_background_task(task: asyncio.Task) -> None:
"""Keep a strong reference to a fire-and-forget task until it finishes."""
_background_tasks.add(task)
task.add_done_callback(_background_tasks.discard)
def _log_budget_lookup_failure(entity: str, error: Exception) -> None:
"""
@ -2148,10 +2159,12 @@ async def _get_fuzzy_user_object(
)
if response is not None and sso_user_id is not None: # update sso_user_id
asyncio.create_task( # background task to update user with sso id
_user_table(UserRepository(prisma_client)).update(
where={"user_id": response.user_id},
data={"sso_user_id": sso_user_id},
_hold_background_task(
asyncio.create_task( # background task to update user with sso id
_user_table(UserRepository(prisma_client)).update(
where={"user_id": response.user_id},
data={"sso_user_id": sso_user_id},
)
)
)
@ -4411,10 +4424,12 @@ async def _virtual_key_max_budget_check(
key_alias=valid_token.key_alias,
event_group=Litellm_EntityType.KEY,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="token_budget",
user_info=call_info,
_hold_background_task(
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="token_budget",
user_info=call_info,
)
)
)
@ -4515,10 +4530,12 @@ async def _virtual_key_soft_budget_check(
event_group=Litellm_EntityType.KEY,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="soft_budget",
user_info=call_info,
_hold_background_task(
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="soft_budget",
user_info=call_info,
)
)
)
@ -4606,10 +4623,12 @@ async def _virtual_key_max_budget_alert_check(
event_group=Litellm_EntityType.KEY,
max_budget_alert_emails=alert_email_config,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="max_budget_alert",
user_info=call_info,
_hold_background_task(
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="max_budget_alert",
user_info=call_info,
)
)
)
else:
@ -4638,10 +4657,12 @@ async def _virtual_key_max_budget_alert_check(
event_group=Litellm_EntityType.KEY,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="max_budget_alert",
user_info=call_info,
_hold_background_task(
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="max_budget_alert",
user_info=call_info,
)
)
)
@ -4802,10 +4823,12 @@ async def _team_max_budget_check(
organization_id=valid_token.org_id,
event_group=Litellm_EntityType.TEAM,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="team_budget",
user_info=call_info,
_hold_background_task(
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="team_budget",
user_info=call_info,
)
)
)
@ -4920,10 +4943,12 @@ async def _team_soft_budget_check(
alert_emails=alert_emails,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="soft_budget",
user_info=call_info,
_hold_background_task(
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="soft_budget",
user_info=call_info,
)
)
)
@ -4964,10 +4989,12 @@ async def _project_max_budget_check(
organization_id=valid_token.org_id,
event_group=Litellm_EntityType.PROJECT,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="project_budget",
user_info=call_info,
_hold_background_task(
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="project_budget",
user_info=call_info,
)
)
)
@ -5016,10 +5043,12 @@ async def _project_soft_budget_check(
organization_id=valid_token.org_id,
event_group=Litellm_EntityType.PROJECT,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="soft_budget",
user_info=call_info,
_hold_background_task(
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="soft_budget",
user_info=call_info,
)
)
)
@ -5174,10 +5203,12 @@ async def _organization_max_budget_check(
organization_id=org_id,
event_group=Litellm_EntityType.ORGANIZATION,
)
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="organization_budget",
user_info=call_info,
_hold_background_task(
asyncio.create_task(
proxy_logging_obj.budget_alerts(
type="organization_budget",
user_info=call_info,
)
)
)