From 0c19414b3673470db5b9e23e285e5e0022cf6ed7 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 28 Jun 2025 14:57:10 -0700 Subject: [PATCH] =?UTF-8?q?[=E2=9A=A1=EF=B8=8F=20Python=20SDK=20import]=20?= =?UTF-8?q?-=20reduce=20python=20sdk=20import=20time=20by=20.3s=20=20(#121?= =?UTF-8?q?40)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * use 1 file for KeyManagementSystem * move key management settings * fix import locs * test_proxy_types_not_imported * test the import loc * fix import item * fix imports * fix import loc * fix imports --- .circleci/config.yml | 1 + litellm/__init__.py | 8 +- litellm/_service_logger.py | 4 +- litellm/integrations/custom_logger.py | 3 +- litellm/proxy/_types.py | 56 +---------- litellm/proxy/proxy_server.py | 8 ++ litellm/secret_managers/main.py | 2 +- litellm/types/llms/base.py | 25 ++++- .../proxy/management_endpoints/ui_sso.py | 22 ++++- litellm/types/secret_managers/main.py | 38 +++++++ litellm/types/utils.py | 27 +---- .../test_proxy_types_import.py | 98 +++++++++++++++++++ .../test_secret_manager.py | 4 +- .../test_key_management.py | 2 +- .../test_usage_endpoints.py | 2 +- .../test_e2e_pod_lock_manager.py | 2 +- .../test_key_generate_prisma.py | 4 +- 17 files changed, 210 insertions(+), 96 deletions(-) create mode 100644 litellm/types/secret_managers/main.py create mode 100644 tests/code_coverage_tests/test_proxy_types_import.py diff --git a/.circleci/config.yml b/.circleci/config.yml index 3294845bd9..52c2115bf5 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -1358,6 +1358,7 @@ jobs: # - run: python ./tests/documentation_tests/test_general_setting_keys.py - run: python ./tests/code_coverage_tests/check_licenses.py - run: python ./tests/code_coverage_tests/router_code_coverage.py + - run: python ./tests/code_coverage_tests/test_proxy_types_import.py - run: python ./tests/code_coverage_tests/callback_manager_test.py - run: python ./tests/code_coverage_tests/recursive_detector.py - run: python ./tests/code_coverage_tests/test_router_strategy_async.py diff --git a/litellm/__init__.py b/litellm/__init__.py index bd61fa6af4..9e932ce470 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -61,12 +61,8 @@ from litellm.constants import ( DEFAULT_ALLOWED_FAILS, ) from litellm.types.guardrails import GuardrailItem -from litellm.proxy._types import ( - KeyManagementSystem, - KeyManagementSettings, - LiteLLM_UpperboundKeyGenerateParams, -) -from litellm.types.proxy.management_endpoints.ui_sso import DefaultTeamSSOParams +from litellm.types.secret_managers.main import KeyManagementSystem, KeyManagementSettings +from litellm.types.proxy.management_endpoints.ui_sso import DefaultTeamSSOParams, LiteLLM_UpperboundKeyGenerateParams from litellm.types.utils import StandardKeyGenerationConfig, LlmProviders from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_callback_manager import LoggingCallbackManager diff --git a/litellm/_service_logger.py b/litellm/_service_logger.py index 969a9ef148..3128f02f40 100644 --- a/litellm/_service_logger.py +++ b/litellm/_service_logger.py @@ -4,7 +4,6 @@ from typing import TYPE_CHECKING, Any, Optional, Union import litellm from litellm._logging import verbose_logger -from litellm.proxy._types import UserAPIKeyAuth from .integrations.custom_logger import CustomLogger from .integrations.datadog.datadog import DataDogLogger @@ -15,11 +14,14 @@ from .types.services import ServiceLoggerPayload, ServiceTypes if TYPE_CHECKING: from opentelemetry.trace import Span as _Span + from litellm.proxy._types import UserAPIKeyAuth + Span = Union[_Span, Any] OTELClass = OpenTelemetry else: Span = Any OTELClass = Any + UserAPIKeyAuth = Any class ServiceLogging(CustomLogger): diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index ce97b9a292..1cbcd360ce 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -16,7 +16,6 @@ from typing import ( from pydantic import BaseModel from litellm.caching.caching import DualCache -from litellm.proxy._types import UserAPIKeyAuth from litellm.types.integrations.argilla import ArgillaItem from litellm.types.llms.openai import AllMessageValues, ChatCompletionRequest from litellm.types.utils import ( @@ -33,11 +32,13 @@ if TYPE_CHECKING: from opentelemetry.trace import Span as _Span from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy._types import UserAPIKeyAuth Span = Union[_Span, Any] else: Span = Any LiteLLMLoggingObj = Any + UserAPIKeyAuth = Any class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callback#callback-class diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index bc0332b33a..98156a5ad6 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -25,6 +25,7 @@ from litellm.types.mcp import ( MCPTransportType, ) from litellm.types.router import RouterErrors, UpdateRouterConfig +from litellm.types.secret_managers.main import KeyManagementSystem from litellm.types.utils import ( CallTypes, EmbeddingResponse, @@ -185,27 +186,6 @@ def hash_token(token: str): return hashed_token -class LiteLLM_UpperboundKeyGenerateParams(LiteLLMPydanticObjectBase): - """ - Set default upperbound to max budget a key called via `/key/generate` can be. - - Args: - max_budget (Optional[float], optional): Max budget a key can be. Defaults to None. - budget_duration (Optional[str], optional): Duration of the budget. Defaults to None. - duration (Optional[str], optional): Duration of the key. Defaults to None. - max_parallel_requests (Optional[int], optional): Max number of requests that can be made in parallel. Defaults to None. - tpm_limit (Optional[int], optional): Tpm limit. Defaults to None. - rpm_limit (Optional[int], optional): Rpm limit. Defaults to None. - """ - - max_budget: Optional[float] = None - budget_duration: Optional[str] = None - duration: Optional[str] = None - max_parallel_requests: Optional[int] = None - tpm_limit: Optional[int] = None - rpm_limit: Optional[int] = None - - class KeyManagementRoutes(str, enum.Enum): """ Enum for key management routes @@ -1398,40 +1378,6 @@ class DeleteOrganizationRequest(LiteLLMPydanticObjectBase): organization_ids: List[str] # required -class KeyManagementSystem(enum.Enum): - GOOGLE_KMS = "google_kms" - AZURE_KEY_VAULT = "azure_key_vault" - AWS_SECRET_MANAGER = "aws_secret_manager" - GOOGLE_SECRET_MANAGER = "google_secret_manager" - HASHICORP_VAULT = "hashicorp_vault" - LOCAL = "local" - AWS_KMS = "aws_kms" - - -class KeyManagementSettings(LiteLLMPydanticObjectBase): - hosted_keys: Optional[List] = None - store_virtual_keys: Optional[bool] = False - """ - If True, virtual keys created by litellm will be stored in the secret manager - """ - prefix_for_stored_virtual_keys: str = "litellm/" - """ - If set, this prefix will be used for stored virtual keys in the secret manager - """ - - access_mode: Literal["read_only", "write_only", "read_and_write"] = "read_only" - """ - Access mode for the secret manager, when write_only will only use for writing secrets - """ - - primary_secret_name: Optional[str] = None - """ - If set, will read secrets from this primary secret in the secret manager - - eg. on AWS you can store multiple secret values as K/V pairs in a single secret - """ - - class TeamDefaultSettings(LiteLLMPydanticObjectBase): team_id: str diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 1aa914585e..4234633c80 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -353,10 +353,18 @@ from litellm.types.llms.anthropic import ( AnthropicResponseUsageBlock, ) from litellm.types.llms.openai import HttpxBinaryResponseContent +from litellm.types.proxy.management_endpoints.ui_sso import ( + DefaultTeamSSOParams, + LiteLLM_UpperboundKeyGenerateParams, +) from litellm.types.router import DeploymentTypedDict from litellm.types.router import ModelInfo as RouterModelInfo from litellm.types.router import RouterGeneralSettings, updateDeployment from litellm.types.scheduler import DefaultPriorities +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 diff --git a/litellm/secret_managers/main.py b/litellm/secret_managers/main.py index b8b0cb2b51..6ecf5e370f 100644 --- a/litellm/secret_managers/main.py +++ b/litellm/secret_managers/main.py @@ -11,10 +11,10 @@ import litellm from litellm._logging import print_verbose, verbose_logger from litellm.caching.caching import DualCache from litellm.llms.custom_httpx.http_handler import HTTPHandler -from litellm.proxy._types import KeyManagementSystem from litellm.secret_managers.get_azure_ad_token_provider import ( get_azure_ad_token_provider, ) +from litellm.types.secret_managers.main import KeyManagementSystem oidc_cache = DualCache() diff --git a/litellm/types/llms/base.py b/litellm/types/llms/base.py index aec1438c48..29c8b049fb 100644 --- a/litellm/types/llms/base.py +++ b/litellm/types/llms/base.py @@ -1,4 +1,27 @@ -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict + + +class LiteLLMPydanticObjectBase(BaseModel): + """ + Implements default functions, all pydantic objects should have. + """ + + def json(self, **kwargs): # type: ignore + try: + return self.model_dump(**kwargs) # noqa + except Exception: + # if using pydantic v1 + return self.dict(**kwargs) + + def fields_set(self): + try: + return self.model_fields_set # noqa + except Exception: + # if using pydantic v1 + return self.__fields_set__ + + model_config = ConfigDict(protected_namespaces=()) + class BaseLiteLLMOpenAIResponseObject(BaseModel): diff --git a/litellm/types/proxy/management_endpoints/ui_sso.py b/litellm/types/proxy/management_endpoints/ui_sso.py index 3cdb5cb639..fe6d7569b2 100644 --- a/litellm/types/proxy/management_endpoints/ui_sso.py +++ b/litellm/types/proxy/management_endpoints/ui_sso.py @@ -2,9 +2,29 @@ from typing import List, Literal, Optional, TypedDict, Union from pydantic import Field -from litellm.proxy._types import LiteLLMPydanticObjectBase, LitellmUserRoles +from litellm.types.utils import LiteLLMPydanticObjectBase +class LiteLLM_UpperboundKeyGenerateParams(LiteLLMPydanticObjectBase): + """ + Set default upperbound to max budget a key called via `/key/generate` can be. + + Args: + max_budget (Optional[float], optional): Max budget a key can be. Defaults to None. + budget_duration (Optional[str], optional): Duration of the budget. Defaults to None. + duration (Optional[str], optional): Duration of the key. Defaults to None. + max_parallel_requests (Optional[int], optional): Max number of requests that can be made in parallel. Defaults to None. + tpm_limit (Optional[int], optional): Tpm limit. Defaults to None. + rpm_limit (Optional[int], optional): Rpm limit. Defaults to None. + """ + + max_budget: Optional[float] = None + budget_duration: Optional[str] = None + duration: Optional[str] = None + max_parallel_requests: Optional[int] = None + tpm_limit: Optional[int] = None + rpm_limit: Optional[int] = None + class MicrosoftGraphAPIUserGroupDirectoryObject(TypedDict, total=False): """Model for Microsoft Graph API directory object""" diff --git a/litellm/types/secret_managers/main.py b/litellm/types/secret_managers/main.py new file mode 100644 index 0000000000..e582df47df --- /dev/null +++ b/litellm/types/secret_managers/main.py @@ -0,0 +1,38 @@ +import enum +from typing import List, Literal, Optional + +from litellm.types.llms.base import LiteLLMPydanticObjectBase + + +class KeyManagementSystem(enum.Enum): + GOOGLE_KMS = "google_kms" + AZURE_KEY_VAULT = "azure_key_vault" + AWS_SECRET_MANAGER = "aws_secret_manager" + GOOGLE_SECRET_MANAGER = "google_secret_manager" + HASHICORP_VAULT = "hashicorp_vault" + LOCAL = "local" + AWS_KMS = "aws_kms" + + +class KeyManagementSettings(LiteLLMPydanticObjectBase): + hosted_keys: Optional[List] = None + store_virtual_keys: Optional[bool] = False + """ + If True, virtual keys created by litellm will be stored in the secret manager + """ + prefix_for_stored_virtual_keys: str = "litellm/" + """ + If set, this prefix will be used for stored virtual keys in the secret manager + """ + + access_mode: Literal["read_only", "write_only", "read_and_write"] = "read_only" + """ + Access mode for the secret manager, when write_only will only use for writing secrets + """ + + primary_secret_name: Optional[str] = None + """ + If set, will read secrets from this primary secret in the secret manager + + eg. on AWS you can store multiple secret values as K/V pairs in a single secret + """ \ No newline at end of file diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 9d76f04d97..a3e49d0b4d 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -33,7 +33,10 @@ from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, model_validator from typing_extensions import Callable, Dict, Required, TypedDict, override import litellm -from litellm.types.llms.base import BaseLiteLLMOpenAIResponseObject +from litellm.types.llms.base import ( + BaseLiteLLMOpenAIResponseObject, + LiteLLMPydanticObjectBase, +) from ..litellm_core_utils.core_helpers import map_finish_reason from .guardrails import GuardrailEventHooks @@ -63,28 +66,6 @@ def _generate_id(): # private helper function return "chatcmpl-" + str(uuid.uuid4()) -class LiteLLMPydanticObjectBase(BaseModel): - """ - Implements default functions, all pydantic objects should have. - """ - - def json(self, **kwargs): # type: ignore - try: - return self.model_dump(**kwargs) # noqa - except Exception: - # if using pydantic v1 - return self.dict(**kwargs) - - def fields_set(self): - try: - return self.model_fields_set # noqa - except Exception: - # if using pydantic v1 - return self.__fields_set__ - - model_config = ConfigDict(protected_namespaces=()) - - class LiteLLMCommonStrings(Enum): redacted_by_litellm = "redacted by litellm. 'litellm.turn_off_message_logging=True'" llm_provider_not_provided = "Unmapped LLM provider for this endpoint. You passed model={model}, custom_llm_provider={custom_llm_provider}. Check supported provider and route: https://docs.litellm.ai/docs/providers" diff --git a/tests/code_coverage_tests/test_proxy_types_import.py b/tests/code_coverage_tests/test_proxy_types_import.py new file mode 100644 index 0000000000..f0027d8870 --- /dev/null +++ b/tests/code_coverage_tests/test_proxy_types_import.py @@ -0,0 +1,98 @@ +import ast +import os +import sys + + +def test_proxy_types_not_imported(): + """ + Test that proxy._types is not directly imported in litellm/__init__.py + by examining the source code using AST parsing. + """ + # Read the litellm/__init__.py file + # local_init_file = "../litellm/" + init_file_path = os.path.join("./litellm", "__init__.py") + if not os.path.exists(init_file_path): + raise Exception(f"Could not find {init_file_path}") + + with open(init_file_path, "r") as f: + content = f.read() + lines = content.splitlines() # Get lines for line number reporting + + try: + tree = ast.parse(content) + except SyntaxError as e: + raise Exception(f"Could not parse {init_file_path}: {e}") + + # Check for direct imports of proxy._types + found_imports = [] + + for node in ast.walk(tree): + if isinstance(node, ast.Import): + for alias in node.names: + if "proxy._types" in alias.name or "proxy/_types" in alias.name: + line_num = node.lineno + line_content = lines[line_num - 1] if line_num <= len(lines) else "Unknown" + import_statement = f"import {alias.name}" + found_imports.append({ + 'type': 'import', + 'line': line_num, + 'content': line_content.strip(), + 'statement': import_statement, + 'module': alias.name + }) + + elif isinstance(node, ast.ImportFrom): + if node.module and ("proxy._types" in node.module or "proxy/_types" in node.module): + line_num = node.lineno + line_content = lines[line_num - 1] if line_num <= len(lines) else "Unknown" + import_names = [alias.name for alias in node.names] + import_statement = f"from {node.module} import {', '.join(import_names)}" + found_imports.append({ + 'type': 'from_import', + 'line': line_num, + 'content': line_content.strip(), + 'statement': import_statement, + 'module': node.module + }) + + if found_imports: + print("āŒ BAD, this can import time to import litellm. Found direct imports of proxy._types in litellm/__init__.py:") + print("=" * 80) + for imp in found_imports: + print(f"Line {imp['line']}: {imp['content']}") + print(f" Type: {imp['type']}") + print(f" Statement: {imp['statement']}") + print(f" Module: {imp['module']}") + print("-" * 80) + print("To fix this, please conditionally import this TYPE using TYPE_CHECKING") + + raise Exception( + f"Found {len(found_imports)} direct import(s) of proxy._types in litellm/__init__.py" + ) + + print("āœ“ No direct imports of proxy._types found in litellm/__init__.py") + return True + + +def main(): + """ + Main function to run the import test + """ + print("=" * 60) + print("Testing litellm import performance") + print("Checking that proxy._types is not directly imported from litellm/__init__.py") + print("=" * 60) + + try: + test_proxy_types_not_imported() + print("\n" + "=" * 60) + print("āœ“ Test passed! proxy._types is not directly imported from litellm/__init__.py") + print("=" * 60) + except Exception as e: + print(f"\nāŒ Test failed: {e}") + print("=" * 60) + sys.exit(1) + + +if __name__ == "__main__": + main() \ No newline at end of file diff --git a/tests/litellm_utils_tests/test_secret_manager.py b/tests/litellm_utils_tests/test_secret_manager.py index cd9230d937..cc0977cc51 100644 --- a/tests/litellm_utils_tests/test_secret_manager.py +++ b/tests/litellm_utils_tests/test_secret_manager.py @@ -296,7 +296,7 @@ def test_should_read_secret_from_secret_manager(): """ Test that _should_read_secret_from_secret_manager returns correct values based on access mode """ - from litellm.proxy._types import KeyManagementSettings + from litellm.types.secret_managers.main import KeyManagementSettings # Test when secret manager client is None litellm.secret_manager_client = None @@ -327,7 +327,7 @@ def test_get_secret_with_access_mode(): """ Test that get_secret respects access mode settings """ - from litellm.proxy._types import KeyManagementSettings + from litellm.types.secret_managers.main import KeyManagementSettings # Set up test environment test_secret_name = "TEST_SECRET_KEY" diff --git a/tests/proxy_admin_ui_tests/test_key_management.py b/tests/proxy_admin_ui_tests/test_key_management.py index 766d4835b3..b3522c8816 100644 --- a/tests/proxy_admin_ui_tests/test_key_management.py +++ b/tests/proxy_admin_ui_tests/test_key_management.py @@ -80,7 +80,6 @@ from litellm.proxy._types import ( DynamoDBArgs, GenerateKeyRequest, KeyRequest, - LiteLLM_UpperboundKeyGenerateParams, NewCustomerRequest, NewTeamRequest, NewUserRequest, @@ -92,6 +91,7 @@ from litellm.proxy._types import ( UpdateUserRequest, UserAPIKeyAuth, ) +from litellm.types.proxy.management_endpoints.ui_sso import LiteLLM_UpperboundKeyGenerateParams proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) diff --git a/tests/proxy_admin_ui_tests/test_usage_endpoints.py b/tests/proxy_admin_ui_tests/test_usage_endpoints.py index ab5e5fb834..61d4e94493 100644 --- a/tests/proxy_admin_ui_tests/test_usage_endpoints.py +++ b/tests/proxy_admin_ui_tests/test_usage_endpoints.py @@ -87,12 +87,12 @@ verbose_proxy_logger.setLevel(level=logging.DEBUG) from starlette.datastructures import URL from litellm.caching.caching import DualCache +from litellm.types.proxy.management_endpoints.ui_sso import LiteLLM_UpperboundKeyGenerateParams from litellm.proxy._types import ( DynamoDBArgs, GenerateKeyRequest, RegenerateKeyRequest, KeyRequest, - LiteLLM_UpperboundKeyGenerateParams, NewCustomerRequest, NewTeamRequest, NewUserRequest, diff --git a/tests/proxy_unit_tests/test_e2e_pod_lock_manager.py b/tests/proxy_unit_tests/test_e2e_pod_lock_manager.py index a9bede3498..5f4a16f4a0 100644 --- a/tests/proxy_unit_tests/test_e2e_pod_lock_manager.py +++ b/tests/proxy_unit_tests/test_e2e_pod_lock_manager.py @@ -75,11 +75,11 @@ verbose_proxy_logger.setLevel(level=logging.DEBUG) from starlette.datastructures import URL from litellm.caching.caching import DualCache, RedisCache +from litellm.types.proxy.management_endpoints.ui_sso import LiteLLM_UpperboundKeyGenerateParams from litellm.proxy._types import ( DynamoDBArgs, GenerateKeyRequest, KeyRequest, - LiteLLM_UpperboundKeyGenerateParams, NewCustomerRequest, NewTeamRequest, NewUserRequest, diff --git a/tests/proxy_unit_tests/test_key_generate_prisma.py b/tests/proxy_unit_tests/test_key_generate_prisma.py index 4b0df3e0f3..d8b826a387 100644 --- a/tests/proxy_unit_tests/test_key_generate_prisma.py +++ b/tests/proxy_unit_tests/test_key_generate_prisma.py @@ -94,11 +94,11 @@ verbose_proxy_logger.setLevel(level=logging.DEBUG) from starlette.datastructures import URL from litellm.caching.caching import DualCache +from litellm.types.proxy.management_endpoints.ui_sso import LiteLLM_UpperboundKeyGenerateParams from litellm.proxy._types import ( DynamoDBArgs, GenerateKeyRequest, KeyRequest, - LiteLLM_UpperboundKeyGenerateParams, NewCustomerRequest, NewTeamRequest, NewUserRequest, @@ -3608,7 +3608,7 @@ async def test_key_generate_with_secret_manager_call(prisma_client): assert it is deleted from the secret manager """ from litellm.secret_managers.aws_secret_manager_v2 import AWSSecretsManagerV2 - from litellm.proxy._types import KeyManagementSystem, KeyManagementSettings + from litellm.types.secret_managers.main import KeyManagementSystem, KeyManagementSettings from litellm.proxy.hooks.key_management_event_hooks import ( LITELLM_PREFIX_STORED_VIRTUAL_KEYS,