[⚡️ Python SDK import] - reduce python sdk import time by .3s (#12140)

* 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

@@ -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

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

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):

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

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
 

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

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()
 

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):

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"""
 

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
+    """

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"

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() 

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"

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())
 

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,

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,

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,