refactor: expose core private helpers under public names

enterprise/enterprise_hooks/__init__.py

@@ -1,15 +1,15 @@
 from typing import Dict, Literal, Type, Union
 
-from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles
+from litellm_enterprise.proxy.hooks.managed_files import PROXY_LiteLLMManagedFiles
 from litellm_enterprise.proxy.hooks.managed_vector_stores import (
-    _PROXY_LiteLLMManagedVectorStores,
+    PROXY_LiteLLMManagedVectorStores,
 )
 
 from litellm.integrations.custom_logger import CustomLogger
 
 ENTERPRISE_PROXY_HOOKS: Dict[str, Type[CustomLogger]] = {
-    "managed_files": _PROXY_LiteLLMManagedFiles,
-    "managed_vector_stores": _PROXY_LiteLLMManagedVectorStores,
+    "managed_files": PROXY_LiteLLMManagedFiles,
+    "managed_vector_stores": PROXY_LiteLLMManagedVectorStores,
 }
 
 

enterprise/enterprise_hooks/banned_keywords.py

@@ -20,7 +20,7 @@
 from fastapi import HTTPException
 
 
-class _ENTERPRISE_BannedKeywords(CustomLogger):
+class ENTERPRISE_BannedKeywords(CustomLogger):
     enforces_request_content: bool = True
     # Class variables or attributes
     def __init__(self):
@@ -114,3 +114,4 @@ async def async_post_call_streaming_hook(
         response: str,
     ):
         self.test_violation(test_str=response)
+_ENTERPRISE_BannedKeywords = ENTERPRISE_BannedKeywords

enterprise/enterprise_hooks/blocked_user_list.py

@@ -21,7 +21,7 @@
 from litellm.proxy.utils import PrismaClient
 
 
-class _ENTERPRISE_BlockedUserList(CustomLogger):
+class ENTERPRISE_BlockedUserList(CustomLogger):
     enforces_request_content: bool = True
     # Class variables or attributes
     def __init__(self, prisma_client: Optional[PrismaClient]):
@@ -128,3 +128,4 @@ async def async_pre_call_hook(
                     str(e)
                 )
             )
+_ENTERPRISE_BlockedUserList = ENTERPRISE_BlockedUserList

enterprise/enterprise_hooks/google_text_moderation.py

@@ -16,7 +16,7 @@
 from litellm.types.utils import CallTypesLiteral
 
 
