From 6772d817af7e13e0676005cb6ccd7d7a7b37dee4 Mon Sep 17 00:00:00 2001 From: Harshit28j Date: Thu, 26 Feb 2026 21:37:19 +0530 Subject: [PATCH] feaat: add tags for projects --- litellm/proxy/_types.py | 72 +++++-- litellm/proxy/litellm_pre_call_utils.py | 24 ++- .../management_endpoints/project_endpoints.py | 14 ++ litellm/proxy/utils.py | 195 +++++++++++------- 4 files changed, 204 insertions(+), 101 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index de1609baf62..e86ded01f31 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1106,7 +1106,9 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): raise ValueError("args is required for stdio transport") elif transport in [MCPTransport.http, MCPTransport.sse]: if not values.get("url") and not values.get("spec_path"): - raise ValueError("url or spec_path is required for HTTP/SSE transport") + raise ValueError( + "url or spec_path is required for HTTP/SSE transport" + ) return values @model_validator(mode="before") @@ -1158,7 +1160,9 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): raise ValueError("args is required for stdio transport") elif transport in [MCPTransport.http, MCPTransport.sse]: if not values.get("url") and not values.get("spec_path"): - raise ValueError("url or spec_path is required for HTTP/SSE transport") + raise ValueError( + "url or spec_path is required for HTTP/SSE transport" + ) return values @@ -1409,12 +1413,12 @@ class NewCustomerRequest(BudgetNewRequest): blocked: bool = False # allow/disallow requests for this end-user budget_id: Optional[str] = None # give either a budget_id or max_budget spend: Optional[float] = None - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @model_validator(mode="before") @@ -1437,12 +1441,12 @@ class UpdateCustomerRequest(LiteLLMPydanticObjectBase): blocked: bool = False # allow/disallow requests for this end-user max_budget: Optional[float] = None budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[AllowedModelRegion] = ( - None # require all user requests to use models in this specific region - ) - default_model: Optional[str] = ( - None # if no equivalent model in allowed region - default all requests to this model - ) + allowed_model_region: Optional[ + AllowedModelRegion + ] = None # require all user requests to use models in this specific region + default_model: Optional[ + str + ] = None # if no equivalent model in allowed region - default all requests to this model object_permission: Optional[LiteLLM_ObjectPermissionBase] = None @@ -2274,6 +2278,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): organization_tpm_limit: Optional[int] = None organization_rpm_limit: Optional[int] = None organization_metadata: Optional[dict] = None + project_metadata: Optional[dict] = None # Time stamps last_refreshed_at: Optional[float] = None # last time joint view was pulled from db @@ -2581,16 +2586,20 @@ class NewProjectRequest(LiteLLM_BudgetTable): model_tpm_limit: Optional[dict] = None blocked: bool = False object_permission: Optional[LiteLLM_ObjectPermissionBase] = None + tags: Optional[list] = None @model_validator(mode="before") @classmethod def set_model_info(cls, values): - for field in LiteLLM_ManagementEndpoint_MetadataFields: + for field in ( + LiteLLM_ManagementEndpoint_MetadataFields + + LiteLLM_ManagementEndpoint_MetadataFields_Premium + ): if values.get(field) is not None: if values.get("metadata") is None: values.update({"metadata": {}}) values["metadata"][field] = values.get(field) - values.pop(field) + values.pop(field, None) return values @@ -2608,16 +2617,20 @@ class UpdateProjectRequest(LiteLLM_BudgetTable): blocked: Optional[bool] = None budget_id: Optional[str] = None object_permission: Optional[LiteLLM_ObjectPermissionBase] = None + tags: Optional[list] = None @model_validator(mode="before") @classmethod def set_model_info(cls, values): - for field in LiteLLM_ManagementEndpoint_MetadataFields: + for field in ( + LiteLLM_ManagementEndpoint_MetadataFields + + LiteLLM_ManagementEndpoint_MetadataFields_Premium + ): if values.get(field) is not None: if values.get("metadata") is None: values.update({"metadata": {}}) values["metadata"][field] = values.get(field) - values.pop(field) + values.pop(field, None) return values @@ -2636,6 +2649,7 @@ class LiteLLM_ProjectTable(LiteLLMPydanticObjectBase): team_id: Optional[str] = None budget_id: Optional[str] = None metadata: Optional[dict] = None + tags: Optional[list] = None models: List[str] = [] spend: float = 0.0 model_spend: Optional[dict] = None @@ -2645,6 +2659,17 @@ class LiteLLM_ProjectTable(LiteLLMPydanticObjectBase): object_permission_id: Optional[str] = None created_by: str updated_by: str + + @model_validator(mode="before") + @classmethod + def surface_tags(cls, values): + """Surfaces tags from metadata to the top level for responses""" + if isinstance(values, dict): + metadata = values.get("metadata") + if isinstance(metadata, dict) and "tags" in metadata: + values["tags"] = metadata["tags"] + return values + litellm_budget_table: Optional[LiteLLM_BudgetTable] = None object_permission: Optional[LiteLLM_ObjectPermissionTable] = None @@ -3056,7 +3081,9 @@ class SpendLogsMetadata(TypedDict): str ] # S3/GCS object key for cold storage retrieval litellm_overhead_time_ms: Optional[float] # LiteLLM overhead time in milliseconds - attempted_retries: Optional[int] # Number of retries attempted (0 = first attempt succeeded) + attempted_retries: Optional[ + int + ] # Number of retries attempted (0 = first attempt succeeded) max_retries: Optional[int] # Max retries configured for this request cost_breakdown: Optional[ CostBreakdown @@ -4117,10 +4144,10 @@ class SpendUpdateQueueItem(TypedDict, total=False): class ToolDiscoveryQueueItem(TypedDict, total=False): tool_name: str - origin: Optional[str] # MCP server name or "user_defined" + origin: Optional[str] # MCP server name or "user_defined" created_by: Optional[str] - key_hash: Optional[str] # hash of virtual key that triggered discovery - team_id: Optional[str] # team that triggered discovery + key_hash: Optional[str] # hash of virtual key that triggered discovery + team_id: Optional[str] # team that triggered discovery key_alias: Optional[str] # human-readable key alias @@ -4144,6 +4171,7 @@ class LiteLLM_ManagedObjectTable(LiteLLMPydanticObjectBase): class LiteLLM_ManagedVectorStoreTable(LiteLLMPydanticObjectBase): """Table for managing vector stores with target_model_names support.""" + unified_resource_id: str resource_object: Optional[Any] = None # VectorStoreCreateResponse model_mappings: Dict[str, str] diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index d2312a00c3b..e8a009921f7 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -248,13 +248,15 @@ def clean_headers( clean_headers = {} litellm_key_lower = ( litellm_key_header_name.lower() if litellm_key_header_name is not None else None - ) + ) for header, value in headers.items(): header_lower = header.lower() - + if header_lower == "authorization" and is_anthropic_oauth_key(value): clean_headers[header] = value - elif forward_llm_provider_auth_headers and header_lower in _SPECIAL_HEADERS_CACHE: + elif ( + forward_llm_provider_auth_headers and header_lower in _SPECIAL_HEADERS_CACHE + ): if litellm_key_lower and header_lower == litellm_key_lower: continue if header_lower == "authorization": @@ -840,11 +842,13 @@ async def add_litellm_data_to_request( # noqa: PLR0915 from litellm.types.proxy.litellm_pre_call_utils import SecretFields _raw_headers: Dict[str, str] = _safe_get_request_headers(request) - + forward_llm_auth = False if general_settings: - forward_llm_auth = general_settings.get("forward_llm_provider_auth_headers", False) - + forward_llm_auth = general_settings.get( + "forward_llm_provider_auth_headers", False + ) + _headers: Dict[str, str] = clean_headers( request.headers, litellm_key_header_name=( @@ -997,6 +1001,14 @@ async def add_litellm_data_to_request( # noqa: PLR0915 request_tags=data[_metadata_variable_name].get("tags"), tags_to_add=team_metadata["tags"], ) + + ## PROJECT-LEVEL SPEND LOGS/TAGS + project_metadata = user_api_key_dict.project_metadata or {} + if "tags" in project_metadata and project_metadata["tags"] is not None: + data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup._merge_tags( + request_tags=data[_metadata_variable_name].get("tags"), + tags_to_add=project_metadata["tags"], + ) if "disable_global_guardrails" in team_metadata and isinstance( team_metadata["disable_global_guardrails"], bool ): diff --git a/litellm/proxy/management_endpoints/project_endpoints.py b/litellm/proxy/management_endpoints/project_endpoints.py index ba3238ebfd5..8b1986f08e2 100644 --- a/litellm/proxy/management_endpoints/project_endpoints.py +++ b/litellm/proxy/management_endpoints/project_endpoints.py @@ -429,6 +429,14 @@ async def new_project( value=getattr(data, field), ) + for field in LiteLLM_ManagementEndpoint_MetadataFields_Premium: + if getattr(data, field, None) is not None: + _set_object_metadata_field( + object_data=project_row, + field_name=field, + value=getattr(data, field), + ) + new_project_row = prisma_client.jsonify_object( project_row.json(exclude_none=True) ) @@ -633,6 +641,12 @@ async def update_project( update_data["metadata"] = {} update_data["metadata"][field] = update_data.pop(field) + for field in LiteLLM_ManagementEndpoint_MetadataFields_Premium: + if field in update_data: + if update_data.get("metadata") is None: + update_data["metadata"] = {} + update_data["metadata"][field] = update_data.pop(field) + # Remove budget fields (following organization_endpoints.py pattern) update_data = _remove_budget_fields_from_project_data(update_data) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index f6613b5548f..c3587ab88f1 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -23,23 +23,31 @@ from typing import ( ) from litellm import _custom_logger_compatible_callbacks_literal -from litellm.constants import (DEFAULT_MODEL_CREATED_AT_TIME, - MAX_TEAM_LIST_LIMIT) -from litellm.proxy._types import (DB_CONNECTION_ERROR_TYPES, CommonProxyErrors, - ProxyErrorTypes, ProxyException, - SpendLogsMetadata, SpendLogsPayload) +from litellm.constants import DEFAULT_MODEL_CREATED_AT_TIME, MAX_TEAM_LIST_LIMIT +from litellm.proxy._types import ( + DB_CONNECTION_ERROR_TYPES, + CommonProxyErrors, + ProxyErrorTypes, + ProxyException, + SpendLogsMetadata, + SpendLogsPayload, +) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import CallTypes, CallTypesLiteral try: - from litellm_enterprise.enterprise_callbacks.send_emails.base_email import \ - BaseEmailLogger - from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import \ - ResendEmailLogger - from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import \ - SendGridEmailLogger - from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import \ - SMTPEmailLogger + from litellm_enterprise.enterprise_callbacks.send_emails.base_email import ( + BaseEmailLogger, + ) + from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ( + ResendEmailLogger, + ) + from litellm_enterprise.enterprise_callbacks.send_emails.sendgrid_email import ( + SendGridEmailLogger, + ) + from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import ( + SMTPEmailLogger, + ) except ImportError: BaseEmailLogger = None # type: ignore SendGridEmailLogger = None # type: ignore @@ -58,56 +66,70 @@ from fastapi import HTTPException, status import litellm import litellm.litellm_core_utils import litellm.litellm_core_utils.litellm_logging -from litellm import (EmbeddingResponse, ImageResponse, ModelResponse, - ModelResponseStream, Router) +from litellm import ( + EmbeddingResponse, + ImageResponse, + ModelResponse, + ModelResponseStream, + Router, +) from litellm._logging import verbose_proxy_logger from litellm._service_logger import ServiceLogging, ServiceTypes from litellm.caching.caching import DualCache, RedisCache from litellm.caching.dual_cache import LimitedSizeOrderedDict from litellm.exceptions import RejectedRequestError -from litellm.integrations.custom_guardrail import (CustomGuardrail, - ModifyResponseException) +from litellm.integrations.custom_guardrail import ( + CustomGuardrail, + ModifyResponseException, +) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting -from litellm.integrations.SlackAlerting.utils import \ - _add_langfuse_trace_id_to_alert +from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.proxy._types import (AlertType, CallInfo, - LiteLLM_VerificationTokenView, Member, - UserAPIKeyAuth) +from litellm.proxy._types import ( + AlertType, + CallInfo, + LiteLLM_VerificationTokenView, + Member, + UserAPIKeyAuth, +) from litellm.proxy.auth.route_checks import RouteChecks -from litellm.proxy.db.create_views import (create_missing_views, - should_create_missing_views) +from litellm.proxy.db.create_views import ( + create_missing_views, + should_create_missing_views, +) from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.db.log_db_metrics import log_db_metrics from litellm.proxy.db.prisma_client import PrismaWrapper -from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import \ - UnifiedLLMGuardrails +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, +) from litellm.proxy.hooks import PROXY_HOOKS, get_proxy_hook from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter -from litellm.proxy.hooks.parallel_request_limiter import \ - _PROXY_MaxParallelRequestsHandler +from litellm.proxy.hooks.parallel_request_limiter import ( + _PROXY_MaxParallelRequestsHandler, +) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor from litellm.secret_managers.main import str_to_bool from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES -from litellm.types.mcp import (MCPDuringCallResponseObject, - MCPPreCallRequestObject, - MCPPreCallResponseObject) -from litellm.types.proxy.policy_engine.pipeline_types import \ - PipelineExecutionResult +from litellm.types.mcp import ( + MCPDuringCallResponseObject, + MCPPreCallRequestObject, + MCPPreCallResponseObject, +) +from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionResult from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - from litellm.litellm_core_utils.litellm_logging import \ - Logging as LiteLLMLoggingObj + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj Span = Union[_Span, Any] else: @@ -1050,9 +1072,10 @@ class ProxyLogging: """Process prompt template if applicable.""" from litellm.proxy.prompts.prompt_endpoints import ( - construct_versioned_prompt_id, get_latest_version_prompt_id) - from litellm.proxy.prompts.prompt_registry import \ - IN_MEMORY_PROMPT_REGISTRY + construct_versioned_prompt_id, + get_latest_version_prompt_id, + ) + from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY from litellm.utils import get_non_default_completion_params if prompt_version is None: @@ -1102,8 +1125,9 @@ class ProxyLogging: def _process_guardrail_metadata(self, data: dict) -> None: """Process guardrails from metadata and add to applied_guardrails.""" - from litellm.proxy.common_utils.callback_utils import \ - add_guardrail_to_applied_guardrails_header + from litellm.proxy.common_utils.callback_utils import ( + add_guardrail_to_applied_guardrails_header, + ) metadata_standard = data.get("metadata") or {} metadata_litellm = data.get("litellm_metadata") or {} @@ -2000,8 +2024,7 @@ class ProxyLogging: if isinstance(response, (ModelResponse, ModelResponseStream)): response_str = litellm.get_response_string(response_obj=response) elif isinstance(response, dict) and self.is_a2a_streaming_response(response): - from litellm.llms.a2a.common_utils import \ - extract_text_from_a2a_response + from litellm.llms.a2a.common_utils import extract_text_from_a2a_response response_str = extract_text_from_a2a_response(response) if response_str is not None: @@ -2010,8 +2033,7 @@ class ProxyLogging: _callback: Optional[CustomLogger] = None if isinstance(callback, CustomGuardrail): # Main - V2 Guardrails implementation - from litellm.types.guardrails import \ - GuardrailEventHooks + from litellm.types.guardrails import GuardrailEventHooks ## CHECK FOR MODEL-LEVEL GUARDRAILS modified_data = _check_and_merge_model_level_guardrails( @@ -2851,9 +2873,11 @@ class PrismaClient: o.metadata as organization_metadata, b2.max_budget as organization_max_budget, b2.tpm_limit as organization_tpm_limit, - b2.rpm_limit as organization_rpm_limit + b2.rpm_limit as organization_rpm_limit, + p.metadata as project_metadata FROM "LiteLLM_VerificationToken" AS v LEFT JOIN "LiteLLM_TeamTable" AS t ON v.team_id = t.team_id + LEFT JOIN "LiteLLM_ProjectTable" AS p ON v.project_id = p.project_id LEFT JOIN "LiteLLM_TeamMembership" AS tm ON v.team_id = tm.team_id AND tm.user_id = v.user_id LEFT JOIN "LiteLLM_ModelTable" m ON t.model_id = m.id LEFT JOIN "LiteLLM_BudgetTable" AS b ON v.budget_id = b.budget_id @@ -3600,13 +3624,15 @@ class PrismaClient: probe_pid, _ = os.waitpid(pid, os.WNOHANG) except ChildProcessError: verbose_proxy_logger.debug( - "PID %s is not a child process; skipping waitpid watch.", pid, + "PID %s is not a child process; skipping waitpid watch.", + pid, ) return False if probe_pid == pid: verbose_proxy_logger.warning( - "prisma-query-engine PID %s already dead at watch start.", pid, + "prisma-query-engine PID %s already dead at watch start.", + pid, ) self._engine_confirmed_dead = True self._reap_all_zombies() @@ -3783,11 +3809,17 @@ class PrismaClient: waitpid thread nor pidfd are available. """ - if self._watching_engine or self._engine_pidfd >= 0 or self._engine_wait_thread is not None: + if ( + self._watching_engine + or self._engine_pidfd >= 0 + or self._engine_wait_thread is not None + ): return pid = self._get_engine_pid() if pid == 0: - verbose_proxy_logger.debug("Could not find prisma-query-engine PID; engine death detection unavailable.") + verbose_proxy_logger.debug( + "Could not find prisma-query-engine PID; engine death detection unavailable." + ) return self._engine_pid = pid self._engine_confirmed_dead = False @@ -3796,15 +3828,18 @@ class PrismaClient: pidfd_ok = False if waitpid_ok else self._try_pidfd_watch(pid) if waitpid_ok: verbose_proxy_logger.info( - "Watching engine PID %s via waitpid thread.", pid, + "Watching engine PID %s via waitpid thread.", + pid, ) elif pidfd_ok: verbose_proxy_logger.info( - "Watching engine PID %s via pidfd.", pid, + "Watching engine PID %s via pidfd.", + pid, ) else: verbose_proxy_logger.info( - "Watching engine PID %s via os.kill polling.", pid, + "Watching engine PID %s via os.kill polling.", + pid, ) self._watching_engine = True asyncio.create_task(self._poll_engine_proc()) @@ -3827,7 +3862,9 @@ class PrismaClient: blip -- disconnect, connect, SELECT 1). """ effective_timeout = ( - timeout_seconds if timeout_seconds is not None else self._db_watchdog_reconnect_timeout_seconds + timeout_seconds + if timeout_seconds is not None + else self._db_watchdog_reconnect_timeout_seconds ) engine_is_dead = self._engine_confirmed_dead or ( @@ -3847,14 +3884,18 @@ class PrismaClient: async def _do_heavy_reconnect() -> None: db_url = os.getenv("DATABASE_URL", "") if not db_url: - verbose_proxy_logger.error("DATABASE_URL not set; cannot recreate Prisma client.") + verbose_proxy_logger.error( + "DATABASE_URL not set; cannot recreate Prisma client." + ) raise RuntimeError("DATABASE_URL not set") await self.db.recreate_prisma_client(db_url) await self._start_engine_watcher() await asyncio.wait_for(_do_heavy_reconnect(), timeout=effective_timeout) else: - verbose_proxy_logger.debug("Performing Prisma DB reconnect (engine alive or unknown).") + verbose_proxy_logger.debug( + "Performing Prisma DB reconnect (engine alive or unknown)." + ) async def _do_direct_reconnect() -> None: try: @@ -3937,7 +3978,9 @@ class PrismaClient: if lock_timeout_seconds is None: async with self._db_reconnect_lock: - return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) + return await self._attempt_reconnect_inside_lock( + force, reason, timeout_seconds + ) lock_acquired_by_timeout_task = False @@ -3986,14 +4029,17 @@ class PrismaClient: return False try: - return await self._attempt_reconnect_inside_lock(force, reason, timeout_seconds) + return await self._attempt_reconnect_inside_lock( + force, reason, timeout_seconds + ) finally: self._db_reconnect_lock.release() async def start_db_health_watchdog_task(self) -> None: """Start background tasks that monitor DB health: - A periodic SELECT 1 probe that triggers reconnect on network/connection failure. - - A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling.""" + - A process-level watcher that detects engine death via waitpid thread, pidfd, or os.kill polling. + """ if self._db_health_watchdog_enabled is not True: verbose_proxy_logger.debug( "Prisma DB health watchdog disabled via PRISMA_HEALTH_WATCHDOG_ENABLED" @@ -4453,9 +4499,9 @@ class ProxyUpdateSpend: :MAX_LOGS_PER_INTERVAL ] # Remove the logs we're about to process - prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[ - len(logs_to_process) : - ] + prisma_client.spend_log_transactions = ( + prisma_client.spend_log_transactions[len(logs_to_process) :] + ) popped_batch = True if len(logs_to_process) > 0: verbose_proxy_logger.info( @@ -4609,9 +4655,7 @@ async def update_spend_logs_job( return async with prisma_client._spend_log_transactions_lock: - logs_to_process = prisma_client.spend_log_transactions[ - :MAX_LOGS_PER_INTERVAL - ] + logs_to_process = prisma_client.spend_log_transactions[:MAX_LOGS_PER_INTERVAL] prisma_client.spend_log_transactions = prisma_client.spend_log_transactions[ len(logs_to_process) : ] @@ -4626,8 +4670,10 @@ async def update_spend_logs_job( # Guardrail/policy usage tracking (same batch, outside spend-logs update) try: - from litellm.proxy.guardrails.usage_tracking import \ - process_spend_logs_guardrail_usage + from litellm.proxy.guardrails.usage_tracking import ( + process_spend_logs_guardrail_usage, + ) + await process_spend_logs_guardrail_usage( prisma_client=prisma_client, logs_to_process=logs_to_process, @@ -4653,8 +4699,10 @@ async def _monitor_spend_logs_queue( db_writer_client: Optional HTTP handler for external spend logs endpoint proxy_logging_obj: Proxy logging object """ - from litellm.constants import (SPEND_LOG_QUEUE_POLL_INTERVAL, - SPEND_LOG_QUEUE_SIZE_THRESHOLD) + from litellm.constants import ( + SPEND_LOG_QUEUE_POLL_INTERVAL, + SPEND_LOG_QUEUE_SIZE_THRESHOLD, + ) threshold = SPEND_LOG_QUEUE_SIZE_THRESHOLD base_interval = SPEND_LOG_QUEUE_POLL_INTERVAL @@ -5175,11 +5223,12 @@ async def get_available_models_for_user( List of model names available to the user """ from litellm.proxy.auth.auth_checks import get_team_object - from litellm.proxy.auth.model_checks import (get_complete_model_list, - get_key_models, - get_team_models) - from litellm.proxy.management_endpoints.team_endpoints import \ - validate_membership + from litellm.proxy.auth.model_checks import ( + get_complete_model_list, + get_key_models, + get_team_models, + ) + from litellm.proxy.management_endpoints.team_endpoints import validate_membership # Get proxy model list and access groups if llm_router is None: