diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 9b07d44deb6..e44c5bea651 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,59 +1,40 @@ import enum import json from datetime import datetime -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Literal, Optional, Union +from typing import (TYPE_CHECKING, Any, Callable, Dict, List, Literal, + Optional, Union) import httpx -from pydantic import ( - BaseModel, - ConfigDict, - Field, - Json, - field_validator, - model_validator, -) +from pydantic import (BaseModel, ConfigDict, Field, Json, field_validator, + model_validator) from typing_extensions import Required, TypedDict from litellm._uuid import uuid from litellm.types.integrations.slack_alerting import AlertType -from litellm.types.llms.openai import ( - AllMessageValues, - OpenAIFileObject, - ResponsesAPIResponse, -) -from litellm.types.mcp import ( - MCPAuthType, - MCPCredentials, - MCPTransport, - MCPTransportType, -) +from litellm.types.llms.openai import (AllMessageValues, OpenAIFileObject, + ResponsesAPIResponse) +from litellm.types.mcp import (MCPAuthType, MCPCredentials, MCPTransport, + MCPTransportType) from litellm.types.mcp_server.mcp_server_manager import MCPInfo from litellm.types.router import RouterErrors, UpdateRouterConfig from litellm.types.secret_managers.main import KeyManagementSystem -from litellm.types.utils import ( - CallTypes, - CostBreakdown, - EmbeddingResponse, - GenericBudgetConfigType, - ImageResponse, - LiteLLMBatch, - LiteLLMFineTuningJob, - LiteLLMPydanticObjectBase, - ModelResponse, - ProviderField, - StandardCallbackDynamicParams, - StandardLoggingGuardrailInformation, - StandardLoggingMCPToolCall, - StandardLoggingModelInformation, - StandardLoggingPayloadErrorInformation, - StandardLoggingPayloadStatus, - StandardLoggingVectorStoreRequest, - StandardPassThroughResponseObject, - TextCompletionResponse, -) +from litellm.types.utils import (CallTypes, CostBreakdown, EmbeddingResponse, + GenericBudgetConfigType, ImageResponse, + LiteLLMBatch, LiteLLMFineTuningJob, + LiteLLMPydanticObjectBase, ModelResponse, + ProviderField, StandardCallbackDynamicParams, + StandardLoggingGuardrailInformation, + StandardLoggingMCPToolCall, + StandardLoggingModelInformation, + StandardLoggingPayloadErrorInformation, + StandardLoggingPayloadStatus, + StandardLoggingVectorStoreRequest, + StandardPassThroughResponseObject, + TextCompletionResponse) from litellm.types.videos.main import VideoObject -from .types_utils.utils import get_instance_fn, validate_custom_validate_return_type +from .types_utils.utils import (get_instance_fn, + validate_custom_validate_return_type) if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -201,6 +182,7 @@ class Litellm_EntityType(enum.Enum): ORGANIZATION = "organization" PROJECT = "project" TAG = "tag" + AGENT = "agent" # global proxy level entity PROXY = "proxy" @@ -2399,7 +2381,8 @@ class UserAPIKeyAuth( This is used to track number of requests/spend for health check calls. """ - from litellm.constants import LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME + from litellm.constants import \ + LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME return cls( api_key=LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, @@ -2431,7 +2414,8 @@ class UserAPIKeyAuth( This is used to track actions performed by automated system jobs. """ - from litellm.constants import LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME + from litellm.constants import \ + LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME return cls( api_key=LITELLM_INTERNAL_JOBS_SERVICE_ACCOUNT_NAME, @@ -2837,7 +2821,8 @@ class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase): @model_validator(mode="after") def mask_api_keys(self): - from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker + from litellm.litellm_core_utils.sensitive_data_masker import \ + SensitiveDataMasker masker = SensitiveDataMasker(sensitive_patterns={"key"}) @@ -4205,6 +4190,7 @@ class DBSpendUpdateTransactions(TypedDict): team_member_list_transactions: Optional[Dict[str, float]] org_list_transactions: Optional[Dict[str, float]] tag_list_transactions: Optional[Dict[str, float]] + agent_list_transactions: Optional[Dict[str, float]] class SpendUpdateQueueItem(TypedDict, total=False): diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index af51b9fe445..6de9e6b1993 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -11,7 +11,8 @@ Run checks for: import asyncio import re import time -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, cast +from typing import (TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union, + cast) from fastapi import HTTPException, Request, status from pydantic import BaseModel @@ -20,42 +21,29 @@ import litellm from litellm._logging import verbose_proxy_logger from litellm.caching.caching import DualCache from litellm.caching.dual_cache import LimitedSizeOrderedDict -from litellm.constants import ( - CLI_JWT_EXPIRATION_HOURS, - CLI_JWT_TOKEN_NAME, - DEFAULT_ACCESS_GROUP_CACHE_TTL, - DEFAULT_IN_MEMORY_TTL, - DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, - DEFAULT_MAX_RECURSE_DEPTH, - EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE, -) +from litellm.constants import (CLI_JWT_EXPIRATION_HOURS, CLI_JWT_TOKEN_NAME, + DEFAULT_ACCESS_GROUP_CACHE_TTL, + DEFAULT_IN_MEMORY_TTL, + DEFAULT_MANAGEMENT_OBJECT_IN_MEMORY_CACHE_TTL, + DEFAULT_MAX_RECURSE_DEPTH, + EMAIL_BUDGET_ALERT_MAX_SPEND_ALERT_PERCENTAGE) from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider -from litellm.proxy._types import ( - RBAC_ROLES, - CallInfo, - LiteLLM_AccessGroupTable, - LiteLLM_BudgetTable, - LiteLLM_EndUserTable, - Litellm_EntityType, - LiteLLM_JWTAuth, - LiteLLM_ObjectPermissionTable, - LiteLLM_OrganizationMembershipTable, - LiteLLM_OrganizationTable, - LiteLLM_ProjectTableCachedObj, - LiteLLM_TagTable, - LiteLLM_TeamMembership, - LiteLLM_TeamTable, - LiteLLM_TeamTableCachedObj, - LiteLLM_UserTable, - LiteLLMRoutes, - LitellmUserRoles, - NewTeamRequest, - ProxyErrorTypes, - ProxyException, - RoleBasedPermissions, - SpecialModelNames, - UserAPIKeyAuth, -) +from litellm.proxy._types import (RBAC_ROLES, CallInfo, + LiteLLM_AccessGroupTable, + LiteLLM_BudgetTable, LiteLLM_EndUserTable, + Litellm_EntityType, LiteLLM_JWTAuth, + LiteLLM_ObjectPermissionTable, + LiteLLM_OrganizationMembershipTable, + LiteLLM_OrganizationTable, + LiteLLM_ProjectTableCachedObj, + LiteLLM_TagTable, LiteLLM_TeamMembership, + LiteLLM_TeamTable, + LiteLLM_TeamTableCachedObj, + LiteLLM_UserTable, LiteLLMRoutes, + LitellmUserRoles, NewTeamRequest, + ProxyErrorTypes, ProxyException, + RoleBasedPermissions, SpecialModelNames, + UserAPIKeyAuth) from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.guardrails.tool_name_extraction import ( @@ -326,6 +314,29 @@ async def common_checks( # noqa: PLR0915 code=status.HTTP_401_UNAUTHORIZED, ) + # Require trace id for agent keys when agent has require_trace_id_on_calls_by_agent + if valid_token is not None and valid_token.agent_id: + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry + from litellm.proxy.litellm_pre_call_utils import \ + get_chain_id_from_headers + + agent = global_agent_registry.get_agent_by_id(agent_id=valid_token.agent_id) + if agent is not None: + require_trace_id = (agent.litellm_params or {}).get( + "require_trace_id_on_calls_by_agent" + ) + if require_trace_id: + headers_dict = dict(request.headers) + trace_id = get_chain_id_from_headers(headers_dict) + if not trace_id: + raise ProxyException( + message="Requests made with this agent's key must include the x-litellm-trace-id header.", + type=ProxyErrorTypes.bad_request_error, + param=None, + code=status.HTTP_400_BAD_REQUEST, + ) + ## 2.1 If user can call model (if personal key) if _model and team_object is None and user_object is not None: await can_user_call_model( @@ -480,7 +491,8 @@ async def common_checks( # noqa: PLR0915 _request_metadata: dict = request_body.get("metadata", {}) or {} if _request_metadata.get("guardrails"): # check if team allowed to modify guardrails - from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails + from litellm.proxy.guardrails.guardrail_helpers import \ + can_modify_guardrails can_modify: bool = can_modify_guardrails(team_object) if can_modify is False: @@ -1930,9 +1942,8 @@ class ExperimentalUIJWTToken: def get_experimental_ui_login_jwt_auth_token(user_info: LiteLLM_UserTable) -> str: from datetime import timedelta - from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - encrypt_value_helper, - ) + from litellm.proxy.common_utils.encrypt_decrypt_utils import \ + encrypt_value_helper if user_info.user_role is None: raise Exception("User role is required for experimental UI login") @@ -1978,9 +1989,8 @@ class ExperimentalUIJWTToken: """ from datetime import timedelta - from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - encrypt_value_helper, - ) + from litellm.proxy.common_utils.encrypt_decrypt_utils import \ + encrypt_value_helper if user_info.user_role is None: raise Exception("User role is required for CLI JWT login") @@ -2019,9 +2029,8 @@ class ExperimentalUIJWTToken: import json from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth - from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - decrypt_value_helper, - ) + from litellm.proxy.common_utils.encrypt_decrypt_utils import \ + decrypt_value_helper decrypted_token = decrypt_value_helper( hashed_token, key="ui_hash_key", exception_type="debug" @@ -2142,13 +2151,11 @@ async def get_key_object( ) # else, check db - _valid_token: Optional[BaseModel] = ( - await _fetch_key_object_from_db_with_reconnect( - hashed_token=hashed_token, - prisma_client=prisma_client, - parent_otel_span=parent_otel_span, - proxy_logging_obj=proxy_logging_obj, - ) + _valid_token: Optional[BaseModel] = await _fetch_key_object_from_db_with_reconnect( + hashed_token=hashed_token, + prisma_client=prisma_client, + parent_otel_span=parent_otel_span, + proxy_logging_obj=proxy_logging_obj, ) if _valid_token is None: @@ -2294,9 +2301,9 @@ async def get_org_object( # Cache the result await user_api_key_cache.async_set_cache( key=cache_key, - value=response.model_dump() - if hasattr(response, "model_dump") - else response, + value=( + response.model_dump() if hasattr(response, "model_dump") else response + ), ttl=DEFAULT_IN_MEMORY_TTL, ) @@ -2339,8 +2346,10 @@ async def _get_resources_from_access_groups( # Lazy import to avoid circular imports if prisma_client is None or user_api_key_cache is None: from litellm.proxy.proxy_server import prisma_client as _prisma_client - from litellm.proxy.proxy_server import proxy_logging_obj as _proxy_logging_obj - from litellm.proxy.proxy_server import user_api_key_cache as _user_api_key_cache + from litellm.proxy.proxy_server import \ + proxy_logging_obj as _proxy_logging_obj + from litellm.proxy.proxy_server import \ + user_api_key_cache as _user_api_key_cache prisma_client = prisma_client or _prisma_client user_api_key_cache = user_api_key_cache or _user_api_key_cache @@ -3296,7 +3305,8 @@ async def _tag_max_budget_check( BudgetExceededError if any tag is over its max budget. Triggers a budget alert if any tag is over its max budget. """ - from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body + from litellm.proxy.common_utils.http_parsing_utils import \ + get_tags_from_request_body if prisma_client is None: return diff --git a/litellm/proxy/db/db_spend_update_writer.py b/litellm/proxy/db/db_spend_update_writer.py index 4c96e079c9e..679ee200f7b 100644 --- a/litellm/proxy/db/db_spend_update_writer.py +++ b/litellm/proxy/db/db_spend_update_writer.py @@ -374,6 +374,18 @@ class DBSpendUpdateWriter: traceback.format_exc(), ) + try: + await self._update_agent_db( + response_cost=response_cost, + agent_id=payload_copy.get("agent_id"), + prisma_client=prisma_client, + ) + except Exception: + verbose_proxy_logger.debug( + "_batch_database_updates: _update_agent_db failed: %s", + traceback.format_exc(), + ) + try: await self.add_spend_log_transaction_to_daily_user_transaction( payload=payload_copy, @@ -604,6 +616,37 @@ class DBSpendUpdateWriter: ) raise e + async def _update_agent_db( + self, + response_cost: Optional[float], + agent_id: Optional[str], + prisma_client: Optional[PrismaClient], + ): + try: + if agent_id is None or prisma_client is None: + verbose_proxy_logger.debug( + "track_cost_callback: agent_id is None or prisma_client is None. Not tracking spend for agent" + ) + return + + await self.spend_update_queue.add_update( + update=SpendUpdateQueueItem( + entity_type=Litellm_EntityType.AGENT, + entity_id=agent_id, + response_cost=response_cost, + ) + ) + except Exception as e: + verbose_proxy_logger.error( + "Spend tracking - failed to enqueue agent spend update. " + "agent_id=%s, response_cost=%s - %s\n%s", + agent_id, + response_cost, + str(e), + traceback.format_exc(), + ) + raise e + async def _update_tag_db( self, response_cost: Optional[float], @@ -765,7 +808,7 @@ class DBSpendUpdateWriter: if db_spend_update_transactions is not None: verbose_proxy_logger.info( "Spend tracking - committing spend updates from Redis to DB: " - "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d", + "keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, tags=%d, agents=%d", len( db_spend_update_transactions.get("key_list_transactions") or {} @@ -798,6 +841,12 @@ class DBSpendUpdateWriter: db_spend_update_transactions.get("tag_list_transactions") or {} ), + len( + db_spend_update_transactions.get( + "agent_list_transactions" + ) + or {} + ), ) await self._commit_spend_updates_to_db( prisma_client=prisma_client, @@ -1279,6 +1328,18 @@ class DBSpendUpdateWriter: proxy_logging_obj=proxy_logging_obj, ) + ### UPDATE AGENT TABLE ### + agent_list_transactions = db_spend_update_transactions["agent_list_transactions"] + await DBSpendUpdateWriter._update_entity_spend_in_db( + entity_name="Agent", + transactions=agent_list_transactions, + table_accessor="litellm_agentstable", + where_field="agent_id", + n_retry_times=n_retry_times, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + @staticmethod async def _update_entity_spend_in_db( entity_name: str, diff --git a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py index 51201f96d77..4f38e71bbfa 100644 --- a/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py +++ b/litellm/proxy/db/db_transaction_queue/redis_update_buffer.py @@ -10,33 +10,31 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast from litellm._logging import verbose_proxy_logger from litellm.caching import RedisCache -from litellm.constants import ( - MAX_REDIS_BUFFER_DEQUEUE_COUNT, - REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, - REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, - REDIS_UPDATE_BUFFER_KEY, -) +from litellm.constants import (MAX_REDIS_BUFFER_DEQUEUE_COUNT, + REDIS_DAILY_AGENT_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_END_USER_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_ORG_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_TAG_SPEND_UPDATE_BUFFER_KEY, + REDIS_DAILY_TEAM_SPEND_UPDATE_BUFFER_KEY, + REDIS_UPDATE_BUFFER_KEY) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy._types import ( - DailyTagSpendTransaction, - DailyTeamSpendTransaction, - DailyUserSpendTransaction, - DailyOrganizationSpendTransaction, - DailyEndUserSpendTransaction, - DBSpendUpdateTransactions, - DailyAgentSpendTransaction, -) -from litellm.proxy.db.db_transaction_queue.base_update_queue import service_logger_obj -from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import ( - DailySpendUpdateQueue, -) -from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue +from litellm.proxy._types import (DailyAgentSpendTransaction, + DailyEndUserSpendTransaction, + DailyOrganizationSpendTransaction, + DailyTagSpendTransaction, + DailyTeamSpendTransaction, + DailyUserSpendTransaction, + DBSpendUpdateTransactions) +from litellm.proxy.db.db_transaction_queue.base_update_queue import \ + service_logger_obj +from litellm.proxy.db.db_transaction_queue.daily_spend_update_queue import \ + DailySpendUpdateQueue +from litellm.proxy.db.db_transaction_queue.spend_update_queue import \ + SpendUpdateQueue from litellm.secret_managers.main import str_to_bool -from litellm.types.caching import RedisPipelineLpopOperation, RedisPipelineRpushOperation +from litellm.types.caching import (RedisPipelineLpopOperation, + RedisPipelineRpushOperation) from litellm.types.services import ServiceTypes if TYPE_CHECKING: @@ -579,6 +577,7 @@ class RedisUpdateBuffer: team_member_list_transactions={}, org_list_transactions={}, tag_list_transactions={}, + agent_list_transactions={}, ) # Define the transaction fields to process @@ -590,6 +589,7 @@ class RedisUpdateBuffer: "team_member_list_transactions", "org_list_transactions", "tag_list_transactions", + "agent_list_transactions", ] # Loop through each transaction and combine the values diff --git a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py index 3e059cf8c1f..b7cd06a64f3 100644 --- a/litellm/proxy/db/db_transaction_queue/spend_update_queue.py +++ b/litellm/proxy/db/db_transaction_queue/spend_update_queue.py @@ -3,15 +3,10 @@ from typing import Dict, List, Optional from litellm._logging import verbose_proxy_logger from litellm.constants import LITELLM_ASYNCIO_QUEUE_MAXSIZE -from litellm.proxy._types import ( - DBSpendUpdateTransactions, - Litellm_EntityType, - SpendUpdateQueueItem, -) +from litellm.proxy._types import (DBSpendUpdateTransactions, + Litellm_EntityType, SpendUpdateQueueItem) from litellm.proxy.db.db_transaction_queue.base_update_queue import ( - BaseUpdateQueue, - service_logger_obj, -) + BaseUpdateQueue, service_logger_obj) from litellm.types.services import ServiceTypes @@ -145,6 +140,7 @@ class SpendUpdateQueue(BaseUpdateQueue): team_member_list_transactions={}, org_list_transactions={}, tag_list_transactions={}, + agent_list_transactions={}, ) # Map entity types to their corresponding transaction dictionary keys @@ -156,6 +152,7 @@ class SpendUpdateQueue(BaseUpdateQueue): Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions", Litellm_EntityType.ORGANIZATION: "org_list_transactions", Litellm_EntityType.TAG: "tag_list_transactions", + Litellm_EntityType.AGENT: "agent_list_transactions", } for update in updates: @@ -207,6 +204,10 @@ class SpendUpdateQueue(BaseUpdateQueue): transactions_dict = db_spend_update_transactions[ "tag_list_transactions" ] + elif dict_key == "agent_list_transactions": + transactions_dict = db_spend_update_transactions[ + "agent_list_transactions" + ] else: continue diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 32eab99fb99..9e06a987721 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -10,16 +10,12 @@ import litellm from litellm._logging import verbose_logger, verbose_proxy_logger from litellm._service_logger import ServiceLogging from litellm.litellm_core_utils.safe_json_loads import safe_json_loads -from litellm.proxy._types import ( - AddTeamCallback, - CommonProxyErrors, - LitellmDataForBackendLLMCall, - LitellmUserRoles, - SpecialHeaders, - TeamCallbackMetadata, - UserAPIKeyAuth, -) -from litellm.proxy.common_utils.http_parsing_utils import _safe_get_request_headers +from litellm.proxy._types import (AddTeamCallback, CommonProxyErrors, + LitellmDataForBackendLLMCall, + LitellmUserRoles, SpecialHeaders, + TeamCallbackMetadata, UserAPIKeyAuth) +from litellm.proxy.common_utils.http_parsing_utils import \ + _safe_get_request_headers # Cache special headers as a frozenset for O(1) lookup performance _SPECIAL_HEADERS_CACHE = frozenset( @@ -28,12 +24,9 @@ _SPECIAL_HEADERS_CACHE = frozenset( from litellm.router import Router from litellm.types.llms.anthropic import ANTHROPIC_API_HEADERS from litellm.types.services import ServiceTypes -from litellm.types.utils import ( - LlmProviders, - ProviderSpecificHeader, - StandardLoggingUserAPIKeyMetadata, - SupportedCacheControls, -) +from litellm.types.utils import (LlmProviders, ProviderSpecificHeader, + StandardLoggingUserAPIKeyMetadata, + SupportedCacheControls) service_logger_obj = ServiceLogging() # used for tracking latency on OTEL @@ -667,6 +660,12 @@ class LiteLLMProxyRequestSetup: "user_api_key" ] = user_api_key_dict.api_key # this is just the hashed token + # Key-owned agent_id for spend attribution; keep existing (e.g. from header) if key has none + data[_metadata_variable_name]["agent_id"] = ( + getattr(user_api_key_dict, "agent_id", None) + or data[_metadata_variable_name].get("agent_id") + ) + data[_metadata_variable_name]["user_api_end_user_max_budget"] = getattr( user_api_key_dict, "end_user_max_budget", None ) @@ -689,8 +688,7 @@ class LiteLLMProxyRequestSetup: return data from litellm.proxy._types import ( LiteLLM_ManagementEndpoint_MetadataFields, - LiteLLM_ManagementEndpoint_MetadataFields_Premium, - ) + LiteLLM_ManagementEndpoint_MetadataFields_Premium) # ignore any special fields added_metadata = {} @@ -859,7 +857,8 @@ async def add_litellm_data_to_request( # noqa: PLR0915 """ from litellm.proxy.proxy_server import llm_router, premium_user - from litellm.types.proxy.litellm_pre_call_utils import RedactedDict, SecretFields + from litellm.types.proxy.litellm_pre_call_utils import (RedactedDict, + SecretFields) _raw_headers: Dict[str, str] = RedactedDict(_safe_get_request_headers(request)) @@ -1543,7 +1542,8 @@ async def move_guardrails_to_metadata( # Only check policy engine if no local config (avoid import + registry lookup) if not (has_key_config or has_team_config or has_request_config): - from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.proxy.policy_engine.policy_registry import \ + get_policy_registry if not get_policy_registry().is_initialized(): # Nothing configured anywhere - clean up request body fields and return @@ -1607,14 +1607,16 @@ async def move_guardrails_to_metadata( def _is_policy_version_id(s: str) -> bool: """Return True if string is a policy version ID (starts with policy_ prefix).""" - from litellm.proxy.policy_engine.policy_registry import POLICY_VERSION_ID_PREFIX + from litellm.proxy.policy_engine.policy_registry import \ + POLICY_VERSION_ID_PREFIX return isinstance(s, str) and s.startswith(POLICY_VERSION_ID_PREFIX) def _extract_policy_id(s: str) -> Optional[str]: """Extract raw UUID from policy_ string, or None if not a valid version ID.""" - from litellm.proxy.policy_engine.policy_registry import POLICY_VERSION_ID_PREFIX + from litellm.proxy.policy_engine.policy_registry import \ + POLICY_VERSION_ID_PREFIX if not _is_policy_version_id(s): return None @@ -1635,10 +1637,9 @@ def _match_and_track_policies( """ from litellm._logging import verbose_proxy_logger from litellm.proxy.common_utils.callback_utils import ( - add_policy_sources_to_metadata, - add_policy_to_applied_policies_header, - ) - from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + add_policy_sources_to_metadata, add_policy_to_applied_policies_header) + from litellm.proxy.policy_engine.attachment_registry import \ + get_attachment_registry from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher # Get matching policies via attachments (with match reasons for attribution) @@ -1783,7 +1784,8 @@ async def add_guardrails_from_policy_engine( user_api_key_dict: The user's API key authentication info """ from litellm._logging import verbose_proxy_logger - from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body + from litellm.proxy.common_utils.http_parsing_utils import \ + get_tags_from_request_body from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.types.proxy.policy_engine import PolicyMatchContext diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index ad31ff33802..99e8fa8e94c 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -14,22 +14,9 @@ import time import traceback import warnings from datetime import datetime, timedelta, timezone -from typing import ( - TYPE_CHECKING, - Any, - AsyncGenerator, - Dict, - List, - Literal, - Optional, - Set, - Tuple, - Union, - cast, - get_args, - get_origin, - get_type_hints, -) +from typing import (TYPE_CHECKING, Any, AsyncGenerator, Dict, List, Literal, + Optional, Set, Tuple, Union, cast, get_args, get_origin, + get_type_hints) import anyio import websockets @@ -38,74 +25,39 @@ from pydantic import BaseModel, Json from litellm._uuid import uuid from litellm.constants import ( - AIOHTTP_CONNECTOR_LIMIT, - AIOHTTP_CONNECTOR_LIMIT_PER_HOST, - AIOHTTP_KEEPALIVE_TIMEOUT, - AIOHTTP_NEEDS_CLEANUP_CLOSED, - AIOHTTP_TTL_DNS_CACHE, - AUDIO_SPEECH_CHUNK_SIZE, - BASE_MCP_ROUTE, - DEFAULT_MAX_RECURSE_DEPTH, - DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, - DEFAULT_SHARED_HEALTH_CHECK_TTL, - DEFAULT_SLACK_ALERTING_THRESHOLD, + AIOHTTP_CONNECTOR_LIMIT, AIOHTTP_CONNECTOR_LIMIT_PER_HOST, + AIOHTTP_KEEPALIVE_TIMEOUT, AIOHTTP_NEEDS_CLEANUP_CLOSED, + AIOHTTP_TTL_DNS_CACHE, AUDIO_SPEECH_CHUNK_SIZE, BASE_MCP_ROUTE, + DEFAULT_MAX_RECURSE_DEPTH, DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL, + DEFAULT_SHARED_HEALTH_CHECK_TTL, DEFAULT_SLACK_ALERTING_THRESHOLD, LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS, - LITELLM_SETTINGS_SAFE_DB_OVERRIDES, - LITELLM_UI_ALLOW_HEADERS, - LITELLM_UI_SESSION_DURATION, -) -from litellm.litellm_core_utils.litellm_logging import ( - _init_custom_logger_compatible_class, -) + LITELLM_SETTINGS_SAFE_DB_OVERRIDES, LITELLM_UI_ALLOW_HEADERS, + LITELLM_UI_SESSION_DURATION) +from litellm.litellm_core_utils.litellm_logging import \ + _init_custom_logger_compatible_class from litellm.litellm_core_utils.safe_json_dumps import safe_dumps -from litellm.proxy._types import ( - CallbackDelete, - CallInfo, - CommonProxyErrors, - ConfigFieldDelete, - ConfigFieldInfo, - ConfigFieldUpdate, - ConfigGeneralSettings, - ConfigList, - ConfigYAML, - EnterpriseLicenseData, - FieldDetail, - InvitationClaim, - InvitationDelete, - InvitationModel, - InvitationNew, - InvitationUpdate, - Litellm_EntityType, - LiteLLM_JWTAuth, - LiteLLM_TeamTable, - LiteLLM_UserTable, - LitellmUserRoles, - PassThroughGenericEndpoint, - ProxyErrorTypes, - ProxyException, - RoleBasedPermissions, - SpecialModelNames, - SupportedDBObjectType, - TeamDefaultSettings, - TokenCountRequest, - TransformRequestBody, - UserAPIKeyAuth, -) +from litellm.proxy._types import (CallbackDelete, CallInfo, CommonProxyErrors, + ConfigFieldDelete, ConfigFieldInfo, + ConfigFieldUpdate, ConfigGeneralSettings, + ConfigList, ConfigYAML, + EnterpriseLicenseData, FieldDetail, + InvitationClaim, InvitationDelete, + InvitationModel, InvitationNew, + InvitationUpdate, Litellm_EntityType, + LiteLLM_JWTAuth, LiteLLM_TeamTable, + LiteLLM_UserTable, LitellmUserRoles, + PassThroughGenericEndpoint, ProxyErrorTypes, + ProxyException, RoleBasedPermissions, + SpecialModelNames, SupportedDBObjectType, + TeamDefaultSettings, TokenCountRequest, + TransformRequestBody, UserAPIKeyAuth) from litellm.proxy.common_utils.callback_utils import ( - normalize_callback_names, - process_callback, -) + normalize_callback_names, process_callback) from litellm.proxy.common_utils.realtime_utils import _realtime_request_body -from litellm.types.utils import ( - ModelResponse, - ModelResponseStream, - TextCompletionResponse, - TokenCountResponse, -) -from litellm.utils import ( - _invalidate_model_cost_lowercase_map, - load_credentials_from_list, -) +from litellm.types.utils import (ModelResponse, ModelResponseStream, + TextCompletionResponse, TokenCountResponse) +from litellm.utils import (_invalidate_model_cost_lowercase_map, + load_credentials_from_list) if TYPE_CHECKING: from aiohttp import ClientSession @@ -202,173 +154,139 @@ from litellm import Router from litellm._logging import verbose_proxy_logger, verbose_router_logger from litellm.caching.caching import DualCache, RedisCache from litellm.caching.redis_cluster_cache import RedisClusterCache -from litellm.constants import ( - _REALTIME_BODY_CACHE_SIZE, - APSCHEDULER_COALESCE, - APSCHEDULER_MAX_INSTANCES, - APSCHEDULER_MISFIRE_GRACE_TIME, - APSCHEDULER_REPLACE_EXISTING, - DAYS_IN_A_MONTH, - DEFAULT_HEALTH_CHECK_INTERVAL, - DEFAULT_MODEL_CREATED_AT_TIME, - LITELLM_PROXY_ADMIN_NAME, - PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS, - PROXY_BATCH_POLLING_INTERVAL, - PROXY_BATCH_WRITE_AT, - PROXY_BUDGET_RESCHEDULER_MAX_TIME, - PROXY_BUDGET_RESCHEDULER_MIN_TIME, -) +from litellm.constants import (_REALTIME_BODY_CACHE_SIZE, APSCHEDULER_COALESCE, + APSCHEDULER_MAX_INSTANCES, + APSCHEDULER_MISFIRE_GRACE_TIME, + APSCHEDULER_REPLACE_EXISTING, DAYS_IN_A_MONTH, + DEFAULT_HEALTH_CHECK_INTERVAL, + DEFAULT_MODEL_CREATED_AT_TIME, + LITELLM_PROXY_ADMIN_NAME, + PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS, + PROXY_BATCH_POLLING_INTERVAL, + PROXY_BATCH_WRITE_AT, + PROXY_BUDGET_RESCHEDULER_MAX_TIME, + PROXY_BUDGET_RESCHEDULER_MIN_TIME) from litellm.exceptions import RejectedRequestError from litellm.integrations.custom_guardrail import ModifyResponseException from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.core_helpers import ( - _get_parent_otel_span_from_kwargs, - get_litellm_metadata_from_kwargs, -) + _get_parent_otel_span_from_kwargs, get_litellm_metadata_from_kwargs) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor -from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.litellm_core_utils.sensitive_data_masker import SensitiveDataMasker -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler +from litellm.litellm_core_utils.litellm_logging import \ + Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.sensitive_data_masker import \ + SensitiveDataMasker +from litellm.llms.custom_httpx.http_handler import (AsyncHTTPHandler, + HTTPHandler) from litellm.llms.vertex_ai.vertex_llm_base import VertexBase -from litellm.proxy._experimental.mcp_server.discoverable_endpoints import ( - router as mcp_discoverable_endpoints_router, -) -from litellm.proxy._experimental.mcp_server.rest_endpoints import ( - router as mcp_rest_endpoints_router, -) +from litellm.proxy._experimental.mcp_server.discoverable_endpoints import \ + router as mcp_discoverable_endpoints_router +from litellm.proxy._experimental.mcp_server.rest_endpoints import \ + router as mcp_rest_endpoints_router from litellm.proxy._experimental.mcp_server.server import app as mcp_app -from litellm.proxy._experimental.mcp_server.tool_registry import ( - global_mcp_tool_registry, -) +from litellm.proxy._experimental.mcp_server.tool_registry import \ + global_mcp_tool_registry from litellm.proxy._types import * from litellm.proxy.agent_endpoints.a2a_endpoints import router as a2a_router from litellm.proxy.agent_endpoints.agent_registry import global_agent_registry -from litellm.proxy.agent_endpoints.endpoints import router as agent_endpoints_router +from litellm.proxy.agent_endpoints.endpoints import \ + router as agent_endpoints_router from litellm.proxy.agent_endpoints.model_list_helpers import ( - append_agents_to_model_group, - append_agents_to_model_info, -) -from litellm.proxy.analytics_endpoints.analytics_endpoints import ( - router as analytics_router, -) -from litellm.proxy.anthropic_endpoints.claude_code_endpoints import ( - claude_code_marketplace_router, -) -from litellm.proxy.anthropic_endpoints.endpoints import router as anthropic_router -from litellm.proxy.anthropic_endpoints.skills_endpoints import ( - router as anthropic_skills_router, -) -from litellm.proxy.auth.auth_checks import ( - ExperimentalUIJWTToken, - get_team_object, - log_db_metrics, -) + append_agents_to_model_group, append_agents_to_model_info) +from litellm.proxy.analytics_endpoints.analytics_endpoints import \ + router as analytics_router +from litellm.proxy.anthropic_endpoints.claude_code_endpoints import \ + claude_code_marketplace_router +from litellm.proxy.anthropic_endpoints.endpoints import \ + router as anthropic_router +from litellm.proxy.anthropic_endpoints.skills_endpoints import \ + router as anthropic_skills_router +from litellm.proxy.auth.auth_checks import (ExperimentalUIJWTToken, + get_team_object, log_db_metrics) from litellm.proxy.auth.auth_utils import check_response_size_is_safe from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.auth.litellm_license import LicenseCheck -from litellm.proxy.auth.model_checks import ( - get_all_fallbacks, - get_complete_model_list, - get_key_models, - get_mcp_server_ids, - get_team_models, -) +from litellm.proxy.auth.model_checks import (get_all_fallbacks, + get_complete_model_list, + get_key_models, + get_mcp_server_ids, + get_team_models) from litellm.proxy.auth.user_api_key_auth import ( - _fetch_global_spend_with_event_coordination, - user_api_key_auth, - user_api_key_auth_websocket, -) + _fetch_global_spend_with_event_coordination, user_api_key_auth, + user_api_key_auth_websocket) from litellm.proxy.batches_endpoints.endpoints import router as batches_router - ## Import All Misc routes here ## from litellm.proxy.caching_routes import router as caching_router from litellm.proxy.common_request_processing import ( - ProxyBaseLLMRequestProcessing, - create_response, -) -from litellm.proxy.common_utils.callback_utils import initialize_callbacks_on_proxy + ProxyBaseLLMRequestProcessing, create_response) +from litellm.proxy.common_utils.callback_utils import \ + initialize_callbacks_on_proxy from litellm.proxy.common_utils.debug_utils import init_verbose_loggers -from litellm.proxy.common_utils.debug_utils import router as debugging_endpoints_router +from litellm.proxy.common_utils.debug_utils import \ + router as debugging_endpoints_router from litellm.proxy.common_utils.encrypt_decrypt_utils import ( - decrypt_value_helper, - encrypt_value_helper, -) + decrypt_value_helper, encrypt_value_helper) from litellm.proxy.common_utils.html_forms.ui_login import build_ui_login_form from litellm.proxy.common_utils.http_parsing_utils import ( - _read_request_body, - _safe_get_request_headers, - check_file_size_under_limit, - get_form_data, -) + _read_request_body, _safe_get_request_headers, check_file_size_under_limit, + get_form_data) from litellm.proxy.common_utils.load_config_utils import ( - get_config_file_contents_from_gcs, - get_file_contents_from_s3, -) -from litellm.proxy.common_utils.openai_endpoint_utils import ( - remove_sensitive_info_from_deployment, -) + get_config_file_contents_from_gcs, get_file_contents_from_s3) +from litellm.proxy.common_utils.openai_endpoint_utils import \ + remove_sensitive_info_from_deployment from litellm.proxy.common_utils.proxy_state import ProxyState from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob from litellm.proxy.common_utils.swagger_utils import ERROR_RESPONSES -from litellm.proxy.container_endpoints.endpoints import router as container_router -from litellm.proxy.credential_endpoints.endpoints import router as credential_router -from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import SpendLogCleanup +from litellm.proxy.container_endpoints.endpoints import \ + router as container_router +from litellm.proxy.credential_endpoints.endpoints import \ + router as credential_router +from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import \ + SpendLogCleanup from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.proxy.discovery_endpoints import ui_discovery_endpoints_router -from litellm.proxy.fine_tuning_endpoints.endpoints import router as fine_tuning_router -from litellm.proxy.fine_tuning_endpoints.endpoints import set_fine_tuning_config +from litellm.proxy.fine_tuning_endpoints.endpoints import \ + router as fine_tuning_router +from litellm.proxy.fine_tuning_endpoints.endpoints import \ + set_fine_tuning_config from litellm.proxy.google_endpoints.endpoints import router as google_router -from litellm.proxy.guardrails.guardrail_endpoints import router as guardrails_router -from litellm.proxy.guardrails.init_guardrails import ( - init_guardrails_v2, - initialize_guardrails, -) +from litellm.proxy.guardrails.guardrail_endpoints import \ + router as guardrails_router +from litellm.proxy.guardrails.init_guardrails import (init_guardrails_v2, + initialize_guardrails) from litellm.proxy.health_check import perform_health_check -from litellm.proxy.health_endpoints._health_endpoints import router as health_router -from litellm.proxy.hooks.model_max_budget_limiter import ( - _PROXY_VirtualKeyModelMaxBudgetLimiter, -) -from litellm.proxy.hooks.prompt_injection_detection import ( - _OPTIONAL_PromptInjectionDetection, -) +from litellm.proxy.health_endpoints._health_endpoints import \ + router as health_router +from litellm.proxy.hooks.model_max_budget_limiter import \ + _PROXY_VirtualKeyModelMaxBudgetLimiter +from litellm.proxy.hooks.prompt_injection_detection import \ + _OPTIONAL_PromptInjectionDetection from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger from litellm.proxy.image_endpoints.endpoints import router as image_router from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request -from litellm.proxy.management_endpoints.access_group_endpoints import ( - router as access_group_router, -) -from litellm.proxy.management_endpoints.budget_management_endpoints import ( - router as budget_management_router, -) -from litellm.proxy.management_endpoints.cache_settings_endpoints import ( - router as cache_settings_router, -) -from litellm.proxy.management_endpoints.callback_management_endpoints import ( - router as callback_management_endpoints_router, -) +from litellm.proxy.management_endpoints.access_group_endpoints import \ + router as access_group_router +from litellm.proxy.management_endpoints.budget_management_endpoints import \ + router as budget_management_router +from litellm.proxy.management_endpoints.cache_settings_endpoints import \ + router as cache_settings_router +from litellm.proxy.management_endpoints.callback_management_endpoints import \ + router as callback_management_endpoints_router from litellm.proxy.management_endpoints.common_utils import ( - _user_has_admin_privileges, - admin_can_invite_user, -) -from litellm.proxy.management_endpoints.compliance_endpoints import ( - router as compliance_router, -) -from litellm.proxy.management_endpoints.cost_tracking_settings import ( - router as cost_tracking_settings_router, -) -from litellm.proxy.management_endpoints.customer_endpoints import ( - router as customer_router, -) -from litellm.proxy.management_endpoints.fallback_management_endpoints import ( - router as fallback_management_router, -) -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, -) + _user_has_admin_privileges, admin_can_invite_user) +from litellm.proxy.management_endpoints.compliance_endpoints import \ + router as compliance_router +from litellm.proxy.management_endpoints.cost_tracking_settings import \ + router as cost_tracking_settings_router +from litellm.proxy.management_endpoints.customer_endpoints import \ + router as customer_router +from litellm.proxy.management_endpoints.fallback_management_endpoints import \ + router as fallback_management_router +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.key_management_endpoints import ( delete_verification_tokens, duration_in_seconds, @@ -387,171 +305,126 @@ from litellm.proxy.management_endpoints.model_access_group_management_endpoints router as model_access_group_management_router, ) from litellm.proxy.management_endpoints.model_management_endpoints import ( - _add_model_to_db, - _add_team_model_to_db, - _deduplicate_litellm_router_models, -) -from litellm.proxy.management_endpoints.model_management_endpoints import ( - router as model_management_router, -) -from litellm.proxy.management_endpoints.organization_endpoints import ( - router as organization_router, -) -from litellm.proxy.management_endpoints.policy_endpoints import router as policy_router -from litellm.proxy.management_endpoints.project_endpoints import ( - router as project_router, -) -from litellm.proxy.management_endpoints.router_settings_endpoints import ( - router as router_settings_router, -) + _add_model_to_db, _add_team_model_to_db, + _deduplicate_litellm_router_models) +from litellm.proxy.management_endpoints.model_management_endpoints import \ + router as model_management_router +from litellm.proxy.management_endpoints.organization_endpoints import \ + router as organization_router +from litellm.proxy.management_endpoints.policy_endpoints import \ + router as policy_router +from litellm.proxy.management_endpoints.project_endpoints import \ + router as project_router +from litellm.proxy.management_endpoints.router_settings_endpoints import \ + router as router_settings_router from litellm.proxy.management_endpoints.scim.scim_v2 import scim_router -from litellm.proxy.management_endpoints.tag_management_endpoints import ( - router as tag_management_router, -) -from litellm.proxy.management_endpoints.team_callback_endpoints import ( - router as team_callback_router, -) -from litellm.proxy.management_endpoints.team_endpoints import router as team_router +from litellm.proxy.management_endpoints.tag_management_endpoints import \ + router as tag_management_router +from litellm.proxy.management_endpoints.team_callback_endpoints import \ + router as team_callback_router +from litellm.proxy.management_endpoints.team_endpoints import \ + router as team_router from litellm.proxy.management_endpoints.team_endpoints import ( - update_team, - validate_membership, -) -from litellm.proxy.management_endpoints.tool_management_endpoints import ( - router as tool_management_router, -) -from litellm.proxy.management_endpoints.ui_sso import ( - get_disabled_non_admin_personal_key_creation, -) + update_team, validate_membership) +from litellm.proxy.management_endpoints.tool_management_endpoints import \ + router as tool_management_router +from litellm.proxy.management_endpoints.ui_sso import \ + get_disabled_non_admin_personal_key_creation from litellm.proxy.management_endpoints.ui_sso import router as ui_sso_router -from litellm.proxy.management_endpoints.usage_endpoints import router as usage_ai_router -from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import ( - router as user_agent_analytics_router, -) -from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update -from litellm.proxy.middleware.in_flight_requests_middleware import ( - InFlightRequestsMiddleware, -) -from litellm.proxy.middleware.prometheus_auth_middleware import PrometheusAuthMiddleware +from litellm.proxy.management_endpoints.usage_endpoints import \ + router as usage_ai_router +from litellm.proxy.management_endpoints.user_agent_analytics_endpoints import \ + router as user_agent_analytics_router +from litellm.proxy.management_helpers.audit_logs import \ + create_audit_log_for_update +from litellm.proxy.middleware.in_flight_requests_middleware import \ + InFlightRequestsMiddleware +from litellm.proxy.middleware.prometheus_auth_middleware import \ + PrometheusAuthMiddleware from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router -from litellm.proxy.openai_evals_endpoints.endpoints import router as evals_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.pass_through_endpoints.llm_passthrough_endpoints import ( - passthrough_endpoint_router, -) -from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( - router as llm_passthrough_router, -) -from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( - vertex_ai_live_websocket_passthrough, -) -from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( - initialize_pass_through_endpoints, -) -from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( - router as pass_through_router, -) -from litellm.proxy.policy_engine.policy_endpoints import router as policy_crud_router -from litellm.proxy.policy_engine.policy_resolve_endpoints import ( - router as policy_resolve_router, -) +from litellm.proxy.openai_evals_endpoints.endpoints import \ + router as evals_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.pass_through_endpoints.llm_passthrough_endpoints import \ + passthrough_endpoint_router +from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import \ + router as llm_passthrough_router +from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import \ + vertex_ai_live_websocket_passthrough +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import \ + initialize_pass_through_endpoints +from litellm.proxy.pass_through_endpoints.pass_through_endpoints import \ + router as pass_through_router +from litellm.proxy.policy_engine.policy_endpoints import \ + router as policy_crud_router +from litellm.proxy.policy_engine.policy_resolve_endpoints import \ + router as policy_resolve_router from litellm.proxy.prompts.prompt_endpoints import router as prompts_router from litellm.proxy.public_endpoints import router as public_endpoints_router from litellm.proxy.rag_endpoints.endpoints import router as rag_router from litellm.proxy.rerank_endpoints.endpoints import router as rerank_router -from litellm.proxy.response_api_endpoints.endpoints import router as response_router +from litellm.proxy.response_api_endpoints.endpoints import \ + router as response_router from litellm.proxy.route_llm_request import route_request from litellm.proxy.search_endpoints.endpoints import router as search_router -from litellm.proxy.search_endpoints.search_tool_management import ( - router as search_tool_management_router, -) -from litellm.proxy.spend_tracking.cloudzero_endpoints import router as cloudzero_router -from litellm.proxy.spend_tracking.spend_management_endpoints import ( - router as spend_management_router, -) -from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload +from litellm.proxy.search_endpoints.search_tool_management import \ + router as search_tool_management_router +from litellm.proxy.spend_tracking.cloudzero_endpoints import \ + router as cloudzero_router +from litellm.proxy.spend_tracking.spend_management_endpoints import \ + router as spend_management_router +from litellm.proxy.spend_tracking.spend_tracking_utils import \ + get_logging_payload from litellm.proxy.types_utils.utils import get_instance_fn -from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( - router as ui_crud_endpoints_router, -) -from litellm.proxy.utils import ( - PrismaClient, - ProxyLogging, - ProxyUpdateSpend, - _cache_user_row, - _get_docs_url, - _get_projected_spend_over_limit, - _get_redoc_url, - _is_projected_spend_over_limit, - _is_valid_team_configs, - get_custom_url, - get_error_message_str, - get_server_root_path, - handle_exception_on_proxy, - hash_token, - model_dump_with_preserved_fields, - update_spend, -) -from litellm.proxy.vector_store_endpoints.endpoints import router as vector_store_router -from litellm.proxy.vector_store_endpoints.management_endpoints import ( - router as vector_store_management_router, -) -from litellm.proxy.vector_store_files_endpoints.endpoints import ( - router as vector_store_files_router, -) -from litellm.proxy.vertex_ai_endpoints.langfuse_endpoints import ( - router as langfuse_router, -) +from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import \ + router as ui_crud_endpoints_router +from litellm.proxy.utils import (PrismaClient, ProxyLogging, ProxyUpdateSpend, + _cache_user_row, _get_docs_url, + _get_projected_spend_over_limit, + _get_redoc_url, + _is_projected_spend_over_limit, + _is_valid_team_configs, get_custom_url, + get_error_message_str, get_server_root_path, + handle_exception_on_proxy, hash_token, + model_dump_with_preserved_fields, + update_spend) +from litellm.proxy.vector_store_endpoints.endpoints import \ + router as vector_store_router +from litellm.proxy.vector_store_endpoints.management_endpoints import \ + router as vector_store_management_router +from litellm.proxy.vector_store_files_endpoints.endpoints import \ + router as vector_store_files_router +from litellm.proxy.vertex_ai_endpoints.langfuse_endpoints import \ + router as langfuse_router from litellm.proxy.video_endpoints.endpoints import router as video_router -from litellm.router import ( - AssistantsTypedDict, - Deployment, - LiteLLM_Params, - ModelGroupInfo, -) +from litellm.router import (AssistantsTypedDict, Deployment, LiteLLM_Params, + ModelGroupInfo) from litellm.scheduler import FlowItem, Scheduler from litellm.secret_managers.aws_secret_manager import load_aws_kms from litellm.secret_managers.google_kms import load_google_kms -from litellm.secret_managers.main import ( - get_secret, - get_secret_bool, - get_secret_str, - str_to_bool, -) +from litellm.secret_managers.main import (get_secret, get_secret_bool, + get_secret_str, str_to_bool) from litellm.types.integrations.slack_alerting import SlackAlertingArgs -from litellm.types.llms.anthropic import ( - AnthropicMessagesRequest, - AnthropicResponse, - AnthropicResponseContentBlockText, - AnthropicResponseUsageBlock, -) +from litellm.types.llms.anthropic import (AnthropicMessagesRequest, + AnthropicResponse, + AnthropicResponseContentBlockText, + AnthropicResponseUsageBlock) from litellm.types.llms.openai import HttpxBinaryResponseContent -from litellm.types.proxy.management_endpoints.model_management_endpoints import ( - ModelGroupInfoProxy, -) +from litellm.types.proxy.management_endpoints.model_management_endpoints import \ + ModelGroupInfoProxy from litellm.types.proxy.management_endpoints.ui_sso import ( - DefaultTeamSSOParams, - LiteLLM_UpperboundKeyGenerateParams, -) + DefaultTeamSSOParams, 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, - SearchToolTypedDict, - updateDeployment, -) +from litellm.types.router import (RouterGeneralSettings, SearchToolTypedDict, + updateDeployment) from litellm.types.scheduler import DefaultPriorities -from litellm.types.secret_managers.main import ( - KeyManagementSettings, - KeyManagementSystem, -) +from litellm.types.secret_managers.main import (KeyManagementSettings, + KeyManagementSystem) from litellm.types.utils import CredentialItem, CustomHuggingfaceTokenizer from litellm.types.utils import ModelInfo as ModelMapInfo from litellm.types.utils import RawRequestTypedDict, StandardLoggingPayload @@ -565,34 +438,15 @@ litellm.suppress_debug_info = True import json from typing import Union -from fastapi import ( - Depends, - FastAPI, - File, - Form, - Header, - HTTPException, - Path, - Query, - Request, - Response, - UploadFile, - WebSocket, - WebSocketDisconnect, - applications, - status, -) +from fastapi import (Depends, FastAPI, File, Form, Header, HTTPException, Path, + Query, Request, Response, UploadFile, WebSocket, + WebSocketDisconnect, applications, status) from fastapi.encoders import jsonable_encoder from fastapi.middleware.cors import CORSMiddleware from fastapi.openapi.docs import get_swagger_ui_html from fastapi.openapi.utils import get_openapi -from fastapi.responses import ( - FileResponse, - JSONResponse, - ORJSONResponse, - RedirectResponse, - StreamingResponse, -) +from fastapi.responses import (FileResponse, JSONResponse, ORJSONResponse, + RedirectResponse, StreamingResponse) from fastapi.routing import APIRouter from fastapi.security import OAuth2PasswordBearer from fastapi.security.api_key import APIKeyHeader @@ -615,7 +469,8 @@ except Exception: ################### # Import enterprise routes try: - from litellm_enterprise.proxy.enterprise_routes import router as _enterprise_router + from litellm_enterprise.proxy.enterprise_routes import \ + router as _enterprise_router from litellm_enterprise.proxy.proxy_server import EnterpriseProxyConfig enterprise_router = _enterprise_router @@ -957,9 +812,8 @@ def get_openapi_schema(): return app.openapi_schema # Use compatibility wrapper for FastAPI 0.120+ schema generation - from litellm.proxy.common_utils.openapi_schema_compat import ( - get_openapi_schema_with_compat, - ) + from litellm.proxy.common_utils.openapi_schema_compat import \ + get_openapi_schema_with_compat openapi_schema = get_openapi_schema_with_compat( get_openapi_func=get_openapi, @@ -1013,7 +867,8 @@ def get_openapi_schema(): } # Add LLM API request schema bodies for documentation - from litellm.proxy.common_utils.custom_openapi_spec import CustomOpenAPISpec + from litellm.proxy.common_utils.custom_openapi_spec import \ + CustomOpenAPISpec openapi_schema = CustomOpenAPISpec.add_llm_api_request_schema_body(openapi_schema) @@ -1039,7 +894,8 @@ def custom_openapi(): openapi_schema["paths"] = paths_to_include # Add LLM API request schema bodies for documentation - from litellm.proxy.common_utils.custom_openapi_spec import CustomOpenAPISpec + from litellm.proxy.common_utils.custom_openapi_spec import \ + CustomOpenAPISpec openapi_schema = CustomOpenAPISpec.add_llm_api_request_schema_body(openapi_schema) @@ -2033,9 +1889,8 @@ def _schedule_background_health_check_db_save( return import time as time_module - from litellm.proxy.health_endpoints._health_endpoints import ( - _save_background_health_checks_to_db, - ) + from litellm.proxy.health_endpoints._health_endpoints import \ + _save_background_health_checks_to_db checked_by = ( shared_health_manager.pod_id @@ -2095,9 +1950,8 @@ async def _run_background_health_check(): # Initialize shared health check manager if Redis is available and feature is enabled shared_health_manager = None if use_shared_health_check and redis_usage_cache is not None: - from litellm.proxy.health_check_utils.shared_health_check_manager import ( - SharedHealthCheckManager, - ) + from litellm.proxy.health_check_utils.shared_health_check_manager import \ + SharedHealthCheckManager shared_health_manager = SharedHealthCheckManager( redis_cache=redis_usage_cache, @@ -2781,25 +2635,24 @@ class ProxyConfig: litellm.guardrail_name_config_map = guardrail_name_config_map elif key == "global_prompt_directory": - from litellm.integrations.dotprompt import ( - set_global_prompt_directory, - ) + from litellm.integrations.dotprompt import \ + set_global_prompt_directory set_global_prompt_directory(value) verbose_proxy_logger.info( f"{blue_color_code}Set Global Prompt Directory on LiteLLM Proxy{reset_color_code}" ) elif key == "global_bitbucket_config": - from litellm.integrations.bitbucket import ( - set_global_bitbucket_config, - ) + from litellm.integrations.bitbucket import \ + set_global_bitbucket_config set_global_bitbucket_config(value) verbose_proxy_logger.info( f"{blue_color_code}Set Global BitBucket Config on LiteLLM Proxy{reset_color_code}" ) elif key == "global_gitlab_config": - from litellm.integrations.gitlab import set_global_gitlab_config + from litellm.integrations.gitlab import \ + set_global_gitlab_config set_global_gitlab_config(value) verbose_proxy_logger.info( @@ -2870,9 +2723,8 @@ class ProxyConfig: callback ) if "prometheus" in callback: - from litellm.integrations.prometheus import ( - PrometheusLogger, - ) + from litellm.integrations.prometheus import \ + PrometheusLogger if PrometheusLogger is not None: verbose_proxy_logger.debug( @@ -3283,9 +3135,8 @@ class ProxyConfig: mcp_servers_config = config.get("mcp_servers", None) if mcp_servers_config: - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import \ + global_mcp_server_manager # Get mcp_aliases from litellm_settings if available litellm_settings = config.get("litellm_settings", {}) @@ -3298,7 +3149,8 @@ class ProxyConfig: ## VECTOR STORES vector_store_registry_config = config.get("vector_store_registry", None) if vector_store_registry_config: - from litellm.vector_stores.vector_store_registry import VectorStoreRegistry + from litellm.vector_stores.vector_store_registry import \ + VectorStoreRegistry if litellm.vector_store_registry is None: litellm.vector_store_registry = VectorStoreRegistry() @@ -3325,7 +3177,8 @@ class ProxyConfig: """ from litellm.proxy.policy_engine.init_policies import init_policies - from litellm.proxy.policy_engine.policy_validator import PolicyValidator + from litellm.proxy.policy_engine.policy_validator import \ + PolicyValidator if config is None: verbose_proxy_logger.debug("Policy engine: config is None, skipping") @@ -3414,9 +3267,8 @@ class ProxyConfig: key_management_system == KeyManagementSystem.AWS_SECRET_MANAGER.value # noqa: F405 ): - from litellm.secret_managers.aws_secret_manager_v2 import ( - AWSSecretsManagerV2, - ) + from litellm.secret_managers.aws_secret_manager_v2 import \ + AWSSecretsManagerV2 AWSSecretsManagerV2.load_aws_secret_manager( use_aws_secret_manager=True, @@ -3427,28 +3279,24 @@ class ProxyConfig: elif ( key_management_system == KeyManagementSystem.GOOGLE_SECRET_MANAGER.value ): - from litellm.secret_managers.google_secret_manager import ( - GoogleSecretManager, - ) + from litellm.secret_managers.google_secret_manager import \ + GoogleSecretManager GoogleSecretManager() elif key_management_system == KeyManagementSystem.HASHICORP_VAULT.value: - from litellm.secret_managers.hashicorp_secret_manager import ( - HashicorpSecretManager, - ) + from litellm.secret_managers.hashicorp_secret_manager import \ + HashicorpSecretManager HashicorpSecretManager() elif key_management_system == KeyManagementSystem.CYBERARK.value: - from litellm.secret_managers.cyberark_secret_manager import ( - CyberArkSecretManager, - ) + from litellm.secret_managers.cyberark_secret_manager import \ + CyberArkSecretManager CyberArkSecretManager() elif key_management_system == KeyManagementSystem.CUSTOM.value: ### LOAD CUSTOM SECRET MANAGER ### - from litellm.secret_managers.custom_secret_manager_loader import ( - load_custom_secret_manager, - ) + from litellm.secret_managers.custom_secret_manager_loader import \ + load_custom_secret_manager load_custom_secret_manager(config_file_path=config_file_path) else: @@ -4030,9 +3878,8 @@ class ProxyConfig: # Schedule new job if retention period is set (not None) retention_period = general_settings.get("maximum_spend_logs_retention_period") if retention_period is not None: - from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import ( - SpendLogCleanup, - ) + from litellm.proxy.db.db_transaction_queue.spend_log_cleanup import \ + SpendLogCleanup spend_log_cleanup = SpendLogCleanup() cleanup_cron = general_settings.get("maximum_spend_logs_cleanup_cron") @@ -4059,9 +3906,8 @@ class ProxyConfig: ) else: # Interval-based scheduling (existing behavior) - from litellm.litellm_core_utils.duration_parser import ( - duration_in_seconds, - ) + from litellm.litellm_core_utils.duration_parser import \ + duration_in_seconds retention_interval = general_settings.get( "maximum_spend_logs_retention_interval", "1d" @@ -4428,9 +4274,8 @@ class ProxyConfig: if self._should_load_db_object(object_type="sso_settings"): await self._init_sso_settings_in_db(prisma_client=prisma_client) if self._should_load_db_object(object_type="cache_settings"): - from litellm.proxy.management_endpoints.cache_settings_endpoints import ( - CacheSettingsManager, - ) + from litellm.proxy.management_endpoints.cache_settings_endpoints import \ + CacheSettingsManager await CacheSettingsManager.init_cache_settings_in_db( prisma_client=prisma_client, proxy_config=self @@ -4447,7 +4292,8 @@ class ProxyConfig: import json import litellm - from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.proxy.hooks.mcp_semantic_filter import \ + SemanticToolFilterHook try: # Load litellm_settings from DB @@ -4612,9 +4458,8 @@ class ProxyConfig: if should_reload: # Perform the reload - from litellm.litellm_core_utils.get_model_cost_map import ( - get_model_cost_map, - ) + from litellm.litellm_core_utils.get_model_cost_map import \ + get_model_cost_map model_cost_map_url = litellm.model_cost_map_url new_model_cost_map = get_model_cost_map(url=model_cost_map_url) @@ -4720,9 +4565,8 @@ class ProxyConfig: if should_reload: # Perform the reload - from litellm.anthropic_beta_headers_manager import ( - reload_beta_headers_config, - ) + from litellm.anthropic_beta_headers_manager import \ + reload_beta_headers_config new_config = reload_beta_headers_config() @@ -4773,12 +4617,14 @@ class ProxyConfig: Returns: The PromptSpec object """ - from litellm.proxy.prompts.prompt_endpoints import create_versioned_prompt_spec + from litellm.proxy.prompts.prompt_endpoints import \ + create_versioned_prompt_spec return create_versioned_prompt_spec(db_prompt=db_prompt) async def _init_prompts_in_db(self, prisma_client: PrismaClient): - from litellm.proxy.prompts.prompt_registry import IN_MEMORY_PROMPT_REGISTRY + from litellm.proxy.prompts.prompt_registry import \ + IN_MEMORY_PROMPT_REGISTRY from litellm.types.prompts.init_prompts import PromptSpec try: @@ -4796,10 +4642,7 @@ class ProxyConfig: async def _init_guardrails_in_db(self, prisma_client: PrismaClient): from litellm.proxy.guardrails.guardrail_registry import ( - IN_MEMORY_GUARDRAIL_HANDLER, - Guardrail, - GuardrailRegistry, - ) + IN_MEMORY_GUARDRAIL_HANDLER, Guardrail, GuardrailRegistry) try: guardrails_in_db: List[ @@ -4825,10 +4668,10 @@ class ProxyConfig: """ Initialize policies and policy attachments from database into the in-memory registries. """ - from litellm.proxy.policy_engine.attachment_registry import ( - get_attachment_registry, - ) - from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.proxy.policy_engine.attachment_registry import \ + get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import \ + get_policy_registry try: # Get the global singleton instances @@ -4872,7 +4715,8 @@ class ProxyConfig: ) async def _init_vector_stores_in_db(self, prisma_client: PrismaClient): - from litellm.vector_stores.vector_store_registry import VectorStoreRegistry + from litellm.vector_stores.vector_store_registry import \ + VectorStoreRegistry try: # read vector stores from db table @@ -4899,7 +4743,8 @@ class ProxyConfig: ) async def _init_vector_store_indexes_in_db(self, prisma_client: PrismaClient): - from litellm.vector_stores.vector_store_registry import VectorStoreIndexRegistry + from litellm.vector_stores.vector_store_registry import \ + VectorStoreIndexRegistry try: # read vector stores from db table @@ -4929,7 +4774,8 @@ class ProxyConfig: ) async def _init_mcp_servers_in_db(self): - from litellm.proxy._experimental.mcp_server.utils import is_mcp_available + from litellm.proxy._experimental.mcp_server.utils import \ + is_mcp_available if not is_mcp_available(): verbose_proxy_logger.debug( @@ -4937,9 +4783,8 @@ class ProxyConfig: ) return - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import \ + global_mcp_server_manager try: await global_mcp_server_manager.reload_servers_from_database() @@ -4951,9 +4796,8 @@ class ProxyConfig: ) async def _init_agents_in_db(self, prisma_client: PrismaClient): - from litellm.proxy.agent_endpoints.agent_registry import ( - global_agent_registry as AGENT_REGISTRY, - ) + from litellm.proxy.agent_endpoints.agent_registry import \ + global_agent_registry as AGENT_REGISTRY try: db_agents = await AGENT_REGISTRY.get_all_agents_from_db( @@ -4976,9 +4820,8 @@ class ProxyConfig: """ global llm_router - from litellm.proxy.search_endpoints.search_tool_registry import ( - SearchToolRegistry, - ) + from litellm.proxy.search_endpoints.search_tool_registry import \ + SearchToolRegistry from litellm.router_utils.search_api_router import SearchAPIRouter try: @@ -5018,9 +4861,8 @@ class ProxyConfig: ) async def _init_pass_through_endpoints_in_db(self): - from litellm.proxy.pass_through_endpoints.pass_through_endpoints import ( - initialize_pass_through_endpoints_in_db, - ) + from litellm.proxy.pass_through_endpoints.pass_through_endpoints import \ + initialize_pass_through_endpoints_in_db await initialize_pass_through_endpoints_in_db() @@ -5118,11 +4960,8 @@ async def initialize( # noqa: PLR0915 if debug is True: # this needs to be first, so users can see Router init debugg import logging - from litellm._logging import ( - verbose_logger, - verbose_proxy_logger, - verbose_router_logger, - ) + from litellm._logging import (verbose_logger, verbose_proxy_logger, + verbose_router_logger) # this must ALWAYS remain logging.INFO, DO NOT MODIFY THIS verbose_logger.setLevel(level=logging.INFO) # sets package logs to info @@ -5131,11 +4970,8 @@ async def initialize( # noqa: PLR0915 if detailed_debug is True: import logging - from litellm._logging import ( - verbose_logger, - verbose_proxy_logger, - verbose_router_logger, - ) + from litellm._logging import (verbose_logger, verbose_proxy_logger, + verbose_router_logger) verbose_logger.setLevel(level=logging.DEBUG) # set package log to debug verbose_router_logger.setLevel(level=logging.DEBUG) # set router logs to debug @@ -5147,7 +4983,8 @@ async def initialize( # noqa: PLR0915 if litellm_log_setting.upper() == "INFO": import logging - from litellm._logging import verbose_proxy_logger, verbose_router_logger + from litellm._logging import (verbose_proxy_logger, + verbose_router_logger) # this must ALWAYS remain logging.INFO, DO NOT MODIFY THIS @@ -5160,11 +4997,9 @@ async def initialize( # noqa: PLR0915 elif litellm_log_setting.upper() == "DEBUG": import logging - from litellm._logging import ( - verbose_logger, - verbose_proxy_logger, - verbose_router_logger, - ) + from litellm._logging import (verbose_logger, + verbose_proxy_logger, + verbose_router_logger) verbose_logger.setLevel(level=logging.DEBUG) # set package log to debug verbose_router_logger.setLevel( @@ -5518,7 +5353,8 @@ class ProxyStartupEvent: litellm_settings: Dict[str, Any], ): """Initialize MCP semantic tool filter if configured""" - from litellm.proxy.hooks.mcp_semantic_filter import SemanticToolFilterHook + from litellm.proxy.hooks.mcp_semantic_filter import \ + SemanticToolFilterHook mcp_semantic_filter_config = litellm_settings.get( "mcp_semantic_tool_filter", None @@ -5630,9 +5466,8 @@ class ProxyStartupEvent: _teams = litellm.default_internal_user_params.get("teams") or [] if _teams and all(isinstance(team, dict) for team in _teams): - from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( - update_default_team_member_budget, - ) + from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import \ + update_default_team_member_budget teams_pydantic_obj = [NewUserRequestTeam(**team) for team in _teams] await update_default_team_member_budget( @@ -5853,9 +5688,8 @@ class ProxyStartupEvent: ### CHECK BATCH COST ### if llm_router is not None: try: - from litellm_enterprise.proxy.common_utils.check_batch_cost import ( - CheckBatchCost, - ) + from litellm_enterprise.proxy.common_utils.check_batch_cost import \ + CheckBatchCost check_batch_cost_job = CheckBatchCost( proxy_logging_obj=proxy_logging_obj, @@ -5884,9 +5718,8 @@ class ProxyStartupEvent: ### CHECK RESPONSES COST ### if llm_router is not None: try: - from litellm_enterprise.proxy.common_utils.check_responses_cost import ( - CheckResponsesCost, - ) + from litellm_enterprise.proxy.common_utils.check_responses_cost import \ + CheckResponsesCost check_responses_cost_job = CheckResponsesCost( proxy_logging_obj=proxy_logging_obj, @@ -5946,7 +5779,8 @@ class ProxyStartupEvent: ######################################################## from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger from litellm.integrations.focus.focus_logger import FocusLogger - from litellm.proxy.spend_tracking.cloudzero_endpoints import is_cloudzero_setup + from litellm.proxy.spend_tracking.cloudzero_endpoints import \ + is_cloudzero_setup if await is_cloudzero_setup(): await CloudZeroLogger.init_cloudzero_background_job(scheduler=scheduler) @@ -5968,17 +5802,15 @@ class ProxyStartupEvent: ######################################################## from litellm.constants import ( LITELLM_KEY_ROTATION_CHECK_INTERVAL_SECONDS, - LITELLM_KEY_ROTATION_ENABLED, - ) + LITELLM_KEY_ROTATION_ENABLED) key_rotation_enabled: Optional[bool] = str_to_bool(LITELLM_KEY_ROTATION_ENABLED) verbose_proxy_logger.debug(f"key_rotation_enabled: {key_rotation_enabled}") if key_rotation_enabled is True: try: - from litellm.proxy.common_utils.key_rotation_manager import ( - KeyRotationManager, - ) + from litellm.proxy.common_utils.key_rotation_manager import \ + KeyRotationManager # Get prisma_client from global scope global prisma_client @@ -6139,9 +5971,7 @@ class ProxyStartupEvent: Doc: https://docs.datadoghq.com/tracing/trace_collection/automatic_instrumentation/dd_libraries/python/ """ from litellm.litellm_core_utils.dd_tracing import ( - _should_use_dd_profiler, - _should_use_dd_tracer, - ) + _should_use_dd_profiler, _should_use_dd_tracer) if _should_use_dd_tracer(): import ddtrace @@ -6253,13 +6083,10 @@ async def model_list( """ global llm_model_list, general_settings, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj - from litellm.proxy.management_endpoints.common_utils import ( - _user_has_admin_privileges, - ) - from litellm.proxy.utils import ( - create_model_info_response, - get_available_models_for_user, - ) + from litellm.proxy.management_endpoints.common_utils import \ + _user_has_admin_privileges + from litellm.proxy.utils import (create_model_info_response, + get_available_models_for_user) # Validate scope parameter if provided if scope is not None and scope != "expand": @@ -6386,11 +6213,9 @@ async def model_info( """ global llm_model_list, general_settings, llm_router, prisma_client, user_api_key_cache, proxy_logging_obj - from litellm.proxy.utils import ( - create_model_info_response, - get_available_models_for_user, - validate_model_access, - ) + from litellm.proxy.utils import (create_model_info_response, + get_available_models_for_user, + validate_model_access) # Get available models for the user all_models = await get_available_models_for_user( @@ -6505,6 +6330,11 @@ async def chat_completion( # noqa: PLR0915 and user_api_key_dict.org_id is not None ): data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id + if ( + hasattr(user_api_key_dict, "agent_id") + and user_api_key_dict.agent_id is not None + ): + data["metadata"]["agent_id"] = user_api_key_dict.agent_id base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) try: result = await base_llm_response_processor.base_process_llm_request( @@ -6674,6 +6504,11 @@ async def completion( # noqa: PLR0915 and user_api_key_dict.org_id is not None ): data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id + if ( + hasattr(user_api_key_dict, "agent_id") + and user_api_key_dict.agent_id is not None + ): + data["metadata"]["agent_id"] = user_api_key_dict.agent_id base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) return await base_llm_response_processor.base_process_llm_request( request=request, @@ -6911,6 +6746,11 @@ async def embeddings( # noqa: PLR0915 and user_api_key_dict.org_id is not None ): data["metadata"]["user_api_key_org_id"] = user_api_key_dict.org_id + if ( + hasattr(user_api_key_dict, "agent_id") + and user_api_key_dict.agent_id is not None + ): + data["metadata"]["agent_id"] = user_api_key_dict.agent_id # Use unified request processor (same as chat/completions and responses) base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data) @@ -8344,7 +8184,8 @@ def _get_provider_token_counter( if deployment is None: return None - from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + from litellm.litellm_core_utils.get_llm_provider_logic import \ + get_llm_provider full_model = deployment.get("litellm_params", {}).get("model", "") model: Optional[str] = None @@ -10700,7 +10541,8 @@ async def fallback_login(request: Request): ) # hidden since this is a helper for UI sso login async def login(request: Request): # noqa: PLR0915 global premium_user, general_settings, master_key - from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.auth.login_utils import (authenticate_user, + create_ui_token_object) from litellm.proxy.utils import get_custom_url form = await request.form() @@ -10750,7 +10592,8 @@ async def login(request: Request): # noqa: PLR0915 ) # hidden helper for UI logins via API async def login_v2(request: Request): # noqa: PLR0915 global premium_user, general_settings, master_key - from litellm.proxy.auth.login_utils import authenticate_user, create_ui_token_object + from litellm.proxy.auth.login_utils import (authenticate_user, + create_ui_token_object) from litellm.proxy.utils import get_custom_url try: @@ -11079,7 +10922,8 @@ async def get_image(): if logo_path.startswith(("http://", "https://")): try: # Download the image and cache it - from litellm.llms.custom_httpx.http_handler import get_async_httpx_client + from litellm.llms.custom_httpx.http_handler import \ + get_async_httpx_client from litellm.types.llms.custom_http import httpxSpecialProvider async_client = get_async_httpx_client( @@ -11188,9 +11032,8 @@ async def new_invitation( ``` """ try: - from litellm.proxy.management_helpers.user_invitation import ( - create_invitation_for_user, - ) + from litellm.proxy.management_helpers.user_invitation import \ + create_invitation_for_user global prisma_client @@ -12277,7 +12120,8 @@ async def reload_model_cost_map( ) # Immediately reload the model cost map in the current pod - from litellm.litellm_core_utils.get_model_cost_map import get_model_cost_map + from litellm.litellm_core_utils.get_model_cost_map import \ + get_model_cost_map model_cost_map_url = litellm.model_cost_map_url new_model_cost_map = get_model_cost_map(url=model_cost_map_url) @@ -12576,9 +12420,8 @@ async def get_model_cost_map_source( ) try: - from litellm.litellm_core_utils.get_model_cost_map import ( - get_model_cost_map_source_info, - ) + from litellm.litellm_core_utils.get_model_cost_map import \ + get_model_cost_map_source_info source_info = get_model_cost_map_source_info() model_count = len(litellm.model_cost) if litellm.model_cost else 0 @@ -12630,7 +12473,8 @@ async def reload_anthropic_beta_headers( ) # Immediately reload the beta headers config in the current pod - from litellm.anthropic_beta_headers_manager import reload_beta_headers_config + from litellm.anthropic_beta_headers_manager import \ + reload_beta_headers_config new_config = reload_beta_headers_config() @@ -13038,9 +12882,8 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): """Handle dynamic MCP server routes like /github_mcp/mcp""" try: # Validate that the MCP server exists - from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( - global_mcp_server_manager, - ) + from litellm.proxy._experimental.mcp_server.mcp_server_manager import \ + global_mcp_server_manager from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.types.mcp import MCPAuth @@ -13059,9 +12902,8 @@ async def dynamic_mcp_route(mcp_server_name: str, request: Request): scope["path"] = f"/mcp/{mcp_server_name}" # Import the MCP handler - from litellm.proxy._experimental.mcp_server.server import ( - handle_streamable_http_mcp, - ) + from litellm.proxy._experimental.mcp_server.server import \ + handle_streamable_http_mcp # Create a custom send function to capture the response response_started = False diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 5e1ba479298..3ce9f054f25 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -66,6 +66,7 @@ model LiteLLM_AgentsTable { agent_access_groups String[] @default([]) object_permission_id String? object_permission LiteLLM_ObjectPermissionTable? @relation(fields: [object_permission_id], references: [object_permission_id]) + spend Float @default(0.0) created_at DateTime @default(now()) @map("created_at") created_by String updated_at DateTime @default(now()) @updatedAt @map("updated_at")