-class _ENTERPRISE_GoogleTextModeration(CustomLogger):
+class ENTERPRISE_GoogleTextModeration(CustomLogger):
     user_api_key_cache = None
     confidence_categories = [
         "toxic",
@@ -125,9 +125,10 @@ async def async_moderation_hook(
                     )
             # Handle the response
             return data
+_ENTERPRISE_GoogleTextModeration = ENTERPRISE_GoogleTextModeration
 
 
-# google_text_moderation_obj = _ENTERPRISE_GoogleTextModeration()
+# google_text_moderation_obj = ENTERPRISE_GoogleTextModeration()
 # asyncio.run(
 #     google_text_moderation_obj.async_moderation_hook(
 #         data={"messages": [{"role": "user", "content": "Hey, how's it going?"}]}

enterprise/enterprise_hooks/openai_moderation.py

@@ -24,7 +24,7 @@
 from litellm.types.utils import CallTypesLiteral
 
 
-class _ENTERPRISE_OpenAI_Moderation(CustomLogger):
+class ENTERPRISE_OpenAI_Moderation(CustomLogger):
     @property
     def model_name(self) -> str:
         return litellm.openai_moderations_model_name or DEFAULT_OPENAI_MODERATIONS_MODEL
@@ -55,3 +55,4 @@ async def async_moderation_hook(
                 status_code=403, detail={"error": "Violated content safety policy"}
             )
         pass
+_ENTERPRISE_OpenAI_Moderation = ENTERPRISE_OpenAI_Moderation

enterprise/litellm_enterprise/enterprise_callbacks/send_emails/endpoints.py

@@ -3,9 +3,10 @@
 """
 
 import json
-from typing import Dict
+from typing import Dict, Final, cast
 
 from fastapi import APIRouter, Depends, HTTPException
+from pydantic import JsonValue
 from litellm_enterprise.types.enterprise_callbacks.send_emails import (
     DefaultEmailSettings,
     EmailEvent,
@@ -38,14 +39,15 @@ async def _get_email_settings(prisma_client) -> Dict[str, bool]:
             and general_settings_entry.param_value is not None
         ):
             # Get general settings value
-            if isinstance(general_settings_entry.param_value, str):
-                general_settings = json.loads(general_settings_entry.param_value)
-            else:
-                general_settings = general_settings_entry.param_value
+            general_settings: Final = (
+                cast(Dict[str, object], json.loads(general_settings_entry.param_value))
+                if isinstance(general_settings_entry.param_value, str)
+                else cast(Dict[str, object], general_settings_entry.param_value)
+            )
 
             # Extract email_settings from general settings if it exists
             if general_settings and "email_settings" in general_settings:
-                email_settings = general_settings["email_settings"]
+                email_settings: Final = cast(Dict[str, bool], general_settings["email_settings"])
                 # Update settings_dict with values from general_settings
                 for event_name, enabled in email_settings.items():
                     settings_dict[event_name] = enabled
@@ -64,7 +66,7 @@ async def _save_email_settings(prisma_client, settings: Dict[str, bool]):
     from litellm.proxy.proxy_server import proxy_config
 
     proxy_config.reject_config_owned_writes(
-        section_name="general_settings", changed_keys={"email_settings": settings}
+        section_name="general_settings", changed_keys={"email_settings": cast(JsonValue, settings)}
     )
     try:
         verbose_proxy_logger.debug(
@@ -77,16 +79,15 @@ async def _save_email_settings(prisma_client, settings: Dict[str, bool]):
         )
 
         # Initialize general settings dict
-        if (
-            general_settings_entry is not None
-            and general_settings_entry.param_value is not None
-        ):
-            if isinstance(general_settings_entry.param_value, str):
-                general_settings = json.loads(general_settings_entry.param_value)
-            else:
-                general_settings = dict(general_settings_entry.param_value)
-        else:
-            general_settings = {}
+        general_settings: Final = (
+            (
+                cast(Dict[str, object], json.loads(general_settings_entry.param_value))
+                if isinstance(general_settings_entry.param_value, str)
+                else cast(Dict[str, object], dict(general_settings_entry.param_value))
+            )
+            if general_settings_entry is not None and general_settings_entry.param_value is not None
+            else {}
+        )
 
         # Update email_settings in general_settings
         general_settings["email_settings"] = settings

enterprise/litellm_enterprise/integrations/custom_guardrail.py

@@ -36,7 +36,7 @@ def _should_run_if_mode_by_tag(
 
         proxy_server_request = data.get("proxy_server_request", {})
 
-        request_tags = StandardLoggingPayloadSetup._get_request_tags(
+        request_tags = StandardLoggingPayloadSetup.get_request_tags(
             litellm_params=data,
             proxy_server_request=proxy_server_request,
         )

enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py

@@ -60,6 +60,32 @@ def created_by(self) -> str | None: ...
     @property
     def file_object(self) -> object: ...
 
+    @property
+    def org_id(self) -> str | None: ...
+
+    @property
+    def api_key(self) -> str | None: ...
+
+    @property
+    def team_id(self) -> str | None: ...
+
+
+class _ReadableFileContent(Protocol):
+    async def read(self) -> bytes: ...
+
+
+class _HasFileContent(Protocol):
+    @property
+    def content(self) -> bytes: ...
+
+
+async def _file_content_bytes(file_content: object) -> bytes:
+    if hasattr(file_content, "content"):
+        return cast(_HasFileContent, file_content).content
+    if hasattr(file_content, "read"):
+        return await cast(_ReadableFileContent, file_content).read()
+    return cast(bytes, file_content)
+
 
 def _managed_object_table(prisma_client: "PrismaClient") -> "TableActions[_ManagedObjectRow]":
     return ManagedObjectRepository(prisma_client).table
@@ -175,17 +201,21 @@ async def _get_team_alias(self, team_id: str | None) -> str | None:
             return None
 
     async def _get_org_id(self, job: "_ManagedObjectRow", batch_id: str) -> str | None:
-        org_id = getattr(job, "org_id", None)
+        org_id: Final = getattr(job, "org_id", None)
         if org_id:
             return org_id
-        api_key = getattr(job, "api_key", None)
-        team_id = getattr(job, "team_id", None)
+        api_key: Final = getattr(job, "api_key", None)
+        team_id: Final = getattr(job, "team_id", None)
         if api_key:
             try:
                 key_row: prisma_models.LiteLLM_VerificationToken | None = await _token_table(
                     self.prisma_client
                 ).find_unique(where={"token": api_key})
-                key_org_id = getattr(key_row, "organization_id", None) if key_row is not None else None
+                key_org_id: Final = (
+                    cast(str | None, getattr(key_row, "organization_id", None))
+                    if key_row is not None
+                    else None
+                )
                 if key_org_id:
                     return key_org_id
             except Exception as e:
@@ -199,7 +229,7 @@ async def _get_org_id(self, job: "_ManagedObjectRow", batch_id: str) -> str | No
             team_row: prisma_models.LiteLLM_TeamTable | None = await _team_table(self.prisma_client).find_unique(
                 where={"team_id": team_id}
             )
-            return getattr(team_row, "organization_id", None) if team_row is not None else None
+            return cast(str | None, getattr(team_row, "organization_id", None)) if team_row is not None else None
         except Exception as e:
             verbose_proxy_logger.error(f"CheckBatchCost: could not resolve the team's org for batch {batch_id}: {e}")
             return None
@@ -673,16 +703,17 @@ def _get_input_file_id(job: "_ManagedObjectRow") -> Optional[str]:
 
         from litellm.types.utils import LiteLLMBatch
 
-        file_object = job.file_object
-        if isinstance(file_object, str):
-            try:
-                file_object = json.loads(file_object)
-            except (json.JSONDecodeError, ValueError):
-                return None
-        if not isinstance(file_object, dict):
+        file_object: Final = job.file_object
+        try:
+            parsed_file_object: Final[object] = (
+                json.loads(file_object) if isinstance(file_object, str) else file_object
+            )
+        except (json.JSONDecodeError, ValueError):
+            return None
+        if not isinstance(parsed_file_object, dict):
             return None
         try:
-            return LiteLLMBatch.model_validate(file_object).input_file_id
+            return LiteLLMBatch.model_validate(parsed_file_object).input_file_id
         except Exception:
             return None
 
@@ -704,10 +735,11 @@ async def _track_completed_batch_cost(
         """
         from litellm.batches.batch_utils import (
             count_error_file_failed_requests,
-            _get_file_content_as_dictionary,
+            get_file_content_as_dictionary,
             calculate_batch_cost_and_usage,
         )
         from litellm.files.main import afile_content
+        from litellm.files.types import FileContentCallOptions
         from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
         from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
         from litellm.litellm_core_utils.litellm_logging import deployment_pricing_model_info, mask_api_base_credentials
@@ -745,30 +777,22 @@ async def _track_completed_batch_cost(
             except (IndexError, AttributeError):
                 pass
 
-        credentials = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
+        credentials: Final = self.llm_router.get_deployment_credentials_with_provider(model_id) or {}
         _file_content = await afile_content(
-            file_id=raw_output_file_id,
+            file_id=cast(str, raw_output_file_id),
             _litellm_internal_model_credentials=MappingProxyType(dict(credentials)),
-            **credentials,
+            **cast(FileContentCallOptions, credentials),
         )
 
-        # Access content - handle both direct attribute and method call
-        if hasattr(_file_content, 'content'):
-            content_bytes = _file_content.content  # type: ignore[union-attr]
-        elif hasattr(_file_content, 'read'):
-            content_bytes = await _file_content.read()  # type: ignore[misc]
-        else:
-            content_bytes = _file_content  # type: ignore[assignment]
+        content_bytes: Final = await _file_content_bytes(_file_content)
 
-        file_content_as_dict = _get_file_content_as_dictionary(
-            content_bytes  # type: ignore[arg-type]
-        )
+        file_content_as_dict = get_file_content_as_dictionary(content_bytes)
 
         # Record output file size
         if prom_logger and content_bytes:
             try:
                 prom_logger.record_managed_file_size(
-                    size_bytes=len(content_bytes),  # type: ignore
+                    size_bytes=len(content_bytes),
                     purpose="batch",
                     file_type="output",
                     model=model_id,
@@ -805,7 +829,7 @@ async def _track_completed_batch_cost(
                 team_id=getattr(job, "team_id", None),
             )
             for _file_attr in ["output_file_id", "error_file_id"]:
-                _raw_file_id = getattr(response, _file_attr, None)
+                _raw_file_id = cast(str | None, getattr(response, _file_attr, None))
                 if _raw_file_id and not _is_base64_encoded_unified_file_id(_raw_file_id):
                     try:
                         _unified_file_id = managed_files_hook.get_unified_output_file_id(

enterprise/litellm_enterprise/proxy/enterprise_routes.py

@@ -8,7 +8,7 @@
 from .audit_logging_endpoints import router as audit_logging_router
 from .liteadmin import router as liteadmin_router
 from .management_endpoints import management_endpoints_router
-from .utils import _should_block_robots
+from .utils import should_block_robots
 
 __all__ = ["router", "ui_crud_endpoints"]
 
@@ -25,7 +25,7 @@ async def get_robots():
     Block all web crawlers from indexing the proxy server endpoints
     This is useful for ensuring that the API endpoints aren't indexed by search engines
     """
-    if _should_block_robots():
+    if should_block_robots():
         return Response(content="User-agent: *\nDisallow: /", media_type="text/plain")
     else:
         return Response(status_code=404)