From c091dd46087bc3a40f18ab0bc48dc08a2b0552a8 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 19:58:35 +0000 Subject: [PATCH 1/7] perf: lazy-load SDK symbols so import litellm stays under 60 MB RSS (#39121) * perf: lazy-load SDK symbols so import litellm stays under 60 MB RSS Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: resolve litellm.proxy submodules lazily so litellm.proxy._types stays importable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: correct SlackAlerting lazy mapping and keep eager encoding path importable Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix: register module-valued public names as module aliases instead of symbol imports Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: justify module-alias cache write with rebind-ok Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/__init__.py | 410 ++-- litellm/_lazy_imports.py | 96 +- litellm/_lazy_imports_registry.py | 2693 +++++++++++++++++++++++ litellm/proxy/__init__.py | 12 +- tests/test_litellm/test_lazy_imports.py | 87 + 5 files changed, 3095 insertions(+), 203 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 42c0ea881fd..b2bf3f09152 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -13,6 +13,7 @@ warnings.filterwarnings("ignore", message=".*`ReadOnly` qualifier.*") ### INIT VARIABLES ######################### import threading import os +import sys # Load .env before any other litellm imports so env vars (e.g. LITELLM_UI_SESSION_DURATION) are available import dotenv as _dotenv @@ -45,8 +46,6 @@ from typing import ( TYPE_CHECKING, Union, ) -from litellm.types.integrations.datadog import DatadogInitParams -from litellm.types.integrations.newrelic import NewRelicInitParams from litellm._logging import ( set_verbose, _turn_on_debug, @@ -95,8 +94,7 @@ from litellm.constants import ( DEFAULT_SOFT_BUDGET, DEFAULT_ALLOWED_FAILS, ) -import httpx - +# httpx is lazy-loaded via __getattr__ # register_async_client_cleanup is lazy-loaded and called on first access litellm_mode = os.getenv("LITELLM_MODE", "DEV") # "PRODUCTION", "DEV" @@ -364,8 +362,6 @@ guardrail_name_config_map: Dict[str, GuardrailItem] = {} include_cost_in_streaming_usage: bool = False reasoning_auto_summary: bool = False ### PROMPTS #### -from litellm.types.prompts.init_prompts import PromptSpec - prompt_name_config_map: Dict[str, PromptSpec] = {} ################## @@ -1271,206 +1267,203 @@ openai_video_generation_models = ["sora-2"] # get_llm_provider is lazy-loaded via __getattr__ # remove_index_from_tool_calls is lazy-loaded via __getattr__ -# Import KeyManagementSettings here (before utils import) because _key_management_settings -# is accessed during import time in secret_managers/main.py (via dd_tracing -> datadog -> _service_logger -> utils) -from litellm.types.secret_managers.main import KeyManagementSettings +# SDK symbols previously imported eagerly here are lazy-loaded via __getattr__ +# (_SDK_SYMBOLS_IMPORT_MAP in _lazy_imports_registry.py); mirrored under TYPE_CHECKING +# so static type checkers still see them +if TYPE_CHECKING: + _key_management_settings: KeyManagementSettings -_key_management_settings: KeyManagementSettings = KeyManagementSettings() + from .utils import client -# client must be imported immediately as it's used as a decorator at function definition time -from .utils import client + from .llms.custom_llm import CustomLLM + from .llms.anthropic.common_utils import AnthropicModelInfo + from .llms.ai21.chat.transformation import AI21ChatConfig, AI21ChatConfig as AI21Config + from .llms.deprecated_providers.palm import ( + PalmConfig, + ) # here to prevent breaking changes + from .llms.deprecated_providers.aleph_alpha import AlephAlphaConfig + from .llms.gemini.common_utils import GeminiModelInfo -# Note: Most other utils imports are lazy-loaded via __getattr__ to avoid loading utils.py -# (which imports tiktoken) at import time + from .llms.vertex_ai.vertex_embeddings.transformation import ( + VertexAITextEmbeddingConfig, + ) -from .llms.custom_llm import CustomLLM -from .llms.anthropic.common_utils import AnthropicModelInfo -from .llms.ai21.chat.transformation import AI21ChatConfig, AI21ChatConfig as AI21Config -from .llms.deprecated_providers.palm import ( - PalmConfig, -) # here to prevent breaking changes -from .llms.deprecated_providers.aleph_alpha import AlephAlphaConfig -from .llms.gemini.common_utils import GeminiModelInfo + vertexAITextEmbeddingConfig = VertexAITextEmbeddingConfig() + from .llms.bedrock.embed.amazon_titan_v2_transformation import ( + AmazonTitanV2Config, + ) + from .llms.topaz.common_utils import TopazModelInfo -from .llms.vertex_ai.vertex_embeddings.transformation import ( - VertexAITextEmbeddingConfig, -) + # OpenAIOSeriesConfig is lazy loaded - openaiOSeriesConfig will be created on first access + # OpenAIGPTConfig, OpenAIGPT5Config, etc. are lazy loaded - instances will be created on first access + from .llms.xai.common_utils import XAIModelInfo -vertexAITextEmbeddingConfig = VertexAITextEmbeddingConfig() + # PublicAI now uses JSON-based configuration (see litellm/llms/openai_like/providers.json) + # All remaining configs are now lazy loaded - see _lazy_imports_registry.py + # Import LlmProviders here (before main import) because it's imported during import time + # in multiple places including openai.py (via main import) -from .llms.bedrock.embed.amazon_titan_v2_transformation import ( - AmazonTitanV2Config, -) -from .llms.topaz.common_utils import TopazModelInfo + ## Lazy loading this is not straightforward, will leave it here for now. + from .main import * + from .compression import compress -# OpenAIOSeriesConfig is lazy loaded - openaiOSeriesConfig will be created on first access -# OpenAIGPTConfig, OpenAIGPT5Config, etc. are lazy loaded - instances will be created on first access -from .llms.xai.common_utils import XAIModelInfo + # Skills API + from .skills.main import ( + create_skill, + acreate_skill, + list_skills, + alist_skills, + get_skill, + aget_skill, + delete_skill, + adelete_skill, + ) + from .evals.main import ( + create_eval, + acreate_eval, + list_evals, + alist_evals, + get_eval, + aget_eval, + delete_eval, + adelete_eval, + cancel_eval, + acancel_eval, + create_run, + acreate_run, + list_runs, + alist_runs, + get_run, + aget_run, + delete_run, + adelete_run, + cancel_run, + acancel_run, + ) + from .integrations import * + from .llms.custom_httpx.async_client_cleanup import close_litellm_async_clients + from .exceptions import ( + AuthenticationError, + InvalidRequestError, + BadRequestError, + ImageFetchError, + NotFoundError, + PermissionDeniedError, + RateLimitError, + RateLimitErrorCategory, + RateLimitType, + ServiceUnavailableError, + BadGatewayError, + OpenAIError, + ContextWindowExceededError, + ContentPolicyViolationError, + BudgetExceededError, + APIError, + Timeout, + APIConnectionError, + UnsupportedParamsError, + APIResponseValidationError, + UnprocessableEntityError, + InternalServerError, + JSONSchemaValidationError, + LITELLM_EXCEPTION_TYPES, + MockException, + ) + from .budget_manager import BudgetManager + from .proxy.proxy_cli import run_server + from .router import Router + from .assistants.main import * + from .batches.main import * + from .images.main import * + from .videos.main import * + from .batch_completion.main import * + from .rerank_api.main import * + from .llms.anthropic.experimental_pass_through.messages.handler import * + from .responses.main import * -# PublicAI now uses JSON-based configuration (see litellm/llms/openai_like/providers.json) -# All remaining configs are now lazy loaded - see _lazy_imports_registry.py + # Interactions API is available as litellm.interactions module + # Usage: litellm.interactions.create(), litellm.interactions.get(), etc. + from . import interactions + from .interactions.agents.main import ( + acreate as acreate_agent, + create as create_agent, + alist as alist_agents, + list as list_agents, + aget as aget_agent, + get as get_agent, + adelete as adelete_agent, + delete as delete_agent, + alist_versions as alist_agent_versions, + list_versions as list_agent_versions, + ) + from .skills.main import ( + create_skill, + acreate_skill, + list_skills, + alist_skills, + get_skill, + aget_skill, + delete_skill, + adelete_skill, + ) + from .containers.main import * + from .ocr.main import * + from .rust_bridge import rust + from .rag.main import * + from .sandbox.main import * + from .search.main import * + from .realtime_api.main import ( + _arealtime, + acreate_realtime_client_secret, + acreate_realtime_transcription_session, + arealtime_calls, + ) + from .responses.main import _aresponses_websocket + from .fine_tuning.main import * + from .files.main import * + from .vector_store_files.main import ( + acreate as avector_store_file_create, + adelete as avector_store_file_delete, + alist as avector_store_file_list, + aretrieve as avector_store_file_retrieve, + aretrieve_content as avector_store_file_content, + aupdate as avector_store_file_update, + create as vector_store_file_create, + delete as vector_store_file_delete, + list as vector_store_file_list, + retrieve as vector_store_file_retrieve, + retrieve_content as vector_store_file_content, + update as vector_store_file_update, + ) + from .scheduler import * -# Import LlmProviders here (before main import) because it's imported during import time -# in multiple places including openai.py (via main import) -from litellm.types.utils import LlmProviders + ### ADAPTERS ### + import litellm.anthropic_interface as anthropic -## Lazy loading this is not straightforward, will leave it here for now. -from .main import * -from .compression import compress + ### Vector Store Registry ### -# Skills API -from .skills.main import ( - create_skill, - acreate_skill, - list_skills, - alist_skills, - get_skill, - aget_skill, - delete_skill, - adelete_skill, -) -from .evals.main import ( - create_eval, - acreate_eval, - list_evals, - alist_evals, - get_eval, - aget_eval, - delete_eval, - adelete_eval, - cancel_eval, - acancel_eval, - create_run, - acreate_run, - list_runs, - alist_runs, - get_run, - aget_run, - delete_run, - adelete_run, - cancel_run, - acancel_run, -) -from .integrations import * -from .llms.custom_httpx.async_client_cleanup import close_litellm_async_clients -from .exceptions import ( - AuthenticationError, - InvalidRequestError, - BadRequestError, - ImageFetchError, - NotFoundError, - PermissionDeniedError, - RateLimitError, - RateLimitErrorCategory, - RateLimitType, - ServiceUnavailableError, - BadGatewayError, - OpenAIError, - ContextWindowExceededError, - ContentPolicyViolationError, - BudgetExceededError, - APIError, - Timeout, - APIConnectionError, - UnsupportedParamsError, - APIResponseValidationError, - UnprocessableEntityError, - InternalServerError, - JSONSchemaValidationError, - LITELLM_EXCEPTION_TYPES, - MockException, -) -from .budget_manager import BudgetManager -from .proxy.proxy_cli import run_server -from .router import Router -from .assistants.main import * -from .batches.main import * -from .images.main import * -from .videos.main import * -from .batch_completion.main import * -from .rerank_api.main import * -from .llms.anthropic.experimental_pass_through.messages.handler import * -from .responses.main import * + ### RAG ### + from . import rag -# Interactions API is available as litellm.interactions module -# Usage: litellm.interactions.create(), litellm.interactions.get(), etc. -from . import interactions -from .interactions.agents.main import ( - acreate as acreate_agent, - create as create_agent, - alist as alist_agents, - list as list_agents, - aget as aget_agent, - get as get_agent, - adelete as adelete_agent, - delete as delete_agent, - alist_versions as alist_agent_versions, - list_versions as list_agent_versions, -) -from .skills.main import ( - create_skill, - acreate_skill, - list_skills, - alist_skills, - get_skill, - aget_skill, - delete_skill, - adelete_skill, -) -from .containers.main import * -from .ocr.main import * -from .rust_bridge import rust -from .rag.main import * -from .sandbox.main import * -from .search.main import * -from .realtime_api.main import ( - _arealtime, - acreate_realtime_client_secret, - acreate_realtime_transcription_session, - arealtime_calls, -) -from .responses.main import _aresponses_websocket -from .fine_tuning.main import * -from .files.main import * -from .vector_store_files.main import ( - acreate as avector_store_file_create, - adelete as avector_store_file_delete, - alist as avector_store_file_list, - aretrieve as avector_store_file_retrieve, - aretrieve_content as avector_store_file_content, - aupdate as avector_store_file_update, - create as vector_store_file_create, - delete as vector_store_file_delete, - list as vector_store_file_list, - retrieve as vector_store_file_retrieve, - retrieve_content as vector_store_file_content, - update as vector_store_file_update, -) -from .scheduler import * + ### CUSTOM LLMs ### + + ### CLI UTILITIES ### + from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_key + + ### PASSTHROUGH ### + from .passthrough import allm_passthrough_route, llm_passthrough_route + from .google_genai import agenerate_content ### ADAPTERS ### -from .types.adapter import AdapterItem -import litellm.anthropic_interface as anthropic - adapters: List[AdapterItem] = [] ### Vector Store Registry ### -from .vector_stores.vector_store_registry import ( - VectorStoreRegistry, - VectorStoreIndexRegistry, -) - vector_store_registry: Optional[VectorStoreRegistry] = None vector_store_index_registry: Optional[VectorStoreIndexRegistry] = None -### RAG ### -from . import rag - ### CUSTOM LLMs ### -from .types.llms.custom_llm import CustomLLMItem - custom_provider_map: List[CustomLLMItem] = [] _custom_providers: List[str] = [] # internal helper util, used to track names of custom providers disable_hf_tokenizer_download: Optional[bool] = ( @@ -1478,13 +1471,6 @@ disable_hf_tokenizer_download: Optional[bool] = ( ) global_disable_no_log_param: bool = False -### CLI UTILITIES ### -from litellm.litellm_core_utils.cli_token_utils import get_litellm_gateway_api_key - -### PASSTHROUGH ### -from .passthrough import allm_passthrough_route, llm_passthrough_route -from .google_genai import agenerate_content - ### GLOBAL CONFIG ### global_bitbucket_config: Optional[Dict[str, Any]] = None @@ -1508,10 +1494,21 @@ def set_global_gitlab_config(config: Dict[str, Any]) -> None: # Lazy loading system for heavy modules to reduce initial import time and memory usage if TYPE_CHECKING: + import httpx + from litellm.types.utils import ModelInfo as _ModelInfoType from litellm.types.utils import PriorityReservationSettings from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.caching.caching import Cache + from litellm.types.adapter import AdapterItem + from litellm.types.integrations.datadog import DatadogInitParams + from litellm.types.integrations.newrelic import NewRelicInitParams + from litellm.types.llms.custom_llm import CustomLLMItem + from litellm.types.prompts.init_prompts import PromptSpec + from litellm.vector_stores.vector_store_registry import ( + VectorStoreIndexRegistry, + VectorStoreRegistry, + ) # Type stubs for lazy-loaded configs to help mypy from .llms.bedrock.chat.converse_transformation import ( @@ -2187,16 +2184,6 @@ if TYPE_CHECKING: # Track if async client cleanup has been registered (for lazy loading) _async_client_cleanup_registered = False -# Eager loading for backwards compatibility with VCR and other HTTP recording tools -# When LITELLM_DISABLE_LAZY_LOADING is set, lazy-loaded attributes are loaded at import time -# For now, this only affects encoding (tiktoken) as it was the only reported issue -# See: https://github.com/BerriAI/litellm/issues/18659 -# This ensures encoding is initialized before VCR starts recording HTTP requests -if os.getenv("LITELLM_DISABLE_LAZY_LOADING", "").lower() in ("1", "true", "yes", "on"): - # Load encoding at import time (pre-#18070 behavior) - # This ensures encoding is initialized before VCR starts recording - from .main import encoding - def __getattr__(name: str) -> Any: """Lazy import handler with cached registry for improved performance.""" @@ -2276,6 +2263,8 @@ def __getattr__(name: str) -> Any: "openAIGPT5Config": "OpenAIGPT5Config", "nvidiaNimConfig": "NvidiaNimConfig", "nvidiaNimEmbeddingConfig": "NvidiaNimEmbeddingConfig", + "vertexAITextEmbeddingConfig": "VertexAITextEmbeddingConfig", + "_key_management_settings": "KeyManagementSettings", } if name in _config_instances: from ._lazy_imports import get_litellm_globals @@ -2393,7 +2382,30 @@ def __getattr__(name: str) -> Any: return locals()[name] + from ._lazy_imports import lazy_import_litellm_submodule + + submodule: Final = lazy_import_litellm_submodule(name) + if submodule is not None: + return submodule + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") +from ._lazy_imports import LiteLLMModule +from ._lazy_imports_registry import STAR_IMPORT_PUBLIC_NAMES + +sys.modules[__name__].__class__ = LiteLLMModule + +__all__ = list(STAR_IMPORT_PUBLIC_NAMES) # mutable-ok: star imports require __all__ to be a list of str + + # ALL_LITELLM_RESPONSE_TYPES is lazy-loaded via __getattr__ to avoid loading utils at import time + +# Eager loading for backwards compatibility with VCR and other HTTP recording tools +# When LITELLM_DISABLE_LAZY_LOADING is set, lazy-loaded attributes are loaded at import time +# For now, this only affects encoding (tiktoken) as it was the only reported issue +# See: https://github.com/BerriAI/litellm/issues/18659 +# This ensures encoding is initialized before VCR starts recording HTTP requests +# This block stays at the bottom so __getattr__ can resolve attributes main.py needs during its import +if os.getenv("LITELLM_DISABLE_LAZY_LOADING", "").lower() in ("1", "true", "yes", "on"): + from .main import encoding diff --git a/litellm/_lazy_imports.py b/litellm/_lazy_imports.py index 553aeb6680d..004297a559e 100644 --- a/litellm/_lazy_imports.py +++ b/litellm/_lazy_imports.py @@ -16,9 +16,10 @@ until they're actually needed. """ import importlib +import importlib.util import sys from collections.abc import Callable, Mapping -from types import ModuleType +from types import MappingProxyType, ModuleType from typing import TYPE_CHECKING, Any, Final, cast from typing_extensions import ReadOnly, TypedDict @@ -34,6 +35,8 @@ from ._lazy_imports_registry import ( _LITELLM_LOGGING_IMPORT_MAP, _LLM_CONFIGS_IMPORT_MAP, _LLM_PROVIDER_LOGIC_IMPORT_MAP, + _SDK_MODULE_ALIASES, + _SDK_SYMBOLS_IMPORT_MAP, _TOKEN_COUNTER_IMPORT_MAP, _TYPES_IMPORT_MAP, _TYPES_UTILS_IMPORT_MAP, @@ -78,7 +81,10 @@ def _get_utils_globals() -> dict[str, object]: This is where we cache imported attributes so we don't import them twice. When you do `litellm.utils.some_function`, it gets stored in this dictionary. """ - return sys.modules["litellm.utils"].__dict__ + cached: Final = sys.modules.get("litellm.utils") + if cached is not None: + return cached.__dict__ + return importlib.import_module("litellm.utils").__dict__ def _get_module_level_client_timeout(litellm_globals: Mapping[str, Any]) -> "float | httpx.Timeout | None": @@ -214,6 +220,10 @@ def _get_lazy_import_registry() -> dict[str, Callable[[str], object]]: _LAZY_IMPORT_REGISTRY[name] = _lazy_import_llm_provider_logic for name in UTILS_MODULE_NAMES: _LAZY_IMPORT_REGISTRY[name] = _lazy_import_utils_module + for name in _SDK_SYMBOLS_IMPORT_MAP: + _LAZY_IMPORT_REGISTRY.setdefault(name, _lazy_import_sdk_symbols) + for name in _SDK_MODULE_ALIASES: + _LAZY_IMPORT_REGISTRY.setdefault(name, _lazy_import_sdk_module_alias) return _LAZY_IMPORT_REGISTRY @@ -229,7 +239,7 @@ def _module_attribute(module: ModuleType, attr_name: str) -> object: return attribute["value"] -def _generic_lazy_import(name: str, import_map: dict[str, tuple[str, str]], category: str) -> object: +def _generic_lazy_import(name: str, import_map: Mapping[str, tuple[str, str]], category: str) -> object: """ Generic function that handles lazy importing for most attributes. @@ -350,6 +360,86 @@ def _lazy_import_llm_provider_logic(name: str) -> object: return _generic_lazy_import(name, _LLM_PROVIDER_LOGIC_IMPORT_MAP, "LLM provider logic") +def _lazy_import_sdk_symbols(name: str) -> object: + """Handler for SDK symbols previously imported eagerly at the bottom of litellm/__init__.py""" + return _generic_lazy_import(name, _SDK_SYMBOLS_IMPORT_MAP, "SDK symbols") + + +def _lazy_import_sdk_module_alias(name: str) -> object: + """Handler for litellm attributes that bind a module (e.g. litellm.anthropic)""" + _globals: Final = get_litellm_globals() + if name in _globals: + return _globals[name] + module: Final = importlib.import_module(_SDK_MODULE_ALIASES[name]) + _globals[name] = module # rebind-ok: caches the resolved module alias on the package + return module + + +_SHADOWABLE_SDK_FUNCTIONS: Final = MappingProxyType( + { + "batch_completion": ("litellm.batch_completion.main", "batch_completion"), + "ocr": ("litellm.ocr.main", "ocr"), + "responses": ("litellm.responses.main", "responses"), + "search": ("litellm.search.main", "search"), + } +) + + +def _shadowable_function_property(name: str) -> property: + """Property keeping litellm. bound to the SDK function even after the import + machinery binds the identically named litellm. subpackage onto the litellm module.""" + module_path, attr_name = _SHADOWABLE_SDK_FUNCTIONS[name] + + def _get(module: ModuleType) -> object: + stored: Final = module.__dict__.get(name) + if stored is not None and not (isinstance(stored, ModuleType) and stored.__name__ == f"litellm.{name}"): + return stored + value: Final = _module_attribute(importlib.import_module(module_path), attr_name) + module.__dict__[name] = value # rebind-ok: caches the resolved function on the litellm module + return value + + def _set(module: ModuleType, value: object) -> None: + module.__dict__[name] = value # rebind-ok: property setter must store assignments on the module + + return property(_get, _set) + + +class LiteLLMModule(ModuleType): + """Module type installed on the litellm package so function names shadowed by + same-named subpackages (litellm.responses, ...) keep resolving to the functions.""" + + batch_completion = _shadowable_function_property("batch_completion") + ocr = _shadowable_function_property("ocr") + responses = _shadowable_function_property("responses") + search = _shadowable_function_property("search") + + +def lazy_import_submodule(package: str, name: str) -> "ModuleType | None": + """Resolve . as a submodule (e.g. litellm.utils) when no other handler matches""" + if name.startswith("__") or not name.isidentifier(): + return None + qualified_name: Final = f"{package}.{name}" + try: + spec: Final = importlib.util.find_spec(qualified_name) + except ModuleNotFoundError: + return None + if spec is None: + return None + try: + module: Final = importlib.import_module(qualified_name) + except ModuleNotFoundError as exc: + if exc.name == qualified_name: + return None + raise + sys.modules[package].__dict__[name] = module # rebind-ok: caches the resolved submodule on the package + return module + + +def lazy_import_litellm_submodule(name: str) -> "ModuleType | None": + """Resolve litellm. as a submodule (e.g. litellm.utils) when no other handler matches""" + return lazy_import_submodule("litellm", name) + + def _lazy_import_utils_module(name: str) -> object: """ Handler for utils module lazy imports. diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index e9199e1ec80..b0e2fb1398c 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -5,6 +5,8 @@ This module contains all the name tuples and import maps used by the lazy import Separated from the handler functions for better organization. """ +from collections.abc import Mapping +from types import MappingProxyType from typing import Final # Cost calculator names that support lazy loading via _lazy_import_cost_calculator @@ -1479,6 +1481,1171 @@ _UTILS_MODULE_IMPORT_MAP: Final = { "LiteLLM_Params": ("litellm.types.router", "LiteLLM_Params"), } +_SDK_SYMBOLS_IMPORT_MAP: Final[Mapping[str, tuple[str, str]]] = MappingProxyType( + { + "AI21Config": ("litellm.llms.ai21.chat.transformation", "AI21ChatConfig"), + "ALL_RESPONSES_API_TOOL_PARAMS": ("litellm.assistants.main", "ALL_RESPONSES_API_TOOL_PARAMS"), + "APIConnectionError": ("litellm.exceptions", "APIConnectionError"), + "APIError": ("litellm.exceptions", "APIError"), + "APIResponseValidationError": ("litellm.exceptions", "APIResponseValidationError"), + "AZURE_OPENAI_AUDIO_PROVIDERS": ("litellm.main", "AZURE_OPENAI_AUDIO_PROVIDERS"), + "AdapterCompletionStreamWrapper": ("litellm.types.utils", "AdapterCompletionStreamWrapper"), + "AdapterItem": ("litellm.types.adapter", "AdapterItem"), + "AdaptiveRouterConfig": ("litellm.types.router", "AdaptiveRouterConfig"), + "AdaptiveRouterPreferences": ("litellm.types.router", "AdaptiveRouterPreferences"), + "AdaptiveRouterWeights": ("litellm.types.router", "AdaptiveRouterWeights"), + "AlephAlphaConfig": ("litellm.llms.deprecated_providers.aleph_alpha", "AlephAlphaConfig"), + "AlertingConfig": ("litellm.types.router", "AlertingConfig"), + "AllEmbeddingInputValues": ("litellm.assistants.main", "AllEmbeddingInputValues"), + "AllMessageValues": ("litellm.assistants.main", "AllMessageValues"), + "AllPromptValues": ("litellm.assistants.main", "AllPromptValues"), + "AllowedFailsPolicy": ("litellm.types.router", "AllowedFailsPolicy"), + "AmazonTitanV2Config": ("litellm.llms.bedrock.embed.amazon_titan_v2_transformation", "AmazonTitanV2Config"), + "Annotated": ("litellm.assistants.main", "Annotated"), + "AnthropicBatchesHandler": ("litellm.llms.anthropic.batches.handler", "AnthropicBatchesHandler"), + "AnthropicChatCompletion": ("litellm.llms.anthropic.chat.handler", "AnthropicChatCompletion"), + "AnthropicMessagesRequestUtils": ( + "litellm.llms.anthropic.experimental_pass_through.messages.utils", + "AnthropicMessagesRequestUtils", + ), + "AnthropicMessagesResponse": ( + "litellm.types.llms.anthropic_messages.anthropic_response", + "AnthropicMessagesResponse", + ), + "AnthropicMetadata": ("litellm.types.llms.anthropic_messages.anthropic_request", "AnthropicMetadata"), + "AnthropicModelInfo": ("litellm.llms.anthropic.common_utils", "AnthropicModelInfo"), + "Assistant": ("litellm.assistants.main", "Assistant"), + "AssistantDeleted": ("litellm.assistants.main", "AssistantDeleted"), + "AssistantEventHandler": ("litellm.assistants.main", "AssistantEventHandler"), + "AssistantStreamManager": ("litellm.assistants.main", "AssistantStreamManager"), + "AssistantToolParam": ("litellm.assistants.main", "AssistantToolParam"), + "AssistantsTypedDict": ("litellm.types.router", "AssistantsTypedDict"), + "AsyncAssistantEventHandler": ("litellm.assistants.main", "AsyncAssistantEventHandler"), + "AsyncAssistantStreamManager": ("litellm.assistants.main", "AsyncAssistantStreamManager"), + "AsyncCompletions": ("litellm.main", "AsyncCompletions"), + "AsyncCursorPage": ("litellm.assistants.main", "AsyncCursorPage"), + "AsyncIterator": ("litellm.llms.anthropic.experimental_pass_through.messages.handler", "AsyncIterator"), + "AsyncOpenAI": ("litellm.assistants.main", "AsyncOpenAI"), + "Attachment": ("litellm.types.llms.openai", "Attachment"), + "AttachmentTool": ("litellm.assistants.main", "AttachmentTool"), + "AuthenticationError": ("litellm.exceptions", "AuthenticationError"), + "AutoRouterCapabilityLimit": ("litellm.types.router", "AutoRouterCapabilityLimit"), + "AzureAIEmbedding": ("litellm.llms.azure_ai.embed.handler", "AzureAIEmbedding"), + "AzureAnthropicChatCompletion": ("litellm.llms.azure_ai.anthropic.handler", "AzureAnthropicChatCompletion"), + "AzureAssistantsAPI": ("litellm.llms.azure.assistants", "AzureAssistantsAPI"), + "AzureAudioTranscription": ("litellm.llms.azure.audio_transcriptions", "AzureAudioTranscription"), + "AzureBatchesAPI": ("litellm.llms.azure.batches.handler", "AzureBatchesAPI"), + "AzureChatCompletion": ("litellm.llms.azure.azure", "AzureChatCompletion"), + "AzureOpenAIFilesAPI": ("litellm.llms.azure.files.handler", "AzureOpenAIFilesAPI"), + "AzureOpenAIFineTuningAPI": ("litellm.llms.azure.fine_tuning.handler", "AzureOpenAIFineTuningAPI"), + "AzureOpenAIO1ChatCompletion": ("litellm.llms.azure.chat.o_series_handler", "AzureOpenAIO1ChatCompletion"), + "AzureTextCompletion": ("litellm.llms.azure.completion.handler", "AzureTextCompletion"), + "BATCH_GUARDRAIL_RESPONSE_FIELD": ("litellm.assistants.main", "BATCH_GUARDRAIL_RESPONSE_FIELD"), + "BadGatewayError": ("litellm.exceptions", "BadGatewayError"), + "BadRequestError": ("litellm.exceptions", "BadRequestError"), + "BaseConfig": ("litellm.llms.base_llm.chat.transformation", "BaseConfig"), + "BaseLLMAIOHTTPHandler": ("litellm.llms.custom_httpx.aiohttp_handler", "BaseLLMAIOHTTPHandler"), + "BaseLLMException": ("litellm.llms.base_llm.chat.transformation", "BaseLLMException"), + "BaseLLMHTTPHandler": ("litellm.llms.custom_httpx.llm_http_handler", "BaseLLMHTTPHandler"), + "BaseLiteLLMOpenAIResponseObject": ("litellm.types.llms.base", "BaseLiteLLMOpenAIResponseObject"), + "BaseModel": ("litellm.scheduler", "BaseModel"), + "BaseResponsesAPIConfig": ("litellm.llms.base_llm.responses.transformation", "BaseResponsesAPIConfig"), + "BaseResponsesAPIStreamingIterator": ( + "litellm.responses.streaming_iterator", + "BaseResponsesAPIStreamingIterator", + ), + "Batch": ("litellm.assistants.main", "Batch"), + "BatchGuardrailRecord": ("litellm.types.llms.openai", "BatchGuardrailRecord"), + "BatchGuardrailReport": ("litellm.types.llms.openai", "BatchGuardrailReport"), + "BatchJobStatus": ("litellm.assistants.main", "BatchJobStatus"), + "BatchRequestCounts": ("litellm.batches.main", "BatchRequestCounts"), + "BedrockBatchesHandler": ("litellm.llms.bedrock.batches.handler", "BedrockBatchesHandler"), + "BedrockConverseLLM": ("litellm.llms.bedrock.chat.converse_handler", "BedrockConverseLLM"), + "BedrockEmbedding": ("litellm.llms.bedrock.embed.embedding", "BedrockEmbedding"), + "BedrockFilesHandler": ("litellm.llms.bedrock.files.handler", "BedrockFilesHandler"), + "BedrockImageEdit": ("litellm.llms.bedrock.image_edit.handler", "BedrockImageEdit"), + "BedrockImageGeneration": ("litellm.llms.bedrock.image_generation.image_handler", "BedrockImageGeneration"), + "BedrockRerankHandler": ("litellm.llms.bedrock.rerank.handler", "BedrockRerankHandler"), + "BudgetExceededError": ("litellm.exceptions", "BudgetExceededError"), + "BudgetManager": ("litellm.budget_manager", "BudgetManager"), + "CARRY_UNMATCHED_MESSAGE_POINTS": ("litellm.responses.main", "CARRY_UNMATCHED_MESSAGE_POINTS"), + "CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS": ("litellm.files.main", "CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS"), + "CREATE_FILE_REQUESTS_PURPOSE": ("litellm.assistants.main", "CREATE_FILE_REQUESTS_PURPOSE"), + "CallTypes": ("litellm.types.utils", "CallTypes"), + "CancelBatchRequest": ("litellm.types.llms.openai", "CancelBatchRequest"), + "CharacterObject": ("litellm.types.videos.main", "CharacterObject"), + "Chat": ("litellm.main", "Chat"), + "ChatCompletionAnnotation": ("litellm.types.llms.openai", "ChatCompletionAnnotation"), + "ChatCompletionAnnotationURLCitation": ("litellm.types.llms.openai", "ChatCompletionAnnotationURLCitation"), + "ChatCompletionAssistantContentValue": ("litellm.assistants.main", "ChatCompletionAssistantContentValue"), + "ChatCompletionAssistantMessage": ("litellm.types.llms.openai", "ChatCompletionAssistantMessage"), + "ChatCompletionAssistantToolCall": ("litellm.types.llms.openai", "ChatCompletionAssistantToolCall"), + "ChatCompletionAudioDelta": ("litellm.types.llms.openai", "ChatCompletionAudioDelta"), + "ChatCompletionAudioObject": ("litellm.types.llms.openai", "ChatCompletionAudioObject"), + "ChatCompletionAudioParam": ("litellm.assistants.main", "ChatCompletionAudioParam"), + "ChatCompletionCachedContent": ("litellm.types.llms.openai", "ChatCompletionCachedContent"), + "ChatCompletionChunk": ("litellm.assistants.main", "ChatCompletionChunk"), + "ChatCompletionContentPartInputAudioParam": ( + "litellm.assistants.main", + "ChatCompletionContentPartInputAudioParam", + ), + "ChatCompletionDeltaChunk": ("litellm.types.llms.openai", "ChatCompletionDeltaChunk"), + "ChatCompletionDeveloperMessage": ("litellm.types.llms.openai", "ChatCompletionDeveloperMessage"), + "ChatCompletionDocumentObject": ("litellm.types.llms.openai", "ChatCompletionDocumentObject"), + "ChatCompletionFileObject": ("litellm.types.llms.openai", "ChatCompletionFileObject"), + "ChatCompletionFileObjectFile": ("litellm.types.llms.openai", "ChatCompletionFileObjectFile"), + "ChatCompletionFunctionMessage": ("litellm.types.llms.openai", "ChatCompletionFunctionMessage"), + "ChatCompletionImageObject": ("litellm.types.llms.openai", "ChatCompletionImageObject"), + "ChatCompletionImageUrlObject": ("litellm.types.llms.openai", "ChatCompletionImageUrlObject"), + "ChatCompletionMessageToolCall": ("litellm.types.utils", "ChatCompletionMessageToolCall"), + "ChatCompletionModality": ("litellm.assistants.main", "ChatCompletionModality"), + "ChatCompletionNamedToolChoiceParam": ("litellm.types.llms.openai", "ChatCompletionNamedToolChoiceParam"), + "ChatCompletionPredictionContentParam": ("litellm.assistants.main", "ChatCompletionPredictionContentParam"), + "ChatCompletionReasoningItem": ("litellm.types.llms.openai", "ChatCompletionReasoningItem"), + "ChatCompletionReasoningSummaryTextBlock": ( + "litellm.types.llms.openai", + "ChatCompletionReasoningSummaryTextBlock", + ), + "ChatCompletionRedactedThinkingBlock": ("litellm.types.llms.openai", "ChatCompletionRedactedThinkingBlock"), + "ChatCompletionRequest": ("litellm.types.llms.openai", "ChatCompletionRequest"), + "ChatCompletionResponseMessage": ("litellm.types.llms.openai", "ChatCompletionResponseMessage"), + "ChatCompletionSystemMessage": ("litellm.types.llms.openai", "ChatCompletionSystemMessage"), + "ChatCompletionTextObject": ("litellm.types.llms.openai", "ChatCompletionTextObject"), + "ChatCompletionThinkingBlock": ("litellm.types.llms.openai", "ChatCompletionThinkingBlock"), + "ChatCompletionToolChoiceFunctionParam": ("litellm.types.llms.openai", "ChatCompletionToolChoiceFunctionParam"), + "ChatCompletionToolChoiceObjectParam": ("litellm.types.llms.openai", "ChatCompletionToolChoiceObjectParam"), + "ChatCompletionToolChoiceStringValues": ("litellm.assistants.main", "ChatCompletionToolChoiceStringValues"), + "ChatCompletionToolChoiceValues": ("litellm.assistants.main", "ChatCompletionToolChoiceValues"), + "ChatCompletionToolMessage": ("litellm.types.llms.openai", "ChatCompletionToolMessage"), + "ChatCompletionToolParam": ("litellm.types.llms.openai", "ChatCompletionToolParam"), + "ChatCompletionToolParamFunctionChunk": ("litellm.types.llms.openai", "ChatCompletionToolParamFunctionChunk"), + "ChatCompletionToolReferenceObject": ("litellm.types.llms.openai", "ChatCompletionToolReferenceObject"), + "ChatCompletionUsageBlock": ("litellm.types.llms.openai", "ChatCompletionUsageBlock"), + "ChatCompletionUserMessage": ("litellm.types.llms.openai", "ChatCompletionUserMessage"), + "ChatCompletionVideoObject": ("litellm.types.llms.openai", "ChatCompletionVideoObject"), + "ChatCompletionVideoUrlObject": ("litellm.types.llms.openai", "ChatCompletionVideoUrlObject"), + "Choices": ("litellm.types.utils", "Choices"), + "ChunkProcessor": ("litellm.litellm_core_utils.streaming_chunk_builder_utils", "ChunkProcessor"), + "CitationsObject": ("litellm.types.llms.openai", "CitationsObject"), + "ClassVar": ("litellm.files.main", "ClassVar"), + "ClassifierPlugin": ("litellm.types.router", "ClassifierPlugin"), + "CodeInterpreterToolParam": ("litellm.types.llms.openai", "CodeInterpreterToolParam"), + "CodestralTextCompletion": ("litellm.llms.codestral.completion.handler", "CodestralTextCompletion"), + "CompletionRequest": ("litellm.types.completion", "CompletionRequest"), + "CompletionTimeout": ("litellm.litellm_core_utils.completion_timeout", "CompletionTimeout"), + "CompletionTokensDetails": ("litellm.main", "CompletionTokensDetails"), + "Completions": ("litellm.main", "Completions"), + "ComputerToolParam": ("litellm.types.llms.openai", "ComputerToolParam"), + "ConfigDict": ("litellm.files.main", "ConfigDict"), + "ConfigurableClientsideParamsCustomAuth": ("litellm.types.router", "ConfigurableClientsideParamsCustomAuth"), + "ConsumedRequestTagsStamp": ("litellm.types.router", "ConsumedRequestTagsStamp"), + "ContentPartAddedEvent": ("litellm.types.llms.openai", "ContentPartAddedEvent"), + "ContentPartDoneEvent": ("litellm.types.llms.openai", "ContentPartDoneEvent"), + "ContentPartDonePartOutputText": ("litellm.types.llms.openai", "ContentPartDonePartOutputText"), + "ContentPartDonePartReasoningText": ("litellm.types.llms.openai", "ContentPartDonePartReasoningText"), + "ContentPartDonePartRefusal": ("litellm.types.llms.openai", "ContentPartDonePartRefusal"), + "ContentPolicyViolationError": ("litellm.exceptions", "ContentPolicyViolationError"), + "ContextManagementEntry": ("litellm.types.llms.openai", "ContextManagementEntry"), + "ContextWindowExceededError": ("litellm.exceptions", "ContextWindowExceededError"), + "Coroutine": ("litellm.files.main", "Coroutine"), + "CreateBatchRequest": ("litellm.types.llms.openai", "CreateBatchRequest"), + "CreateFileRequest": ("litellm.types.llms.openai", "CreateFileRequest"), + "CreateVideoRequest": ("litellm.types.llms.openai", "CreateVideoRequest"), + "CredentialLiteLLMParams": ("litellm.types.router", "CredentialLiteLLMParams"), + "CustomLLM": ("litellm.llms.custom_llm", "CustomLLM"), + "CustomLLMItem": ("litellm.types.llms.custom_llm", "CustomLLMItem"), + "CustomPricingLiteLLMParams": ("litellm.types.utils", "CustomPricingLiteLLMParams"), + "CustomRoutingStrategyBase": ("litellm.types.router", "CustomRoutingStrategyBase"), + "CustomToolCallOutputItem": ("litellm.types.responses.main", "CustomToolCallOutputItem"), + "DEFAULT_IMAGE_ENDPOINT_MODEL": ("litellm.images.main", "DEFAULT_IMAGE_ENDPOINT_MODEL"), + "DEFAULT_IN_MEMORY_TTL": ("litellm.scheduler", "DEFAULT_IN_MEMORY_TTL"), + "DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT": ( + "litellm.main", + "DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT", + ), + "DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT": ("litellm.main", "DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT"), + "DEFAULT_POLLING_INTERVAL": ("litellm.scheduler", "DEFAULT_POLLING_INTERVAL"), + "DEFAULT_REQUEST_TIMEOUT": ("litellm.videos.main", "DEFAULT_REQUEST_TIMEOUT"), + "DEFAULT_VIDEO_ENDPOINT_MODEL": ("litellm.videos.main", "DEFAULT_VIDEO_ENDPOINT_MODEL"), + "DatabricksEmbeddingHandler": ("litellm.llms.databricks.embed.handler", "DatabricksEmbeddingHandler"), + "DatadogInitParams": ("litellm.types.integrations.datadog", "DatadogInitParams"), + "DecodedResponseId": ("litellm.types.responses.main", "DecodedResponseId"), + "DeleteResponseResult": ("litellm.types.responses.main", "DeleteResponseResult"), + "Deployment": ("litellm.types.router", "Deployment"), + "DeploymentTypedDict": ("litellm.types.router", "DeploymentTypedDict"), + "Discriminator": ("litellm.assistants.main", "Discriminator"), + "DocumentObject": ("litellm.types.llms.openai", "DocumentObject"), + "EmbeddingCreateParams": ("litellm.assistants.main", "EmbeddingCreateParams"), + "EmbeddingInput": ("litellm.assistants.main", "EmbeddingInput"), + "EmbeddingRequest": ("litellm.types.embedding", "EmbeddingRequest"), + "Enum": ("litellm.assistants.main", "Enum"), + "ErrorEvent": ("litellm.types.llms.openai", "ErrorEvent"), + "ErrorEventError": ("litellm.types.llms.openai", "ErrorEventError"), + "FIRST_COMPLETED": ("litellm.batch_completion.main", "FIRST_COMPLETED"), + "FORWARDED_KWARGS_KEYS": ("litellm.main", "FORWARDED_KWARGS_KEYS"), + "FallbackAccessCheck": ("litellm.types.router", "FallbackAccessCheck"), + "Field": ("litellm.files.main", "Field"), + "FileContent": ("litellm.videos.main", "FileContent"), + "FileContentProvider": ("litellm.files.main", "FileContentProvider"), + "FileContentRequest": ("litellm.types.llms.openai", "FileContentRequest"), + "FileContentStreamingResponse": ("litellm.files.streaming", "FileContentStreamingResponse"), + "FileContentStreamingResult": ("litellm.files.types", "FileContentStreamingResult"), + "FileCreateProvider": ("litellm.files.main", "FileCreateProvider"), + "FileDeleteProvider": ("litellm.files.main", "FileDeleteProvider"), + "FileDeleted": ("litellm.files.main", "FileDeleted"), + "FileExpiresAfter": ("litellm.types.llms.openai", "FileExpiresAfter"), + "FileListPage": ("litellm.types.llms.openai", "FileListPage"), + "FileListProvider": ("litellm.files.main", "FileListProvider"), + "FileObject": ("litellm.files.main", "FileObject"), + "FileRetrieveProvider": ("litellm.files.main", "FileRetrieveProvider"), + "FileSearchCallCompletedEvent": ("litellm.types.llms.openai", "FileSearchCallCompletedEvent"), + "FileSearchCallInProgressEvent": ("litellm.types.llms.openai", "FileSearchCallInProgressEvent"), + "FileSearchCallSearchingEvent": ("litellm.types.llms.openai", "FileSearchCallSearchingEvent"), + "FileSearchTool": ("litellm.types.llms.openai", "FileSearchTool"), + "FileSearchToolParam": ("litellm.types.llms.openai", "FileSearchToolParam"), + "FileTypes": ("litellm.files.main", "FileTypes"), + "FineTuningConfig": ("litellm.types.router", "FineTuningConfig"), + "FineTuningJob": ("litellm.assistants.main", "FineTuningJob"), + "FineTuningJobCreate": ("litellm.types.llms.openai", "FineTuningJobCreate"), + "FlowItem": ("litellm.scheduler", "FlowItem"), + "Function": ("litellm.types.llms.openai", "Function"), + "FunctionCallArgumentsDeltaEvent": ("litellm.types.llms.openai", "FunctionCallArgumentsDeltaEvent"), + "FunctionCallArgumentsDoneEvent": ("litellm.types.llms.openai", "FunctionCallArgumentsDoneEvent"), + "GeminiModelInfo": ("litellm.llms.gemini.common_utils", "GeminiModelInfo"), + "GenAIHubOrchestration": ("litellm.llms.sap.chat.handler", "GenAIHubOrchestration"), + "Generator": ("litellm.responses.main", "Generator"), + "Generic": ("litellm.files.main", "Generic"), + "GenericBudgetWindowDetails": ("litellm.types.router", "GenericBudgetWindowDetails"), + "GenericChatCompletionMessage": ("litellm.types.llms.openai", "GenericChatCompletionMessage"), + "GenericEvent": ("litellm.types.llms.openai", "GenericEvent"), + "GenericLiteLLMParams": ("litellm.types.router", "GenericLiteLLMParams"), + "GenericResponseOutputItem": ("litellm.types.responses.main", "GenericResponseOutputItem"), + "GenericResponseOutputItemContentAnnotation": ( + "litellm.types.responses.main", + "GenericResponseOutputItemContentAnnotation", + ), + "GoogleBatchEmbeddings": ( + "litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_handler", + "GoogleBatchEmbeddings", + ), + "GroqChatCompletion": ("litellm.llms.groq.chat.handler", "GroqChatCompletion"), + "GuardrailLiteLLMParams": ("litellm.types.router", "GuardrailLiteLLMParams"), + "GuardrailTypedDict": ("litellm.types.router", "GuardrailTypedDict"), + "HiddenParams": ("litellm.types.llms.base", "HiddenParams"), + "HttpxBinaryResponseContent": ("litellm.types.llms.openai", "HttpxBinaryResponseContent"), + "HuggingFaceEmbedding": ("litellm.llms.huggingface.embedding.handler", "HuggingFaceEmbedding"), + "Hyperparameters": ("litellm.types.llms.openai", "Hyperparameters"), + "IBMWatsonXMixin": ("litellm.llms.watsonx.common_utils", "IBMWatsonXMixin"), + "IO": ("litellm.assistants.main", "IO"), + "IOBase": ("litellm.ocr.main", "IOBase"), + "ImageEditOptionalRequestParams": ("litellm.types.images.main", "ImageEditOptionalRequestParams"), + "ImageFetchError": ("litellm.exceptions", "ImageFetchError"), + "ImageFileObject": ("litellm.types.llms.openai", "ImageFileObject"), + "ImageGenerationPartialImageEvent": ("litellm.types.llms.openai", "ImageGenerationPartialImageEvent"), + "ImageGenerationRequestQuality": ("litellm.types.llms.openai", "ImageGenerationRequestQuality"), + "ImageURLListItem": ("litellm.types.llms.openai", "ImageURLListItem"), + "ImageURLObject": ("litellm.types.llms.openai", "ImageURLObject"), + "IncompleteDetails": ("litellm.assistants.main", "IncompleteDetails"), + "InputTokensDetails": ("litellm.types.llms.openai", "InputTokensDetails"), + "InternalServerError": ("litellm.exceptions", "InternalServerError"), + "InvalidRequestError": ("litellm.exceptions", "InvalidRequestError"), + "Iterable": ("litellm.responses.main", "Iterable"), + "Iterator": ("litellm.llms.anthropic.experimental_pass_through.messages.handler", "Iterator"), + "JSONProviderRegistry": ("litellm.llms.openai_like.json_loader", "JSONProviderRegistry"), + "JSONSchemaValidationError": ("litellm.exceptions", "JSONSchemaValidationError"), + "KeyManagementSettings": ("litellm.types.secret_managers.main", "KeyManagementSettings"), + "LIST_BATCHES_SUPPORTED_PROVIDERS": ("litellm.batches.main", "LIST_BATCHES_SUPPORTED_PROVIDERS"), + "LITELLM_EXCEPTION_TYPES": ("litellm.exceptions", "LITELLM_EXCEPTION_TYPES"), + "LITELLM_IMAGE_VARIATION_PROVIDERS": ("litellm.types.utils", "LITELLM_IMAGE_VARIATION_PROVIDERS"), + "ListBatchRequest": ("litellm.types.llms.openai", "ListBatchRequest"), + "ListBatchesSupportedProvider": ("litellm.batches.main", "ListBatchesSupportedProvider"), + "LiteLLM": ("litellm.main", "LiteLLM"), + "LiteLLMBatch": ("litellm.types.utils", "LiteLLMBatch"), + "LiteLLMBatchCreateRequest": ("litellm.types.llms.openai", "LiteLLMBatchCreateRequest"), + "LiteLLMCompletionTransformationHandler": ( + "litellm.responses.litellm_completion_transformation.handler", + "LiteLLMCompletionTransformationHandler", + ), + "LiteLLMFineTuningJob": ("litellm.types.utils", "LiteLLMFineTuningJob"), + "LiteLLMFineTuningJobCreate": ("litellm.types.llms.openai", "LiteLLMFineTuningJobCreate"), + "LiteLLMLoggingObj": ("litellm.files.main", "LiteLLMLoggingObj"), + "LiteLLMMessagesToCompletionTransformationHandler": ( + "litellm.llms.anthropic.experimental_pass_through.adapters.handler", + "LiteLLMMessagesToCompletionTransformationHandler", + ), + "LiteLLMMessagesToResponsesAPIHandler": ( + "litellm.llms.anthropic.experimental_pass_through.responses_adapters.handler", + "LiteLLMMessagesToResponsesAPIHandler", + ), + "LiteLLMParamsTypedDict": ("litellm.types.router", "LiteLLMParamsTypedDict"), + "LiteLLMResponsesTransformationHandler": ( + "litellm.completion_extras.litellm_responses_transformation.transformation", + "LiteLLMResponsesTransformationHandler", + ), + "LiteLLMUnknownProvider": ("litellm.exceptions", "LiteLLMUnknownProvider"), + "LiteLLM_RouterFileObject": ("litellm.types.router", "LiteLLM_RouterFileObject"), + "LlmProviders": ("litellm.types.utils", "LlmProviders"), + "MCPCallArgumentsDeltaEvent": ("litellm.types.llms.openai", "MCPCallArgumentsDeltaEvent"), + "MCPCallArgumentsDoneEvent": ("litellm.types.llms.openai", "MCPCallArgumentsDoneEvent"), + "MCPCallCompletedEvent": ("litellm.types.llms.openai", "MCPCallCompletedEvent"), + "MCPCallFailedEvent": ("litellm.types.llms.openai", "MCPCallFailedEvent"), + "MCPCallInProgressEvent": ("litellm.types.llms.openai", "MCPCallInProgressEvent"), + "MCPListToolsCompletedEvent": ("litellm.types.llms.openai", "MCPListToolsCompletedEvent"), + "MCPListToolsFailedEvent": ("litellm.types.llms.openai", "MCPListToolsFailedEvent"), + "MCPListToolsInProgressEvent": ("litellm.types.llms.openai", "MCPListToolsInProgressEvent"), + "MCPTool": ("litellm.responses.main", "MCPTool"), + "MOCK_RESPONSE_TYPE": ("litellm.main", "MOCK_RESPONSE_TYPE"), + "Mapping": ("litellm.files.main", "Mapping"), + "MappingProxyType": ("litellm.main", "MappingProxyType"), + "Message": ("litellm.types.utils", "Message"), + "MessageContent": ("litellm.assistants.main", "MessageContent"), + "MessageContentImageFileObject": ("litellm.types.llms.openai", "MessageContentImageFileObject"), + "MessageContentImageURLObject": ("litellm.types.llms.openai", "MessageContentImageURLObject"), + "MessageContentTextObject": ("litellm.types.llms.openai", "MessageContentTextObject"), + "MessageData": ("litellm.types.llms.openai", "MessageData"), + "MirroredPricingParams": ("litellm.types.utils", "MirroredPricingParams"), + "MockException": ("litellm.exceptions", "MockException"), + "MockRouterTestingParams": ("litellm.types.router", "MockRouterTestingParams"), + "ModelConfig": ("litellm.types.router", "ModelConfig"), + "ModelGroupInfo": ("litellm.types.router", "ModelGroupInfo"), + "ModelGroupSettings": ("litellm.types.router", "ModelGroupSettings"), + "ModelInfo": ("litellm.types.router", "ModelInfo"), + "NOT_GIVEN": ("litellm.types.llms.openai", "NOT_GIVEN"), + "NewRelicInitParams": ("litellm.types.integrations.newrelic", "NewRelicInitParams"), + "NonNegativeInt": ("litellm.assistants.main", "NonNegativeInt"), + "NotFoundError": ("litellm.exceptions", "NotFoundError"), + "NotGiven": ("litellm.types.llms.openai", "NotGiven"), + "NotRequired": ("litellm.assistants.main", "NotRequired"), + "NvidiaRivaAudioTranscription": ( + "litellm.llms.nvidia_riva.audio_transcription.handler", + "NvidiaRivaAudioTranscription", + ), + "NvidiaRivaAudioTranscriptionConfig": ( + "litellm.llms.nvidia_riva.audio_transcription.transformation", + "NvidiaRivaAudioTranscriptionConfig", + ), + "OCRResponse": ("litellm.llms.base_llm.ocr.transformation", "OCRResponse"), + "OCR_REQUEST_FORMAT_PARAM": ("litellm.ocr.main", "OCR_REQUEST_FORMAT_PARAM"), + "OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS": ( + "litellm.files.main", + "OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS", + ), + "OPTIONAL_KWARGS_KEYS": ("litellm.main", "OPTIONAL_KWARGS_KEYS"), + "Omit": ("litellm.assistants.main", "Omit"), + "OpenAI": ("litellm.assistants.main", "OpenAI"), + "OpenAIAssistantsAPI": ("litellm.llms.openai.openai", "OpenAIAssistantsAPI"), + "OpenAIAudioTranscription": ("litellm.llms.openai.transcriptions.handler", "OpenAIAudioTranscription"), + "OpenAIAudioTranscriptionOptionalParams": ("litellm.assistants.main", "OpenAIAudioTranscriptionOptionalParams"), + "OpenAIBatchResponse": ("litellm.types.llms.openai", "OpenAIBatchResponse"), + "OpenAIBatchResult": ("litellm.types.llms.openai", "OpenAIBatchResult"), + "OpenAIBatchesAPI": ("litellm.llms.openai.openai", "OpenAIBatchesAPI"), + "OpenAIChatCompletion": ("litellm.llms.openai.openai", "OpenAIChatCompletion"), + "OpenAIChatCompletionAssistantMessage": ("litellm.types.llms.openai", "OpenAIChatCompletionAssistantMessage"), + "OpenAIChatCompletionChoices": ("litellm.types.llms.openai", "OpenAIChatCompletionChoices"), + "OpenAIChatCompletionChunk": ("litellm.types.llms.openai", "OpenAIChatCompletionChunk"), + "OpenAIChatCompletionDeveloperMessage": ("litellm.types.llms.openai", "OpenAIChatCompletionDeveloperMessage"), + "OpenAIChatCompletionFinishReason": ("litellm.assistants.main", "OpenAIChatCompletionFinishReason"), + "OpenAIChatCompletionLogprobs": ("litellm.types.llms.openai", "OpenAIChatCompletionLogprobs"), + "OpenAIChatCompletionLogprobsContent": ("litellm.types.llms.openai", "OpenAIChatCompletionLogprobsContent"), + "OpenAIChatCompletionLogprobsContentTopLogprobs": ( + "litellm.types.llms.openai", + "OpenAIChatCompletionLogprobsContentTopLogprobs", + ), + "OpenAIChatCompletionResponse": ("litellm.types.llms.openai", "OpenAIChatCompletionResponse"), + "OpenAIChatCompletionSystemMessage": ("litellm.types.llms.openai", "OpenAIChatCompletionSystemMessage"), + "OpenAIChatCompletionTextObject": ("litellm.types.llms.openai", "OpenAIChatCompletionTextObject"), + "OpenAIChatCompletionToolParam": ("litellm.types.llms.openai", "OpenAIChatCompletionToolParam"), + "OpenAIChatCompletionUserMessage": ("litellm.types.llms.openai", "OpenAIChatCompletionUserMessage"), + "OpenAICreateFileRequestOptionalParams": ("litellm.assistants.main", "OpenAICreateFileRequestOptionalParams"), + "OpenAICreateThreadParamsMessage": ("litellm.assistants.main", "OpenAICreateThreadParamsMessage"), + "OpenAICreateThreadParamsToolResources": ("litellm.types.llms.openai", "OpenAICreateThreadParamsToolResources"), + "OpenAIEmbedding": ("litellm.assistants.main", "OpenAIEmbedding"), + "OpenAIError": ("litellm.exceptions", "OpenAIError"), + "OpenAIErrorBody": ("litellm.types.llms.openai", "OpenAIErrorBody"), + "OpenAIFileObject": ("litellm.types.llms.openai", "OpenAIFileObject"), + "OpenAIFilesAPI": ("litellm.llms.openai.openai", "OpenAIFilesAPI"), + "OpenAIFilesPurpose": ("litellm.assistants.main", "OpenAIFilesPurpose"), + "OpenAIFineTuningAPI": ("litellm.llms.openai.fine_tuning.handler", "OpenAIFineTuningAPI"), + "OpenAIImageEditOptionalParams": ("litellm.assistants.main", "OpenAIImageEditOptionalParams"), + "OpenAIImageGenerationOptionalParams": ("litellm.assistants.main", "OpenAIImageGenerationOptionalParams"), + "OpenAIImageVariationOptionalParams": ("litellm.assistants.main", "OpenAIImageVariationOptionalParams"), + "OpenAIImageVariationsHandler": ( + "litellm.llms.openai.image_variations.handler", + "OpenAIImageVariationsHandler", + ), + "OpenAILikeChatHandler": ("litellm.llms.openai_like.chat.handler", "OpenAILikeChatHandler"), + "OpenAILikeEmbeddingHandler": ("litellm.llms.openai_like.embedding.handler", "OpenAILikeEmbeddingHandler"), + "OpenAILikeResponsesConfig": ( + "litellm.llms.openai_like.responses.transformation", + "OpenAILikeResponsesConfig", + ), + "OpenAIMcpServerTool": ("litellm.types.llms.openai", "OpenAIMcpServerTool"), + "OpenAIMessage": ("litellm.assistants.main", "OpenAIMessage"), + "OpenAIMessageContent": ("litellm.assistants.main", "OpenAIMessageContent"), + "OpenAIMessageContentListBlock": ("litellm.assistants.main", "OpenAIMessageContentListBlock"), + "OpenAIModerationResponse": ("litellm.types.llms.openai", "OpenAIModerationResponse"), + "OpenAIModerationResult": ("litellm.types.llms.openai", "OpenAIModerationResult"), + "OpenAIRealtimeContentPartDone": ("litellm.types.llms.openai", "OpenAIRealtimeContentPartDone"), + "OpenAIRealtimeConversationCreated": ("litellm.types.llms.openai", "OpenAIRealtimeConversationCreated"), + "OpenAIRealtimeConversationItemAdded": ("litellm.types.llms.openai", "OpenAIRealtimeConversationItemAdded"), + "OpenAIRealtimeConversationItemCreated": ("litellm.types.llms.openai", "OpenAIRealtimeConversationItemCreated"), + "OpenAIRealtimeConversationItemDone": ("litellm.types.llms.openai", "OpenAIRealtimeConversationItemDone"), + "OpenAIRealtimeConversationObject": ("litellm.types.llms.openai", "OpenAIRealtimeConversationObject"), + "OpenAIRealtimeDoneEvent": ("litellm.types.llms.openai", "OpenAIRealtimeDoneEvent"), + "OpenAIRealtimeEventTypes": ("litellm.types.llms.openai", "OpenAIRealtimeEventTypes"), + "OpenAIRealtimeEvents": ("litellm.assistants.main", "OpenAIRealtimeEvents"), + "OpenAIRealtimeFunctionCallArgumentsDone": ( + "litellm.types.llms.openai", + "OpenAIRealtimeFunctionCallArgumentsDone", + ), + "OpenAIRealtimeInputAudioBufferSpeechEvent": ( + "litellm.types.llms.openai", + "OpenAIRealtimeInputAudioBufferSpeechEvent", + ), + "OpenAIRealtimeInputAudioTranscriptionCompleted": ( + "litellm.types.llms.openai", + "OpenAIRealtimeInputAudioTranscriptionCompleted", + ), + "OpenAIRealtimeInputAudioTranscriptionDelta": ( + "litellm.types.llms.openai", + "OpenAIRealtimeInputAudioTranscriptionDelta", + ), + "OpenAIRealtimeOutputItemDone": ("litellm.types.llms.openai", "OpenAIRealtimeOutputItemDone"), + "OpenAIRealtimeResponseAudioDone": ("litellm.types.llms.openai", "OpenAIRealtimeResponseAudioDone"), + "OpenAIRealtimeResponseContentPart": ("litellm.types.llms.openai", "OpenAIRealtimeResponseContentPart"), + "OpenAIRealtimeResponseContentPartAdded": ( + "litellm.types.llms.openai", + "OpenAIRealtimeResponseContentPartAdded", + ), + "OpenAIRealtimeResponseDelta": ("litellm.types.llms.openai", "OpenAIRealtimeResponseDelta"), + "OpenAIRealtimeResponseDoneObject": ("litellm.types.llms.openai", "OpenAIRealtimeResponseDoneObject"), + "OpenAIRealtimeResponseTextDone": ("litellm.types.llms.openai", "OpenAIRealtimeResponseTextDone"), + "OpenAIRealtimeResponseUsage": ("litellm.types.llms.openai", "OpenAIRealtimeResponseUsage"), + "OpenAIRealtimeStreamList": ("litellm.assistants.main", "OpenAIRealtimeStreamList"), + "OpenAIRealtimeStreamResponseBaseObject": ( + "litellm.types.llms.openai", + "OpenAIRealtimeStreamResponseBaseObject", + ), + "OpenAIRealtimeStreamResponseOutputItem": ( + "litellm.types.llms.openai", + "OpenAIRealtimeStreamResponseOutputItem", + ), + "OpenAIRealtimeStreamResponseOutputItemAdded": ( + "litellm.types.llms.openai", + "OpenAIRealtimeStreamResponseOutputItemAdded", + ), + "OpenAIRealtimeStreamResponseOutputItemContent": ( + "litellm.types.llms.openai", + "OpenAIRealtimeStreamResponseOutputItemContent", + ), + "OpenAIRealtimeStreamSession": ("litellm.types.llms.openai", "OpenAIRealtimeStreamSession"), + "OpenAIRealtimeStreamSessionEvents": ("litellm.types.llms.openai", "OpenAIRealtimeStreamSessionEvents"), + "OpenAIRealtimeTurnDetection": ("litellm.types.llms.openai", "OpenAIRealtimeTurnDetection"), + "OpenAIRealtimeUsageTokenDetails": ("litellm.types.llms.openai", "OpenAIRealtimeUsageTokenDetails"), + "OpenAITextCompletion": ("litellm.llms.openai.completion.handler", "OpenAITextCompletion"), + "OpenAITextCompletionUserMessage": ("litellm.types.llms.openai", "OpenAITextCompletionUserMessage"), + "OpenAIVideoObject": ("litellm.types.llms.openai", "OpenAIVideoObject"), + "OpenAIWebSearchOptions": ("litellm.types.llms.openai", "OpenAIWebSearchOptions"), + "OpenAIWebSearchUserLocation": ("litellm.types.llms.openai", "OpenAIWebSearchUserLocation"), + "OpenAIWebSearchUserLocationApproximate": ( + "litellm.types.llms.openai", + "OpenAIWebSearchUserLocationApproximate", + ), + "OptionalPreCallChecks": ("litellm.files.main", "OptionalPreCallChecks"), + "OutputCodeInterpreterCall": ("litellm.types.responses.main", "OutputCodeInterpreterCall"), + "OutputCodeInterpreterCallLog": ("litellm.types.responses.main", "OutputCodeInterpreterCallLog"), + "OutputFunctionToolCall": ("litellm.types.responses.main", "OutputFunctionToolCall"), + "OutputImageGenerationCall": ("litellm.types.responses.main", "OutputImageGenerationCall"), + "OutputItemAddedEvent": ("litellm.types.llms.openai", "OutputItemAddedEvent"), + "OutputItemDoneEvent": ("litellm.types.llms.openai", "OutputItemDoneEvent"), + "OutputText": ("litellm.types.responses.main", "OutputText"), + "OutputTextAnnotationAddedEvent": ("litellm.types.llms.openai", "OutputTextAnnotationAddedEvent"), + "OutputTextDeltaEvent": ("litellm.types.llms.openai", "OutputTextDeltaEvent"), + "OutputTextDoneEvent": ("litellm.types.llms.openai", "OutputTextDoneEvent"), + "OutputTokensDetails": ("litellm.types.llms.openai", "OutputTokensDetails"), + "PART_UNION_TYPES": ("litellm.assistants.main", "PART_UNION_TYPES"), + "PalmConfig": ("litellm.llms.deprecated_providers.palm", "PalmConfig"), + "PathLike": ("litellm.assistants.main", "PathLike"), + "PermissionDeniedError": ("litellm.exceptions", "PermissionDeniedError"), + "Phase": ("litellm.responses.main", "Phase"), + "PreRoutingHookResponse": ("litellm.types.router", "PreRoutingHookResponse"), + "PreRoutingStrategy": ("litellm.types.router", "PreRoutingStrategy"), + "PredibaseChatCompletion": ("litellm.llms.predibase.chat.handler", "PredibaseChatCompletion"), + "PrivateAttr": ("litellm.responses.main", "PrivateAttr"), + "PromptCacheBreakpoint": ("litellm.types.llms.openai", "PromptCacheBreakpoint"), + "PromptCacheOptions": ("litellm.types.llms.openai", "PromptCacheOptions"), + "PromptObject": ("litellm.types.llms.openai", "PromptObject"), + "PromptSpec": ("litellm.types.prompts.init_prompts", "PromptSpec"), + "PromptTokensDetails": ("litellm.main", "PromptTokensDetails"), + "Protocol": ("litellm.files.main", "Protocol"), + "ProviderConfigManager": ("litellm.utils", "ProviderConfigManager"), + "ProviderSpecificHeader": ("litellm.types.utils", "ProviderSpecificHeader"), + "ProviderSpecificHeaderUtils": ( + "litellm.litellm_core_utils.get_provider_specific_headers", + "ProviderSpecificHeaderUtils", + ), + "REASONING_EFFORT": ("litellm.assistants.main", "REASONING_EFFORT"), + "RateLimitError": ("litellm.exceptions", "RateLimitError"), + "RateLimitErrorCategory": ("litellm.exceptions", "RateLimitErrorCategory"), + "RateLimitType": ("litellm.exceptions", "RateLimitType"), + "RawRequestTypedDict": ("litellm.types.utils", "RawRequestTypedDict"), + "ReadOnly": ("litellm.files.main", "ReadOnly"), + "Reasoning": ("litellm.responses.main", "Reasoning"), + "ReasoningSummaryPartDoneEvent": ("litellm.types.llms.openai", "ReasoningSummaryPartDoneEvent"), + "ReasoningSummaryTextDeltaEvent": ("litellm.types.llms.openai", "ReasoningSummaryTextDeltaEvent"), + "ReasoningSummaryTextDoneEvent": ("litellm.types.llms.openai", "ReasoningSummaryTextDoneEvent"), + "RefusalDeltaEvent": ("litellm.types.llms.openai", "RefusalDeltaEvent"), + "RefusalDoneEvent": ("litellm.types.llms.openai", "RefusalDoneEvent"), + "RequestType": ("litellm.types.router", "RequestType"), + "Required": ("litellm.files.main", "Required"), + "Response": ("litellm.assistants.main", "Response"), + "ResponseAPIUsage": ("litellm.types.llms.openai", "ResponseAPIUsage"), + "ResponseCompletedEvent": ("litellm.types.llms.openai", "ResponseCompletedEvent"), + "ResponseCreatedEvent": ("litellm.types.llms.openai", "ResponseCreatedEvent"), + "ResponseFailedEvent": ("litellm.types.llms.openai", "ResponseFailedEvent"), + "ResponseFunctionToolCall": ("litellm.responses.main", "ResponseFunctionToolCall"), + "ResponseInProgressEvent": ("litellm.types.llms.openai", "ResponseInProgressEvent"), + "ResponseIncludable": ("litellm.responses.main", "ResponseIncludable"), + "ResponseIncompleteEvent": ("litellm.types.llms.openai", "ResponseIncompleteEvent"), + "ResponseInputParam": ("litellm.responses.main", "ResponseInputParam"), + "ResponseOutputItem": ("litellm.assistants.main", "ResponseOutputItem"), + "ResponsePartAddedEvent": ("litellm.types.llms.openai", "ResponsePartAddedEvent"), + "ResponseText": ("litellm.responses.main", "ResponseText"), + "ResponsesAPIOptionalRequestParams": ("litellm.types.llms.openai", "ResponsesAPIOptionalRequestParams"), + "ResponsesAPIRequestParams": ("litellm.types.llms.openai", "ResponsesAPIRequestParams"), + "ResponsesAPIRequestUtils": ("litellm.responses.utils", "ResponsesAPIRequestUtils"), + "ResponsesAPIResponse": ("litellm.types.llms.openai", "ResponsesAPIResponse"), + "ResponsesAPIStatus": ("litellm.assistants.main", "ResponsesAPIStatus"), + "ResponsesAPIStreamEvents": ("litellm.types.llms.openai", "ResponsesAPIStreamEvents"), + "ResponsesAPIStreamOptions": ("litellm.types.llms.openai", "ResponsesAPIStreamOptions"), + "ResponsesAPIStreamingResponse": ("litellm.assistants.main", "ResponsesAPIStreamingResponse"), + "ResponsesToolUsage": ("litellm.types.llms.openai", "ResponsesToolUsage"), + "RetrieveBatchRequest": ("litellm.types.llms.openai", "RetrieveBatchRequest"), + "RetryPolicy": ("litellm.types.router", "RetryPolicy"), + "Router": ("litellm.router", "Router"), + "RouterCacheEnum": ("litellm.types.router", "RouterCacheEnum"), + "RouterConfig": ("litellm.types.router", "RouterConfig"), + "RouterErrors": ("litellm.types.router", "RouterErrors"), + "RouterGeneralSettings": ("litellm.types.router", "RouterGeneralSettings"), + "RouterModelGroupAliasItem": ("litellm.types.router", "RouterModelGroupAliasItem"), + "RouterRateLimitError": ("litellm.types.router", "RouterRateLimitError"), + "RouterRateLimitErrorBasic": ("litellm.types.router", "RouterRateLimitErrorBasic"), + "RoutingContext": ("litellm.types.router", "RoutingContext"), + "RoutingGroup": ("litellm.types.router", "RoutingGroup"), + "RoutingPlugin": ("litellm.types.router", "RoutingPlugin"), + "RoutingStrategy": ("litellm.types.router", "RoutingStrategy"), + "Run": ("litellm.assistants.main", "Run"), + "SPECIAL_MODEL_INFO_PARAMS": ("litellm.files.main", "SPECIAL_MODEL_INFO_PARAMS"), + "SagemakerChatHandler": ("litellm.llms.sagemaker.chat.handler", "SagemakerChatHandler"), + "SagemakerLLM": ("litellm.llms.sagemaker.completion.handler", "SagemakerLLM"), + "Scheduler": ("litellm.scheduler", "Scheduler"), + "SchedulerCacheKeys": ("litellm.scheduler", "SchedulerCacheKeys"), + "SearchProvider": ("litellm.files.main", "SearchProvider"), + "SearchResponse": ("litellm.llms.base_llm.search.transformation", "SearchResponse"), + "SearchToolInfoTypedDict": ("litellm.types.router", "SearchToolInfoTypedDict"), + "SearchToolLiteLLMParams": ("litellm.types.router", "SearchToolLiteLLMParams"), + "SearchToolTypedDict": ("litellm.types.router", "SearchToolTypedDict"), + "SerializerFunctionWrapHandler": ("litellm.assistants.main", "SerializerFunctionWrapHandler"), + "ServiceUnavailableError": ("litellm.exceptions", "ServiceUnavailableError"), + "ShellToolParam": ("litellm.types.llms.openai", "ShellToolParam"), + "SlackAlerting": ("litellm.integrations.SlackAlerting.slack_alerting", "SlackAlerting"), + "StandardLoggingRoutingDecision": ("litellm.types.utils", "StandardLoggingRoutingDecision"), + "StreamingChoices": ("litellm.types.utils", "StreamingChoices"), + "SyncCursorPage": ("litellm.assistants.main", "SyncCursorPage"), + "TaggedPreRoutingStrategy": ("litellm.types.router", "TaggedPreRoutingStrategy"), + "TextChoices": ("litellm.types.utils", "TextChoices"), + "TextCompletionStreamWrapper": ("litellm.utils", "TextCompletionStreamWrapper"), + "Thread": ("litellm.types.llms.openai", "Thread"), + "ThreadPoolExecutor": ("litellm.batch_completion.main", "ThreadPoolExecutor"), + "Timeout": ("litellm.exceptions", "Timeout"), + "TogetherAIRerank": ("litellm.llms.together_ai.rerank.handler", "TogetherAIRerank"), + "Tool": ("litellm.assistants.main", "Tool"), + "ToolChoice": ("litellm.responses.main", "ToolChoice"), + "ToolMessageContentPart": ("litellm.assistants.main", "ToolMessageContentPart"), + "ToolParam": ("litellm.responses.main", "ToolParam"), + "ToolResourcesCodeInterpreter": ("litellm.types.llms.openai", "ToolResourcesCodeInterpreter"), + "ToolResourcesFileSearch": ("litellm.types.llms.openai", "ToolResourcesFileSearch"), + "ToolResourcesFileSearchVectorStore": ("litellm.types.llms.openai", "ToolResourcesFileSearchVectorStore"), + "TopazModelInfo": ("litellm.llms.topaz.common_utils", "TopazModelInfo"), + "TypeAlias": ("litellm.assistants.main", "TypeAlias"), + "TypeVar": ("litellm.files.main", "TypeVar"), + "TypedDict": ("litellm.files.main", "TypedDict"), + "UnprocessableEntityError": ("litellm.exceptions", "UnprocessableEntityError"), + "UnsupportedParamsError": ("litellm.exceptions", "UnsupportedParamsError"), + "UpdateRouterConfig": ("litellm.types.router", "UpdateRouterConfig"), + "Usage": ("litellm.types.utils", "Usage"), + "VALID_LITELLM_ENVIRONMENTS": ("litellm.files.main", "VALID_LITELLM_ENVIRONMENTS"), + "ValidAssistantMessageContentTypes": ("litellm.assistants.main", "ValidAssistantMessageContentTypes"), + "ValidAssistantMessageContentTypesLiteral": ( + "litellm.assistants.main", + "ValidAssistantMessageContentTypesLiteral", + ), + "ValidChatCompletionMessageContentTypes": ("litellm.assistants.main", "ValidChatCompletionMessageContentTypes"), + "ValidChatCompletionMessageContentTypesLiteral": ( + "litellm.assistants.main", + "ValidChatCompletionMessageContentTypesLiteral", + ), + "ValidUserMessageContentTypes": ("litellm.assistants.main", "ValidUserMessageContentTypes"), + "ValidUserMessageContentTypesLiteral": ("litellm.assistants.main", "ValidUserMessageContentTypesLiteral"), + "VectorStoreIndexRegistry": ("litellm.vector_stores.vector_store_registry", "VectorStoreIndexRegistry"), + "VectorStoreRegistry": ("litellm.vector_stores.vector_store_registry", "VectorStoreRegistry"), + "VertexAIBatchPrediction": ("litellm.llms.vertex_ai.batches.handler", "VertexAIBatchPrediction"), + "VertexAIFilesHandler": ("litellm.llms.vertex_ai.files.handler", "VertexAIFilesHandler"), + "VertexAIGemmaModels": ("litellm.llms.vertex_ai.vertex_gemma_models.main", "VertexAIGemmaModels"), + "VertexAIModelGardenModels": ("litellm.llms.vertex_ai.vertex_model_garden.main", "VertexAIModelGardenModels"), + "VertexAIModelRoute": ("litellm.llms.vertex_ai.common_utils", "VertexAIModelRoute"), + "VertexAIPartnerModels": ("litellm.llms.vertex_ai.vertex_ai_partner_models.main", "VertexAIPartnerModels"), + "VertexAITextEmbeddingConfig": ( + "litellm.llms.vertex_ai.vertex_embeddings.transformation", + "VertexAITextEmbeddingConfig", + ), + "VertexEmbedding": ("litellm.llms.vertex_ai.vertex_embeddings.embedding_handler", "VertexEmbedding"), + "VertexFineTuningAPI": ("litellm.llms.vertex_ai.fine_tuning.handler", "VertexFineTuningAPI"), + "VertexImageGeneration": ( + "litellm.llms.vertex_ai.image_generation.image_generation_handler", + "VertexImageGeneration", + ), + "VertexLLM": ("litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini", "VertexLLM"), + "VertexMultimodalEmbedding": ( + "litellm.llms.vertex_ai.multimodal_embeddings.embedding_handler", + "VertexMultimodalEmbedding", + ), + "VideoCreateOptionalRequestParams": ("litellm.types.videos.main", "VideoCreateOptionalRequestParams"), + "VideoGenerationRequestUtils": ("litellm.videos.utils", "VideoGenerationRequestUtils"), + "VideoObject": ("litellm.types.videos.main", "VideoObject"), + "WatsonXChatHandler": ("litellm.llms.watsonx.chat.handler", "WatsonXChatHandler"), + "WebSearchCallCompletedEvent": ("litellm.types.llms.openai", "WebSearchCallCompletedEvent"), + "WebSearchCallInProgressEvent": ("litellm.types.llms.openai", "WebSearchCallInProgressEvent"), + "WebSearchCallSearchingEvent": ("litellm.types.llms.openai", "WebSearchCallSearchingEvent"), + "WebSearchOptions": ("litellm.types.llms.openai", "WebSearchOptions"), + "WebSearchOptionsUserLocation": ("litellm.types.llms.openai", "WebSearchOptionsUserLocation"), + "WebSearchOptionsUserLocationApproximate": ( + "litellm.types.llms.openai", + "WebSearchOptionsUserLocationApproximate", + ), + "WebSearchToolUsage": ("litellm.types.llms.openai", "WebSearchToolUsage"), + "XAIModelInfo": ("litellm.llms.xai.common_utils", "XAIModelInfo"), + "_arealtime": ("litellm.realtime_api.main", "_arealtime"), + "_aresponses_websocket": ("litellm.responses.main", "_aresponses_websocket"), + "a_add_message": ("litellm.assistants.main", "a_add_message"), + "aadapter_completion": ("litellm.main", "aadapter_completion"), + "aadapter_generate_content": ("litellm.main", "aadapter_generate_content"), + "acancel_batch": ("litellm.batches.main", "acancel_batch"), + "acancel_fine_tuning_job": ("litellm.fine_tuning.main", "acancel_fine_tuning_job"), + "acancel_responses": ("litellm.responses.main", "acancel_responses"), + "acode_interpreter_tool": ("litellm.sandbox.main", "acode_interpreter_tool"), + "acompact_responses": ("litellm.responses.main", "acompact_responses"), + "acompletion": ("litellm.main", "acompletion"), + "acompletion_with_retries": ("litellm.main", "acompletion_with_retries"), + "acount_tokens": ("litellm.main", "acount_tokens"), + "acreate_agent": ("litellm.interactions.agents.main", "acreate"), + "acreate_assistants": ("litellm.assistants.main", "acreate_assistants"), + "acreate_batch": ("litellm.batches.main", "acreate_batch"), + "acreate_container": ("litellm.containers.main", "acreate_container"), + "acreate_file": ("litellm.files.main", "acreate_file"), + "acreate_fine_tuning_job": ("litellm.fine_tuning.main", "acreate_fine_tuning_job"), + "acreate_realtime_client_secret": ("litellm.realtime_api.main", "acreate_realtime_client_secret"), + "acreate_realtime_transcription_session": ( + "litellm.realtime_api.main", + "acreate_realtime_transcription_session", + ), + "acreate_sandbox": ("litellm.sandbox.main", "acreate_sandbox"), + "acreate_skill": ("litellm.skills.main", "acreate_skill"), + "acreate_thread": ("litellm.assistants.main", "acreate_thread"), + "adapter_completion": ("litellm.main", "adapter_completion"), + "add_message": ("litellm.assistants.main", "add_message"), + "add_provider_specific_params_to_optional_params": ( + "litellm.utils", + "add_provider_specific_params_to_optional_params", + ), + "add_system_prompt_to_messages": ( + "litellm.litellm_core_utils.prompt_templates.common_utils", + "add_system_prompt_to_messages", + ), + "add_trusted_model_credentials_to_litellm_params": ( + "litellm.litellm_core_utils.get_litellm_params", + "add_trusted_model_credentials_to_litellm_params", + ), + "adelete_agent": ("litellm.interactions.agents.main", "adelete"), + "adelete_assistant": ("litellm.assistants.main", "adelete_assistant"), + "adelete_container": ("litellm.containers.main", "adelete_container"), + "adelete_responses": ("litellm.responses.main", "adelete_responses"), + "adelete_sandbox": ("litellm.sandbox.main", "adelete_sandbox"), + "adelete_skill": ("litellm.skills.main", "adelete_skill"), + "aembedding": ("litellm.main", "aembedding"), + "afile_content": ("litellm.files.main", "afile_content"), + "afile_delete": ("litellm.files.main", "afile_delete"), + "afile_list": ("litellm.files.main", "afile_list"), + "afile_retrieve": ("litellm.files.main", "afile_retrieve"), + "agenerate_content": ("litellm.google_genai.main", "agenerate_content"), + "aget_agent": ("litellm.interactions.agents.main", "aget"), + "aget_assistants": ("litellm.assistants.main", "aget_assistants"), + "aget_messages": ("litellm.assistants.main", "aget_messages"), + "aget_responses": ("litellm.responses.main", "aget_responses"), + "aget_skill": ("litellm.skills.main", "aget_skill"), + "aget_thread": ("litellm.assistants.main", "aget_thread"), + "ahealth_check": ("litellm.main", "ahealth_check"), + "aimage_edit": ("litellm.images.main", "aimage_edit"), + "aimage_generation": ("litellm.images.main", "aimage_generation"), + "aimage_variation": ("litellm.images.main", "aimage_variation"), + "aingest": ("litellm.rag.main", "aingest"), + "alist_agent_versions": ("litellm.interactions.agents.main", "alist_versions"), + "alist_agents": ("litellm.interactions.agents.main", "alist"), + "alist_batches": ("litellm.batches.main", "alist_batches"), + "alist_container_files": ("litellm.containers.main", "alist_container_files"), + "alist_containers": ("litellm.containers.main", "alist_containers"), + "alist_fine_tuning_jobs": ("litellm.fine_tuning.main", "alist_fine_tuning_jobs"), + "alist_input_items": ("litellm.responses.main", "alist_input_items"), + "alist_skills": ("litellm.skills.main", "alist_skills"), + "allm_passthrough_route": ("litellm.passthrough.main", "allm_passthrough_route"), + "amoderation": ("litellm.main", "amoderation"), + "anthropic_batches_instance": ("litellm.batches.main", "anthropic_batches_instance"), + "anthropic_chat_completions": ("litellm.main", "anthropic_chat_completions"), + "anthropic_messages": ( + "litellm.llms.anthropic.experimental_pass_through.messages.handler", + "anthropic_messages", + ), + "anthropic_messages_handler": ( + "litellm.llms.anthropic.experimental_pass_through.messages.handler", + "anthropic_messages_handler", + ), + "aocr": ("litellm.ocr.main", "aocr"), + "aquery": ("litellm.rag.main", "aquery"), + "arealtime_calls": ("litellm.realtime_api.main", "arealtime_calls"), + "arerank": ("litellm.rerank_api.main", "arerank"), + "aresponses": ("litellm.responses.main", "aresponses"), + "aresponses_api_with_mcp": ("litellm.responses.main", "aresponses_api_with_mcp"), + "aresponses_with_retries": ("litellm.main", "aresponses_with_retries"), + "aretrieve_batch": ("litellm.batches.main", "aretrieve_batch"), + "aretrieve_container": ("litellm.containers.main", "aretrieve_container"), + "aretrieve_fine_tuning_job": ("litellm.fine_tuning.main", "aretrieve_fine_tuning_job"), + "arun_code": ("litellm.sandbox.main", "arun_code"), + "arun_thread": ("litellm.assistants.main", "arun_thread"), + "arun_thread_stream": ("litellm.assistants.main", "arun_thread_stream"), + "asearch": ("litellm.search.main", "asearch"), + "aspeech": ("litellm.main", "aspeech"), + "async_completion_with_fallbacks": ( + "litellm.litellm_core_utils.fallback_utils", + "async_completion_with_fallbacks", + ), + "async_mock_completion_streaming_obj": ("litellm.utils", "async_mock_completion_streaming_obj"), + "atext_completion": ("litellm.main", "atext_completion"), + "atranscription": ("litellm.main", "atranscription"), + "aupload_container_file": ("litellm.containers.main", "aupload_container_file"), + "avector_store_file_content": ("litellm.vector_store_files.main", "aretrieve_content"), + "avector_store_file_create": ("litellm.vector_store_files.main", "acreate"), + "avector_store_file_delete": ("litellm.vector_store_files.main", "adelete"), + "avector_store_file_list": ("litellm.vector_store_files.main", "alist"), + "avector_store_file_retrieve": ("litellm.vector_store_files.main", "aretrieve"), + "avector_store_file_update": ("litellm.vector_store_files.main", "aupdate"), + "avideo_content": ("litellm.videos.main", "avideo_content"), + "avideo_create_character": ("litellm.videos.main", "avideo_create_character"), + "avideo_edit": ("litellm.videos.main", "avideo_edit"), + "avideo_extension": ("litellm.videos.main", "avideo_extension"), + "avideo_generation": ("litellm.videos.main", "avideo_generation"), + "avideo_get_character": ("litellm.videos.main", "avideo_get_character"), + "avideo_list": ("litellm.videos.main", "avideo_list"), + "avideo_remix": ("litellm.videos.main", "avideo_remix"), + "avideo_status": ("litellm.videos.main", "avideo_status"), + "azure_ai_embedding": ("litellm.main", "azure_ai_embedding"), + "azure_anthropic_chat_completions": ("litellm.main", "azure_anthropic_chat_completions"), + "azure_assistants_api": ("litellm.assistants.main", "azure_assistants_api"), + "azure_audio_transcriptions": ("litellm.main", "azure_audio_transcriptions"), + "azure_batches_instance": ("litellm.batches.main", "azure_batches_instance"), + "azure_chat_completions": ("litellm.images.main", "azure_chat_completions"), + "azure_files_instance": ("litellm.files.main", "azure_files_instance"), + "azure_fine_tuning_apis_instance": ("litellm.fine_tuning.main", "azure_fine_tuning_apis_instance"), + "azure_o1_chat_completions": ("litellm.main", "azure_o1_chat_completions"), + "azure_text_completions": ("litellm.main", "azure_text_completions"), + "base_llm_aiohttp_handler": ("litellm.images.main", "base_llm_aiohttp_handler"), + "base_llm_http_handler": ("litellm.files.main", "base_llm_http_handler"), + "batch_completion": ("litellm.batch_completion.main", "batch_completion"), + "batch_completion_models": ("litellm.batch_completion.main", "batch_completion_models"), + "batch_completion_models_all_responses": ( + "litellm.batch_completion.main", + "batch_completion_models_all_responses", + ), + "bedrock_converse_chat_completion": ("litellm.main", "bedrock_converse_chat_completion"), + "bedrock_embedding": ("litellm.main", "bedrock_embedding"), + "bedrock_files_instance": ("litellm.files.main", "bedrock_files_instance"), + "bedrock_image_edit": ("litellm.images.main", "bedrock_image_edit"), + "bedrock_image_generation": ("litellm.images.main", "bedrock_image_generation"), + "bedrock_rerank": ("litellm.rerank_api.main", "bedrock_rerank"), + "bfl_image_edit": ("litellm.llms.black_forest_labs.image_edit.handler", "bfl_image_edit"), + "bfl_image_generation": ("litellm.llms.black_forest_labs.image_generation.handler", "bfl_image_generation"), + "build_code_interpreter_log_outputs": ("litellm.types.responses.main", "build_code_interpreter_log_outputs"), + "bytez_transformation": ("litellm.main", "bytez_transformation"), + "calculate_request_duration": ("litellm.litellm_core_utils.audio_utils.utils", "calculate_request_duration"), + "cancel_batch": ("litellm.batches.main", "cancel_batch"), + "cancel_fine_tuning_job": ("litellm.fine_tuning.main", "cancel_fine_tuning_job"), + "cancel_responses": ("litellm.responses.main", "cancel_responses"), + "cast": ("litellm.files.main", "cast"), + "client": ("litellm.utils", "client"), + "close_litellm_async_clients": ( + "litellm.llms.custom_httpx.async_client_cleanup", + "close_litellm_async_clients", + ), + "codestral_text_completions": ("litellm.main", "codestral_text_completions"), + "compact_responses": ("litellm.responses.main", "compact_responses"), + "completion": ("litellm.main", "completion"), + "completion_with_fallbacks": ("litellm.litellm_core_utils.fallback_utils", "completion_with_fallbacks"), + "completion_with_retries": ("litellm.main", "completion_with_retries"), + "compress": ("litellm.compression.compress", "compress"), + "config_completion": ("litellm.main", "config_completion"), + "contextmanager": ("litellm.responses.main", "contextmanager"), + "convert_file_document_to_url_document": ("litellm.ocr.main", "convert_file_document_to_url_document"), + "convert_model_response_to_streaming": ( + "litellm.llms.base_llm.base_model_iterator", + "convert_model_response_to_streaming", + ), + "create_agent": ("litellm.interactions.agents.main", "create"), + "create_assistants": ("litellm.assistants.main", "create_assistants"), + "create_batch": ("litellm.batches.main", "create_batch"), + "create_container": ("litellm.containers.main", "create_container"), + "create_file": ("litellm.files.main", "create_file"), + "create_fine_tuning_job": ("litellm.fine_tuning.main", "create_fine_tuning_job"), + "create_skill": ("litellm.skills.main", "create_skill"), + "create_thread": ("litellm.assistants.main", "create_thread"), + "custom_chat_llm_router": ("litellm.llms.custom_llm", "custom_chat_llm_router"), + "custom_prompt": ("litellm.litellm_core_utils.prompt_templates.factory", "custom_prompt"), + "databricks_embedding": ("litellm.main", "databricks_embedding"), + "dataclass": ("litellm.files.main", "dataclass"), + "decode_video_id_with_provider": ("litellm.types.videos.utils", "decode_video_id_with_provider"), + "declared_authenticating_provider": ( + "litellm.litellm_core_utils.get_llm_provider_logic", + "declared_authenticating_provider", + ), + "deepcopy": ("litellm.main", "deepcopy"), + "delete_agent": ("litellm.interactions.agents.main", "delete"), + "delete_assistant": ("litellm.assistants.main", "delete_assistant"), + "delete_container": ("litellm.containers.main", "delete_container"), + "delete_responses": ("litellm.responses.main", "delete_responses"), + "delete_skill": ("litellm.skills.main", "delete_skill"), + "disable_cache": ("litellm.caching.caching", "disable_cache"), + "embedding": ("litellm.main", "embedding"), + "enable_cache": ("litellm.caching.caching", "enable_cache"), + "field_serializer": ("litellm.assistants.main", "field_serializer"), + "field_validator": ("litellm.files.main", "field_validator"), + "file_content": ("litellm.files.main", "file_content"), + "file_content_streaming": ("litellm.files.main", "file_content_streaming"), + "file_delete": ("litellm.files.main", "file_delete"), + "file_list": ("litellm.files.main", "file_list"), + "file_retrieve": ("litellm.files.main", "file_retrieve"), + "filter_out_litellm_params": ("litellm.utils", "filter_out_litellm_params"), + "flatten_form_field_values": ("litellm.litellm_core_utils.llm_request_utils", "flatten_form_field_values"), + "flatten_unencrypted_web_search_results_in_anthropic_messages": ( + "litellm.llms.anthropic.common_utils", + "flatten_unencrypted_web_search_results_in_anthropic_messages", + ), + "function_call_prompt": ("litellm.litellm_core_utils.prompt_templates.factory", "function_call_prompt"), + "gdc_transformation": ("litellm.main", "gdc_transformation"), + "get_agent": ("litellm.interactions.agents.main", "get"), + "get_api_key_from_env": ("litellm.llms.gemini.common_utils", "get_api_key_from_env"), + "get_assistants": ("litellm.assistants.main", "get_assistants"), + "get_audio_file_for_health_check": ( + "litellm.litellm_core_utils.audio_utils.utils", + "get_audio_file_for_health_check", + ), + "get_azure_credentials": ("litellm.llms.azure.common_utils", "get_azure_credentials"), + "get_completion_messages": ( + "litellm.litellm_core_utils.prompt_templates.common_utils", + "get_completion_messages", + ), + "get_configured_request_timeout": ( + "litellm.litellm_core_utils.request_timeout_resolver", + "get_configured_request_timeout", + ), + "get_content_from_model_response": ( + "litellm.litellm_core_utils.prompt_templates.common_utils", + "get_content_from_model_response", + ), + "get_litellm_gateway_api_key": ("litellm.litellm_core_utils.cli_token_utils", "get_litellm_gateway_api_key"), + "get_messages": ("litellm.assistants.main", "get_messages"), + "get_messages_interceptors": ( + "litellm.llms.anthropic.experimental_pass_through.messages.interceptors", + "get_messages_interceptors", + ), + "get_mime_type": ("litellm.ocr.main", "get_mime_type"), + "get_non_default_completion_params": ("litellm.utils", "get_non_default_completion_params"), + "get_non_default_transcription_params": ("litellm.utils", "get_non_default_transcription_params"), + "get_openai_credentials": ("litellm.llms.openai.common_utils", "get_openai_credentials"), + "get_optional_params_add_message": ("litellm.assistants.utils", "get_optional_params_add_message"), + "get_optional_params_embeddings": ("litellm.utils", "get_optional_params_embeddings"), + "get_optional_params_image_gen": ("litellm.utils", "get_optional_params_image_gen"), + "get_optional_params_transcription": ("litellm.utils", "get_optional_params_transcription"), + "get_optional_rerank_params": ("litellm.rerank_api.rerank_utils", "get_optional_rerank_params"), + "get_requester_metadata": ("litellm.utils", "get_requester_metadata"), + "get_responses": ("litellm.responses.main", "get_responses"), + "get_secret": ("litellm.secret_managers.main", "get_secret"), + "get_secret_bool": ("litellm.secret_managers.main", "get_secret_bool"), + "get_secret_str": ("litellm.secret_managers.main", "get_secret_str"), + "get_skill": ("litellm.skills.main", "get_skill"), + "get_standard_openai_params": ("litellm.utils", "get_standard_openai_params"), + "get_thread": ("litellm.assistants.main", "get_thread"), + "get_type_hints": ("litellm.files.main", "get_type_hints"), + "get_vertex_ai_model_route": ("litellm.llms.vertex_ai.common_utils", "get_vertex_ai_model_route"), + "google_batch_embeddings": ("litellm.main", "google_batch_embeddings"), + "groq_chat_completions": ("litellm.main", "groq_chat_completions"), + "heroku_transformation": ("litellm.main", "heroku_transformation"), + "huggingface_embed": ("litellm.main", "huggingface_embed"), + "image_edit": ("litellm.images.main", "image_edit"), + "image_generation": ("litellm.images.main", "image_generation"), + "image_variation": ("litellm.images.main", "image_variation"), + "infer_openai_data_residency": ("litellm.llms.openai.data_residency", "infer_openai_data_residency"), + "ingest": ("litellm.rag.main", "ingest"), + "is_azure_document_intelligence_model": ( + "litellm.llms.azure_ai.ocr.common_utils", + "is_azure_document_intelligence_model", + ), + "is_reasoning_auto_summary_enabled": ( + "litellm.llms.anthropic.experimental_pass_through.utils", + "is_reasoning_auto_summary_enabled", + ), + "lemonade_transformation": ("litellm.main", "lemonade_transformation"), + "list_agent_versions": ("litellm.interactions.agents.main", "list_versions"), + "list_agents": ("litellm.interactions.agents.main", "list"), + "list_batches": ("litellm.batches.main", "list_batches"), + "list_container_files": ("litellm.containers.main", "list_container_files"), + "list_containers": ("litellm.containers.main", "list_containers"), + "list_fine_tuning_jobs": ("litellm.fine_tuning.main", "list_fine_tuning_jobs"), + "list_input_items": ("litellm.responses.main", "list_input_items"), + "list_skills": ("litellm.skills.main", "list_skills"), + "litellm_completion_transformation_handler": ( + "litellm.responses.main", + "litellm_completion_transformation_handler", + ), + "llm_http_handler": ("litellm.videos.main", "llm_http_handler"), + "llm_passthrough_route": ("litellm.passthrough.main", "llm_passthrough_route"), + "map_system_message_pt": ("litellm.litellm_core_utils.prompt_templates.factory", "map_system_message_pt"), + "maybe_run_chat_completion_agentic_loop": ( + "litellm.litellm_core_utils.chat_completion_agentic_loop", + "maybe_run_chat_completion_agentic_loop", + ), + "mock_completion": ("litellm.main", "mock_completion"), + "mock_completion_streaming_obj": ("litellm.utils", "mock_completion_streaming_obj"), + "mock_embedding": ("litellm.litellm_core_utils.mock_functions", "mock_embedding"), + "mock_image_generation": ("litellm.litellm_core_utils.mock_functions", "mock_image_generation"), + "mock_response": ("litellm.llms.anthropic.experimental_pass_through.messages.utils", "mock_response"), + "mock_responses_api_response": ("litellm.responses.main", "mock_responses_api_response"), + "model_serializer": ("litellm.assistants.main", "model_serializer"), + "model_validator": ("litellm.files.main", "model_validator"), + "moderation": ("litellm.main", "moderation"), + "nlp_cloud_chat_completion": ("litellm.main", "nlp_cloud_chat_completion"), + "nvidia_riva_audio_transcriptions": ("litellm.main", "nvidia_riva_audio_transcriptions"), + "oci_transformation": ("litellm.main", "oci_transformation"), + "ocr": ("litellm.ocr.main", "ocr"), + "ollama_pt": ("litellm.litellm_core_utils.prompt_templates.factory", "ollama_pt"), + "openai_assistants_api": ("litellm.assistants.main", "openai_assistants_api"), + "openai_audio_transcriptions": ("litellm.main", "openai_audio_transcriptions"), + "openai_batches_instance": ("litellm.batches.main", "openai_batches_instance"), + "openai_chat_completions": ("litellm.images.main", "openai_chat_completions"), + "openai_files_instance": ("litellm.files.main", "openai_files_instance"), + "openai_fine_tuning_apis_instance": ("litellm.fine_tuning.main", "openai_fine_tuning_apis_instance"), + "openai_image_variations": ("litellm.images.main", "openai_image_variations"), + "openai_like_chat_completion": ("litellm.main", "openai_like_chat_completion"), + "openai_like_embedding": ("litellm.main", "openai_like_embedding"), + "openai_text_completions": ("litellm.main", "openai_text_completions"), + "override": ("litellm.assistants.main", "override"), + "ovhcloud_transformation": ("litellm.main", "ovhcloud_transformation"), + "parse_ocr_request_format": ("litellm.llms.base_llm.ocr.transformation", "parse_ocr_request_format"), + "partial": ("litellm.files.main", "partial"), + "peek_reasoning_summary_aliases": ("litellm.utils", "peek_reasoning_summary_aliases"), + "pre_process_non_default_params": ("litellm.utils", "pre_process_non_default_params"), + "predibase_chat_completions": ("litellm.main", "predibase_chat_completions"), + "print_verbose": ("litellm.main", "print_verbose"), + "prompt_factory": ("litellm.litellm_core_utils.prompt_templates.factory", "prompt_factory"), + "query": ("litellm.rag.main", "query"), + "read_config_args": ("litellm.utils", "read_config_args"), + "replicate_chat_completion": ("litellm.main", "replicate_chat_completion"), + "rerank": ("litellm.rerank_api.main", "rerank"), + "responses": ("litellm.responses.main", "responses"), + "responses_api_bridge_check": ("litellm.main", "responses_api_bridge_check"), + "responses_with_retries": ("litellm.main", "responses_with_retries"), + "retrieve_batch": ("litellm.batches.main", "retrieve_batch"), + "retrieve_container": ("litellm.containers.main", "retrieve_container"), + "retrieve_fine_tuning_job": ("litellm.fine_tuning.main", "retrieve_fine_tuning_job"), + "run_async_function": ("litellm.litellm_core_utils.asyncify", "run_async_function"), + "run_server": ("litellm.proxy.proxy_cli", "run_server"), + "run_thread": ("litellm.assistants.main", "run_thread"), + "run_thread_stream": ("litellm.assistants.main", "run_thread_stream"), + "runtime_checkable": ("litellm.files.main", "runtime_checkable"), + "rust": ("litellm.rust_bridge", "rust"), + "safe_deep_copy": ("litellm.litellm_core_utils.core_helpers", "safe_deep_copy"), + "sagemaker_chat_completion": ("litellm.main", "sagemaker_chat_completion"), + "sagemaker_llm": ("litellm.main", "sagemaker_llm"), + "sanitize_tool_use_ids_in_anthropic_messages": ( + "litellm.llms.anthropic.common_utils", + "sanitize_tool_use_ids_in_anthropic_messages", + ), + "sap_gen_ai_hub_chat_completions": ("litellm.main", "sap_gen_ai_hub_chat_completions"), + "sap_gen_ai_hub_emb": ("litellm.main", "sap_gen_ai_hub_emb"), + "search": ("litellm.search.main", "search"), + "should_run_mock_completion": ("litellm.utils", "should_run_mock_completion"), + "speech": ("litellm.main", "speech"), + "stream_chunk_builder": ("litellm.main", "stream_chunk_builder"), + "stream_chunk_builder_text_completion": ("litellm.main", "stream_chunk_builder_text_completion"), + "stringify_json_tool_call_content": ( + "litellm.litellm_core_utils.prompt_templates.factory", + "stringify_json_tool_call_content", + ), + "strip_empty_content_blocks_from_anthropic_messages": ( + "litellm.llms.anthropic.common_utils", + "strip_empty_content_blocks_from_anthropic_messages", + ), + "strip_reasoning_summary_aliases_from_optional_params": ( + "litellm.utils", + "strip_reasoning_summary_aliases_from_optional_params", + ), + "supports_httpx_timeout": ("litellm.utils", "supports_httpx_timeout"), + "text_completion": ("litellm.main", "text_completion"), + "together_rerank": ("litellm.rerank_api.main", "together_rerank"), + "tracer": ("litellm.litellm_core_utils.dd_tracing", "tracer"), + "transcription": ("litellm.main", "transcription"), + "updateDeployment": ("litellm.types.router", "updateDeployment"), + "updateLiteLLMParams": ("litellm.types.router", "updateLiteLLMParams"), + "update_cache": ("litellm.caching.caching", "update_cache"), + "update_messages_with_model_file_ids": ( + "litellm.litellm_core_utils.prompt_templates.common_utils", + "update_messages_with_model_file_ids", + ), + "update_responses_input_with_model_file_ids": ( + "litellm.litellm_core_utils.prompt_templates.common_utils", + "update_responses_input_with_model_file_ids", + ), + "update_responses_tools_with_model_file_ids": ( + "litellm.litellm_core_utils.prompt_templates.common_utils", + "update_responses_tools_with_model_file_ids", + ), + "upload_container_file": ("litellm.containers.main", "upload_container_file"), + "urlsplit": ("litellm.main", "urlsplit"), + "validate_and_fix_openai_messages": ("litellm.utils", "validate_and_fix_openai_messages"), + "validate_and_fix_openai_tools": ("litellm.utils", "validate_and_fix_openai_tools"), + "validate_and_fix_thinking_param": ("litellm.utils", "validate_and_fix_thinking_param"), + "validate_anthropic_api_metadata": ( + "litellm.llms.anthropic.experimental_pass_through.messages.handler", + "validate_anthropic_api_metadata", + ), + "validate_chat_completion_tool_choice": ("litellm.utils", "validate_chat_completion_tool_choice"), + "validate_openai_optional_params": ("litellm.utils", "validate_openai_optional_params"), + "vector_store_file_content": ("litellm.vector_store_files.main", "retrieve_content"), + "vector_store_file_create": ("litellm.vector_store_files.main", "create"), + "vector_store_file_delete": ("litellm.vector_store_files.main", "delete"), + "vector_store_file_list": ("litellm.vector_store_files.main", "list"), + "vector_store_file_retrieve": ("litellm.vector_store_files.main", "retrieve"), + "vector_store_file_update": ("litellm.vector_store_files.main", "update"), + "vertex_ai_batches_instance": ("litellm.batches.main", "vertex_ai_batches_instance"), + "vertex_ai_files_instance": ("litellm.files.main", "vertex_ai_files_instance"), + "vertex_chat_completion": ("litellm.main", "vertex_chat_completion"), + "vertex_embedding": ("litellm.main", "vertex_embedding"), + "vertex_fine_tuning_apis_instance": ("litellm.fine_tuning.main", "vertex_fine_tuning_apis_instance"), + "vertex_gemma_chat_completion": ("litellm.main", "vertex_gemma_chat_completion"), + "vertex_image_generation": ("litellm.main", "vertex_image_generation"), + "vertex_model_garden_chat_completion": ("litellm.main", "vertex_model_garden_chat_completion"), + "vertex_multimodal_embedding": ("litellm.main", "vertex_multimodal_embedding"), + "vertex_partner_models_chat_completion": ("litellm.main", "vertex_partner_models_chat_completion"), + "video_content": ("litellm.videos.main", "video_content"), + "video_create_character": ("litellm.videos.main", "video_create_character"), + "video_edit": ("litellm.videos.main", "video_edit"), + "video_extension": ("litellm.videos.main", "video_extension"), + "video_generation": ("litellm.videos.main", "video_generation"), + "video_get_character": ("litellm.videos.main", "video_get_character"), + "video_list": ("litellm.videos.main", "video_list"), + "video_remix": ("litellm.videos.main", "video_remix"), + "video_status": ("litellm.videos.main", "video_status"), + "wait": ("litellm.batch_completion.main", "wait"), + "watsonx_chat_completion": ("litellm.main", "watsonx_chat_completion"), + } +) + +_SDK_MODULE_ALIASES: Final[Mapping[str, str]] = MappingProxyType( + { + "additional_logging_utils": "litellm.integrations.additional_logging_utils", + "agentops": "litellm.integrations.agentops", + "aleph_alpha": "litellm.llms.deprecated_providers.aleph_alpha", + "anthropic_cache_control_hook": "litellm.integrations.anthropic_cache_control_hook", + "argilla": "litellm.integrations.argilla", + "arize": "litellm.integrations.arize", + "asyncio": "asyncio", + "athina": "litellm.integrations.athina", + "azure_sentinel": "litellm.integrations.azure_sentinel", + "azure_storage": "litellm.integrations.azure_storage", + "base64": "base64", + "cohere_embed": "litellm.llms.cohere.embed.handler", + "contextvars": "contextvars", + "custom_batch_logger": "litellm.integrations.custom_batch_logger", + "custom_guardrail": "litellm.integrations.custom_guardrail", + "custom_logger": "litellm.integrations.custom_logger", + "custom_prompt_management": "litellm.integrations.custom_prompt_management", + "datadog": "litellm.integrations.datadog", + "datetime": "datetime", + "deepeval": "litellm.integrations.deepeval", + "dotenv": "dotenv", + "dotprompt": "litellm.integrations.dotprompt", + "dynamodb": "litellm.integrations.dynamodb", + "email_templates": "litellm.integrations.email_templates", + "enum": "enum", + "futures": "concurrent.futures", + "galileo": "litellm.integrations.galileo", + "gcs_bucket": "litellm.integrations.gcs_bucket", + "gcs_pubsub": "litellm.integrations.gcs_pubsub", + "generic_api": "litellm.integrations.generic_api", + "greenscale": "litellm.integrations.greenscale", + "heapq": "heapq", + "helicone": "litellm.integrations.helicone", + "helicone_mock_client": "litellm.integrations.helicone_mock_client", + "humanloop": "litellm.integrations.humanloop", + "importlib": "importlib", + "inspect": "inspect", + "json": "json", + "lago": "litellm.integrations.lago", + "langfuse": "litellm.integrations.langfuse", + "langsmith": "litellm.integrations.langsmith", + "langsmith_mock_client": "litellm.integrations.langsmith_mock_client", + "litellm": "litellm", + "litellm_agent": "litellm.integrations.litellm_agent", + "literal_ai": "litellm.integrations.literal_ai", + "logfire_logger": "litellm.integrations.logfire_logger", + "lunary": "litellm.integrations.lunary", + "mimetypes": "mimetypes", + "mlflow": "litellm.integrations.mlflow", + "mock_client_factory": "litellm.integrations.mock_client_factory", + "newrelic": "litellm.integrations.newrelic", + "ollama": "litellm.llms.ollama.completion.handler", + "oobabooga": "litellm.llms.oobabooga.chat.oobabooga", + "openai": "openai", + "openmeter": "litellm.integrations.openmeter", + "opentelemetry": "litellm.integrations.opentelemetry", + "opentelemetry_utils": "litellm.integrations.opentelemetry_utils", + "opik": "litellm.integrations.opik", + "otel": "litellm.integrations.otel", + "palm": "litellm.llms.deprecated_providers.palm", + "petals_handler": "litellm.llms.petals.completion.handler", + "posthog": "litellm.integrations.posthog", + "posthog_mock_client": "litellm.integrations.posthog_mock_client", + "prompt_layer": "litellm.integrations.prompt_layer", + "prompt_management_base": "litellm.integrations.prompt_management_base", + "random": "random", + "rust_ocr_bridge": "litellm.rust_bridge.ocr", + "s3": "litellm.integrations.s3", + "s3_v2": "litellm.integrations.s3_v2", + "sqs": "litellm.integrations.sqs", + "supabase": "litellm.integrations.supabase", + "sys": "sys", + "tiktoken": "tiktoken", + "time": "time", + "traceback": "traceback", + "traceloop": "litellm.integrations.traceloop", + "uuid": "fastuuid", + "uuid_module": "uuid", + "vertex_ai_non_gemini": "litellm.llms.vertex_ai.vertex_ai_non_gemini", + "vllm_handler": "litellm.llms.vllm.completion.handler", + "anthropic": "litellm.anthropic_interface", + "httpx": "httpx", + "interactions": "litellm.interactions", + "rag": "litellm.rag", + } +) + # Export all name tuples and import maps for use in _lazy_imports.py __all__ = [ "BEDROCK_TYPES_NAMES", @@ -1490,6 +2657,7 @@ __all__ = [ "LLM_CLIENT_CACHE_NAMES", "LLM_CONFIG_NAMES", "LLM_PROVIDER_LOGIC_NAMES", + "STAR_IMPORT_PUBLIC_NAMES", "TOKEN_COUNTER_NAMES", "TYPES_NAMES", "TYPES_UTILS_NAMES", @@ -1502,9 +2670,1534 @@ __all__ = [ "_LITELLM_LOGGING_IMPORT_MAP", "_LLM_CONFIGS_IMPORT_MAP", "_LLM_PROVIDER_LOGIC_IMPORT_MAP", + "_SDK_MODULE_ALIASES", + "_SDK_SYMBOLS_IMPORT_MAP", "_TOKEN_COUNTER_IMPORT_MAP", "_TYPES_IMPORT_MAP", "_TYPES_UTILS_IMPORT_MAP", "_UTILS_IMPORT_MAP", "_UTILS_MODULE_IMPORT_MAP", ] + + +STAR_IMPORT_PUBLIC_NAMES: Final = ( + "AI21ChatConfig", + "AI21Config", + "ALL_RESPONSES_API_TOOL_PARAMS", + "APIConnectionError", + "APIError", + "APIResponseValidationError", + "AZURE_DEFAULT_API_VERSION", + "AZURE_OPENAI_AUDIO_PROVIDERS", + "AdapterCompletionStreamWrapper", + "AdapterItem", + "AdaptiveRouterConfig", + "AdaptiveRouterPreferences", + "AdaptiveRouterWeights", + "AlephAlphaConfig", + "AlertingConfig", + "AllEmbeddingInputValues", + "AllMessageValues", + "AllPromptValues", + "AllowedFailsPolicy", + "AmazonTitanV2Config", + "Annotated", + "AnthropicBatchesHandler", + "AnthropicChatCompletion", + "AnthropicMessagesRequestUtils", + "AnthropicMessagesResponse", + "AnthropicMetadata", + "AnthropicModelInfo", + "AnthropicThinkingParam", + "Any", + "Assistant", + "AssistantDeleted", + "AssistantEventHandler", + "AssistantStreamManager", + "AssistantToolParam", + "AssistantsTypedDict", + "AsyncAssistantEventHandler", + "AsyncAssistantStreamManager", + "AsyncCompletions", + "AsyncCursorPage", + "AsyncHTTPHandler", + "AsyncIterator", + "AsyncOpenAI", + "Attachment", + "AttachmentTool", + "AuthenticationError", + "AutoRouterCapabilityLimit", + "AzureAIEmbedding", + "AzureAnthropicChatCompletion", + "AzureAssistantsAPI", + "AzureAudioTranscription", + "AzureBatchesAPI", + "AzureChatCompletion", + "AzureOpenAIFilesAPI", + "AzureOpenAIFineTuningAPI", + "AzureOpenAIO1ChatCompletion", + "AzureTextCompletion", + "BATCH_GUARDRAIL_RESPONSE_FIELD", + "BEDROCK_CONVERSE_MODELS", + "BEDROCK_EMBEDDING_PROVIDERS_LITERAL", + "BEDROCK_INVOKE_PROVIDERS_LITERAL", + "BadGatewayError", + "BadRequestError", + "BaseAnthropicMessagesConfig", + "BaseConfig", + "BaseImageEditConfig", + "BaseImageGenerationConfig", + "BaseLLMAIOHTTPHandler", + "BaseLLMException", + "BaseLLMHTTPHandler", + "BaseLiteLLMOpenAIResponseObject", + "BaseModel", + "BaseOCRConfig", + "BaseRerankConfig", + "BaseResponsesAPIConfig", + "BaseResponsesAPIStreamingIterator", + "BaseSearchConfig", + "BaseVideoConfig", + "Batch", + "BatchGuardrailRecord", + "BatchGuardrailReport", + "BatchJobStatus", + "BatchRequestCounts", + "BedrockBatchesHandler", + "BedrockConverseLLM", + "BedrockEmbedding", + "BedrockFilesHandler", + "BedrockImageEdit", + "BedrockImageGeneration", + "BedrockModelInfo", + "BedrockRerankHandler", + "BudgetExceededError", + "BudgetManager", + "BytezChatConfig", + "CALLBACK_TYPES", + "CARRY_UNMATCHED_MESSAGE_POINTS", + "COHERE_DEFAULT_EMBEDDING_INPUT_TYPE", + "CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS", + "CREATE_FILE_REQUESTS_PURPOSE", + "CallTypes", + "Callable", + "CancelBatchRequest", + "CharacterObject", + "Chat", + "ChatCompletionAnnotation", + "ChatCompletionAnnotationURLCitation", + "ChatCompletionAssistantContentValue", + "ChatCompletionAssistantMessage", + "ChatCompletionAssistantToolCall", + "ChatCompletionAudioDelta", + "ChatCompletionAudioObject", + "ChatCompletionAudioParam", + "ChatCompletionCachedContent", + "ChatCompletionChunk", + "ChatCompletionContentPartInputAudioParam", + "ChatCompletionDeltaChunk", + "ChatCompletionDeltaToolCallChunk", + "ChatCompletionDeveloperMessage", + "ChatCompletionDocumentObject", + "ChatCompletionFileObject", + "ChatCompletionFileObjectFile", + "ChatCompletionFunctionMessage", + "ChatCompletionImageObject", + "ChatCompletionImageUrlObject", + "ChatCompletionMessageToolCall", + "ChatCompletionModality", + "ChatCompletionNamedToolChoiceParam", + "ChatCompletionPredictionContentParam", + "ChatCompletionReasoningItem", + "ChatCompletionReasoningSummaryTextBlock", + "ChatCompletionRedactedThinkingBlock", + "ChatCompletionRequest", + "ChatCompletionResponseMessage", + "ChatCompletionSystemMessage", + "ChatCompletionTextObject", + "ChatCompletionThinkingBlock", + "ChatCompletionToolCallChunk", + "ChatCompletionToolCallFunctionChunk", + "ChatCompletionToolChoiceFunctionParam", + "ChatCompletionToolChoiceObjectParam", + "ChatCompletionToolChoiceStringValues", + "ChatCompletionToolChoiceValues", + "ChatCompletionToolMessage", + "ChatCompletionToolParam", + "ChatCompletionToolParamFunctionChunk", + "ChatCompletionToolReferenceObject", + "ChatCompletionUsageBlock", + "ChatCompletionUserMessage", + "ChatCompletionVideoObject", + "ChatCompletionVideoUrlObject", + "Choices", + "ChunkProcessor", + "CitationsObject", + "ClarifaiConfig", + "ClassVar", + "ClassifierPlugin", + "CodeInterpreterToolParam", + "CodestralTextCompletion", + "CohereModelInfo", + "CompletionRequest", + "CompletionTimeout", + "CompletionTokensDetails", + "Completions", + "ComputerToolParam", + "ConfigDict", + "ConfigurableClientsideParamsCustomAuth", + "ConsumedRequestTagsStamp", + "ContentPartAddedEvent", + "ContentPartDoneEvent", + "ContentPartDonePartOutputText", + "ContentPartDonePartReasoningText", + "ContentPartDonePartRefusal", + "ContentPolicyViolationError", + "ContextManagementEntry", + "ContextWindowExceededError", + "Coroutine", + "CreateBatchRequest", + "CreateFileRequest", + "CreateVideoRequest", + "CredentialLiteLLMParams", + "CustomLLM", + "CustomLLMItem", + "CustomLogger", + "CustomPricingLiteLLMParams", + "CustomRoutingStrategyBase", + "CustomStreamWrapper", + "CustomToolCallOutputItem", + "DEFAULT_ALLOWED_FAILS", + "DEFAULT_BATCH_SIZE", + "DEFAULT_FLUSH_INTERVAL_SECONDS", + "DEFAULT_IMAGE_ENDPOINT_MODEL", + "DEFAULT_IN_MEMORY_TTL", + "DEFAULT_MAX_RETRIES", + "DEFAULT_MAX_TOKENS", + "DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT", + "DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT", + "DEFAULT_POLLING_INTERVAL", + "DEFAULT_REPLICATE_POLLING_DELAY_SECONDS", + "DEFAULT_REPLICATE_POLLING_RETRIES", + "DEFAULT_REQUEST_TIMEOUT", + "DEFAULT_SOFT_BUDGET", + "DEFAULT_VIDEO_ENDPOINT_MODEL", + "DatabricksEmbeddingHandler", + "DatadogInitParams", + "DecodedResponseId", + "DeleteResponseResult", + "Deployment", + "DeploymentTypedDict", + "Dict", + "Discriminator", + "DocumentObject", + "DualCache", + "EmbeddingCreateParams", + "EmbeddingInput", + "EmbeddingRequest", + "EmbeddingResponse", + "Enum", + "ErrorEvent", + "ErrorEventError", + "FIRST_COMPLETED", + "FORWARDED_KWARGS_KEYS", + "FallbackAccessCheck", + "Field", + "FileContent", + "FileContentProvider", + "FileContentRequest", + "FileContentStreamingResponse", + "FileContentStreamingResult", + "FileCreateProvider", + "FileDeleteProvider", + "FileDeleted", + "FileExpiresAfter", + "FileListPage", + "FileListProvider", + "FileObject", + "FileRetrieveProvider", + "FileSearchCallCompletedEvent", + "FileSearchCallInProgressEvent", + "FileSearchCallSearchingEvent", + "FileSearchTool", + "FileSearchToolParam", + "FileTypes", + "Final", + "FineTuningConfig", + "FineTuningJob", + "FineTuningJobCreate", + "FlowItem", + "Function", + "FunctionCallArgumentsDeltaEvent", + "FunctionCallArgumentsDoneEvent", + "GDCGeminiConfig", + "GeminiModelInfo", + "GenAIHubOrchestration", + "Generator", + "Generic", + "GenericBudgetWindowDetails", + "GenericChatCompletionMessage", + "GenericEvent", + "GenericLiteLLMParams", + "GenericResponseOutputItem", + "GenericResponseOutputItemContentAnnotation", + "GoogleBatchEmbeddings", + "GroqChatCompletion", + "GuardrailLiteLLMParams", + "GuardrailTypedDict", + "HTTPHandler", + "HUMANLOOP_PROMPT_CACHE_TTL_SECONDS", + "HerokuChatConfig", + "HiddenParams", + "HttpxBinaryResponseContent", + "HuggingFaceEmbedding", + "Hyperparameters", + "IBMWatsonXMixin", + "IO", + "IOBase", + "ImageEditOptionalRequestParams", + "ImageFetchError", + "ImageFileObject", + "ImageGenerationPartialImageEvent", + "ImageGenerationRequestQuality", + "ImageResponse", + "ImageURLListItem", + "ImageURLObject", + "IncompleteDetails", + "InputTokensDetails", + "InternalServerError", + "InvalidRequestError", + "Iterable", + "Iterator", + "JSONProviderRegistry", + "JSONSchemaValidationError", + "KeyManagementSettings", + "LIST_BATCHES_SUPPORTED_PROVIDERS", + "LITELLM_CHAT_PROVIDERS", + "LITELLM_EXCEPTION_TYPES", + "LITELLM_IMAGE_VARIATION_PROVIDERS", + "LemonadeChatConfig", + "List", + "ListBatchRequest", + "ListBatchesSupportedProvider", + "LiteLLM", + "LiteLLMBatch", + "LiteLLMBatchCreateRequest", + "LiteLLMCompletionTransformationHandler", + "LiteLLMFineTuningJob", + "LiteLLMFineTuningJobCreate", + "LiteLLMLoggingObj", + "LiteLLMMessagesToCompletionTransformationHandler", + "LiteLLMMessagesToResponsesAPIHandler", + "LiteLLMParamsTypedDict", + "LiteLLMResponsesTransformationHandler", + "LiteLLMUnknownProvider", + "LiteLLM_Params", + "LiteLLM_RouterFileObject", + "Literal", + "LlmProviders", + "Logging", + "MCPCallArgumentsDeltaEvent", + "MCPCallArgumentsDoneEvent", + "MCPCallCompletedEvent", + "MCPCallFailedEvent", + "MCPCallInProgressEvent", + "MCPListToolsCompletedEvent", + "MCPListToolsFailedEvent", + "MCPListToolsInProgressEvent", + "MCPTool", + "MOCK_RESPONSE_TYPE", + "Mapping", + "MappingProxyType", + "Message", + "MessageContent", + "MessageContentImageFileObject", + "MessageContentImageURLObject", + "MessageContentTextObject", + "MessageData", + "MirroredPricingParams", + "MockException", + "MockRouterTestingParams", + "ModelConfig", + "ModelGroupInfo", + "ModelGroupSettings", + "ModelInfo", + "ModelResponse", + "ModelResponseStream", + "MyLocal", + "NOT_GIVEN", + "NewRelicInitParams", + "NonNegativeInt", + "NotFoundError", + "NotGiven", + "NotRequired", + "NvidiaRivaAudioTranscription", + "NvidiaRivaAudioTranscriptionConfig", + "OCIChatConfig", + "OCRResponse", + "OCR_REQUEST_FORMAT_PARAM", + "OPENAI_CHAT_COMPLETION_PARAMS", + "OPENAI_COMPATIBLE_BATCH_AND_FILES_PROVIDERS", + "OPENAI_FINISH_REASONS", + "OPTIONAL_KWARGS_KEYS", + "OVHCloudChatConfig", + "Omit", + "OpenAI", + "OpenAIAssistantsAPI", + "OpenAIAudioTranscription", + "OpenAIAudioTranscriptionOptionalParams", + "OpenAIBatchResponse", + "OpenAIBatchResult", + "OpenAIBatchesAPI", + "OpenAIChatCompletion", + "OpenAIChatCompletionAssistantMessage", + "OpenAIChatCompletionChoices", + "OpenAIChatCompletionChunk", + "OpenAIChatCompletionDeveloperMessage", + "OpenAIChatCompletionFinishReason", + "OpenAIChatCompletionLogprobs", + "OpenAIChatCompletionLogprobsContent", + "OpenAIChatCompletionLogprobsContentTopLogprobs", + "OpenAIChatCompletionResponse", + "OpenAIChatCompletionSystemMessage", + "OpenAIChatCompletionTextObject", + "OpenAIChatCompletionToolParam", + "OpenAIChatCompletionUserMessage", + "OpenAICreateFileRequestOptionalParams", + "OpenAICreateThreadParamsMessage", + "OpenAICreateThreadParamsToolResources", + "OpenAIEmbedding", + "OpenAIError", + "OpenAIErrorBody", + "OpenAIFileObject", + "OpenAIFilesAPI", + "OpenAIFilesPurpose", + "OpenAIFineTuningAPI", + "OpenAIGPT5Config", + "OpenAIImageEditOptionalParams", + "OpenAIImageGenerationOptionalParams", + "OpenAIImageVariationOptionalParams", + "OpenAIImageVariationsHandler", + "OpenAILikeChatHandler", + "OpenAILikeEmbeddingHandler", + "OpenAILikeResponsesConfig", + "OpenAIMcpServerTool", + "OpenAIMessage", + "OpenAIMessageContent", + "OpenAIMessageContentListBlock", + "OpenAIModerationResponse", + "OpenAIModerationResult", + "OpenAIRealtimeContentPartDone", + "OpenAIRealtimeConversationCreated", + "OpenAIRealtimeConversationItemAdded", + "OpenAIRealtimeConversationItemCreated", + "OpenAIRealtimeConversationItemDone", + "OpenAIRealtimeConversationObject", + "OpenAIRealtimeDoneEvent", + "OpenAIRealtimeEventTypes", + "OpenAIRealtimeEvents", + "OpenAIRealtimeFunctionCallArgumentsDone", + "OpenAIRealtimeInputAudioBufferSpeechEvent", + "OpenAIRealtimeInputAudioTranscriptionCompleted", + "OpenAIRealtimeInputAudioTranscriptionDelta", + "OpenAIRealtimeOutputItemDone", + "OpenAIRealtimeResponseAudioDone", + "OpenAIRealtimeResponseContentPart", + "OpenAIRealtimeResponseContentPartAdded", + "OpenAIRealtimeResponseDelta", + "OpenAIRealtimeResponseDoneObject", + "OpenAIRealtimeResponseTextDone", + "OpenAIRealtimeResponseUsage", + "OpenAIRealtimeStreamList", + "OpenAIRealtimeStreamResponseBaseObject", + "OpenAIRealtimeStreamResponseOutputItem", + "OpenAIRealtimeStreamResponseOutputItemAdded", + "OpenAIRealtimeStreamResponseOutputItemContent", + "OpenAIRealtimeStreamSession", + "OpenAIRealtimeStreamSessionEvents", + "OpenAIRealtimeTurnDetection", + "OpenAIRealtimeUsageTokenDetails", + "OpenAITextCompletion", + "OpenAITextCompletionUserMessage", + "OpenAIVideoObject", + "OpenAIWebSearchOptions", + "OpenAIWebSearchUserLocation", + "OpenAIWebSearchUserLocationApproximate", + "Optional", + "OptionalPreCallChecks", + "OutputCodeInterpreterCall", + "OutputCodeInterpreterCallLog", + "OutputFunctionToolCall", + "OutputImageGenerationCall", + "OutputItemAddedEvent", + "OutputItemDoneEvent", + "OutputText", + "OutputTextAnnotationAddedEvent", + "OutputTextDeltaEvent", + "OutputTextDoneEvent", + "OutputTokensDetails", + "PART_UNION_TYPES", + "PalmConfig", + "PathLike", + "PermissionDeniedError", + "Phase", + "PreRoutingHookResponse", + "PreRoutingStrategy", + "PredibaseChatCompletion", + "PrivateAttr", + "PromptCacheBreakpoint", + "PromptCacheOptions", + "PromptObject", + "PromptSpec", + "PromptTokensDetails", + "Protocol", + "ProviderConfigManager", + "ProviderSpecificHeader", + "ProviderSpecificHeaderUtils", + "REASONING_EFFORT", + "REPEATED_STREAMING_CHUNK_LIMIT", + "ROUTER_MAX_FALLBACKS", + "RateLimitError", + "RateLimitErrorCategory", + "RateLimitType", + "RawRequestTypedDict", + "ReadOnly", + "Reasoning", + "ReasoningSummaryPartDoneEvent", + "ReasoningSummaryTextDeltaEvent", + "ReasoningSummaryTextDoneEvent", + "RedisCache", + "RefusalDeltaEvent", + "RefusalDoneEvent", + "RequestType", + "Required", + "RerankResponse", + "Response", + "ResponseAPIUsage", + "ResponseCompletedEvent", + "ResponseCreatedEvent", + "ResponseFailedEvent", + "ResponseFunctionToolCall", + "ResponseInProgressEvent", + "ResponseIncludable", + "ResponseIncompleteEvent", + "ResponseInputParam", + "ResponseOutputItem", + "ResponsePartAddedEvent", + "ResponseText", + "ResponsesAPIOptionalRequestParams", + "ResponsesAPIRequestParams", + "ResponsesAPIRequestUtils", + "ResponsesAPIResponse", + "ResponsesAPIStatus", + "ResponsesAPIStreamEvents", + "ResponsesAPIStreamOptions", + "ResponsesAPIStreamingResponse", + "ResponsesToolUsage", + "RetrieveBatchRequest", + "RetryPolicy", + "Router", + "RouterCacheEnum", + "RouterConfig", + "RouterErrors", + "RouterGeneralSettings", + "RouterModelGroupAliasItem", + "RouterRateLimitError", + "RouterRateLimitErrorBasic", + "RoutingContext", + "RoutingGroup", + "RoutingPlugin", + "RoutingStrategy", + "Run", + "SPECIAL_MODEL_INFO_PARAMS", + "SagemakerChatHandler", + "SagemakerLLM", + "Scheduler", + "SchedulerCacheKeys", + "SearchProvider", + "SearchProviders", + "SearchResponse", + "SearchToolInfoTypedDict", + "SearchToolLiteLLMParams", + "SearchToolTypedDict", + "Sequence", + "SerializerFunctionWrapHandler", + "ServiceUnavailableError", + "Set", + "ShellToolParam", + "SlackAlerting", + "StandardLoggingRoutingDecision", + "StreamingChoices", + "SyncCursorPage", + "TYPE_CHECKING", + "TaggedPreRoutingStrategy", + "TextChoices", + "TextCompletionResponse", + "TextCompletionStreamWrapper", + "Thread", + "ThreadPoolExecutor", + "Timeout", + "TogetherAIRerank", + "Tool", + "ToolChoice", + "ToolMessageContentPart", + "ToolParam", + "ToolResourcesCodeInterpreter", + "ToolResourcesFileSearch", + "ToolResourcesFileSearchVectorStore", + "TopazModelInfo", + "TranscriptionResponse", + "Tuple", + "Type", + "TypeAlias", + "TypeVar", + "TypedDict", + "Union", + "UnprocessableEntityError", + "UnsupportedParamsError", + "UpdateRouterConfig", + "Usage", + "VALID_LITELLM_ENVIRONMENTS", + "ValidAssistantMessageContentTypes", + "ValidAssistantMessageContentTypesLiteral", + "ValidChatCompletionMessageContentTypes", + "ValidChatCompletionMessageContentTypesLiteral", + "ValidUserMessageContentTypes", + "ValidUserMessageContentTypesLiteral", + "VectorStoreIndexRegistry", + "VectorStoreRegistry", + "VertexAIBatchPrediction", + "VertexAIFilesHandler", + "VertexAIGemmaModels", + "VertexAIModelGardenModels", + "VertexAIModelRoute", + "VertexAIPartnerModels", + "VertexAITextEmbeddingConfig", + "VertexEmbedding", + "VertexFineTuningAPI", + "VertexImageGeneration", + "VertexLLM", + "VertexMultimodalEmbedding", + "VideoCreateOptionalRequestParams", + "VideoGenerationRequestUtils", + "VideoObject", + "WANDB_MODELS", + "WATSONX_DEFAULT_API_VERSION", + "WatsonXChatHandler", + "WebSearchCallCompletedEvent", + "WebSearchCallInProgressEvent", + "WebSearchCallSearchingEvent", + "WebSearchOptions", + "WebSearchOptionsUserLocation", + "WebSearchOptionsUserLocationApproximate", + "WebSearchToolUsage", + "XAIModelInfo", + "a_add_message", + "aadapter_completion", + "aadapter_generate_content", + "acancel_batch", + "acancel_eval", + "acancel_fine_tuning_job", + "acancel_responses", + "acancel_run", + "aclient_session", + "acode_interpreter_tool", + "acompact_responses", + "acompletion", + "acompletion_with_retries", + "acount_tokens", + "acreate_agent", + "acreate_assistants", + "acreate_batch", + "acreate_container", + "acreate_eval", + "acreate_file", + "acreate_fine_tuning_job", + "acreate_realtime_client_secret", + "acreate_realtime_transcription_session", + "acreate_run", + "acreate_sandbox", + "acreate_skill", + "acreate_thread", + "adapter_completion", + "adapters", + "add_function_to_prompt", + "add_known_models", + "add_message", + "add_provider_specific_params_to_optional_params", + "add_system_prompt_to_messages", + "add_trusted_model_credentials_to_litellm_params", + "add_user_information_to_llm_headers", + "additional_logging_utils", + "adelete_agent", + "adelete_assistant", + "adelete_container", + "adelete_eval", + "adelete_responses", + "adelete_run", + "adelete_sandbox", + "adelete_skill", + "aembedding", + "afile_content", + "afile_delete", + "afile_list", + "afile_retrieve", + "agenerate_content", + "agent_search_embedding_model", + "agentops", + "aget_agent", + "aget_assistants", + "aget_eval", + "aget_messages", + "aget_responses", + "aget_run", + "aget_skill", + "aget_thread", + "ahealth_check", + "ai21_chat_models", + "ai21_key", + "ai21_models", + "aimage_edit", + "aimage_generation", + "aimage_variation", + "aiml_models", + "aingest", + "aiohttp_trust_env", + "aleph_alpha", + "aleph_alpha_key", + "aleph_alpha_models", + "alist_agent_versions", + "alist_agents", + "alist_batches", + "alist_container_files", + "alist_containers", + "alist_evals", + "alist_fine_tuning_jobs", + "alist_input_items", + "alist_runs", + "alist_skills", + "all_embedding_models", + "all_litellm_params", + "allm_passthrough_route", + "allow_dynamic_callback_disabling", + "allowed_fails", + "amazon_nova_api_key", + "amazon_nova_models", + "amoderation", + "annotations", + "anthropic", + "anthropic_batches_instance", + "anthropic_beta_headers_manager", + "anthropic_beta_headers_url", + "anthropic_cache_control_hook", + "anthropic_chat_completions", + "anthropic_interface", + "anthropic_key", + "anthropic_messages", + "anthropic_messages_handler", + "anthropic_models", + "anthropic_prompt_caching_ttl", + "anthropic_sse_ping_interval_seconds", + "anyscale_models", + "aocr", + "api_base", + "api_key", + "api_version", + "aquery", + "arealtime_calls", + "arerank", + "aresponses", + "aresponses_api_with_mcp", + "aresponses_with_retries", + "aretrieve_batch", + "aretrieve_container", + "aretrieve_fine_tuning_job", + "argilla", + "argilla_batch_size", + "argilla_transformation_object", + "arize", + "arun_code", + "arun_thread", + "arun_thread_stream", + "asearch", + "aspeech", + "assemblyai_models", + "assistants", + "async_completion_with_fallbacks", + "async_mock_completion_streaming_obj", + "asyncio", + "atext_completion", + "athina", + "atranscription", + "audit_log_callbacks", + "aupload_container_file", + "autorouter_presets_url", + "avector_store_file_content", + "avector_store_file_create", + "avector_store_file_delete", + "avector_store_file_list", + "avector_store_file_retrieve", + "avector_store_file_update", + "avideo_content", + "avideo_create_character", + "avideo_edit", + "avideo_extension", + "avideo_generation", + "avideo_get_character", + "avideo_list", + "avideo_remix", + "avideo_status", + "aws_polly_models", + "aws_sqs_callback_params", + "azure_ai_embedding", + "azure_ai_models", + "azure_anthropic_chat_completions", + "azure_anthropic_models", + "azure_assistants_api", + "azure_audio_transcriptions", + "azure_batches_instance", + "azure_chat_completions", + "azure_embedding_models", + "azure_files_instance", + "azure_fine_tuning_apis_instance", + "azure_key", + "azure_llms", + "azure_models", + "azure_o1_chat_completions", + "azure_sentinel", + "azure_storage", + "azure_text_completions", + "azure_text_models", + "banned_keywords_list", + "base64", + "base_llm_aiohttp_handler", + "base_llm_http_handler", + "baseten_key", + "baseten_models", + "batch_completion", + "batch_completion_models", + "batch_completion_models_all_responses", + "batches", + "bedrock_converse_chat_completion", + "bedrock_converse_models", + "bedrock_embedding", + "bedrock_embedding_models", + "bedrock_files_instance", + "bedrock_image_edit", + "bedrock_image_generation", + "bedrock_mantle_models", + "bedrock_models", + "bedrock_request_metadata_fields", + "bedrock_rerank", + "bfl_image_edit", + "bfl_image_generation", + "black_forest_labs_models", + "block_requests_for_models_without_pricing", + "blocked_user_list", + "blog_posts_url", + "budget_duration", + "budget_exceeded_throttle_percentage", + "budget_manager", + "budget_rollover", + "build_code_interpreter_log_outputs", + "bytez_key", + "bytez_transformation", + "cache", + "caching", + "caching_with_models", + "calculate_request_duration", + "callback_settings", + "callbacks", + "cancel_batch", + "cancel_eval", + "cancel_fine_tuning_job", + "cancel_responses", + "cancel_run", + "cast", + "cerebras_models", + "chatgpt_models", + "check_provider_endpoint", + "clarifai_key", + "clarifai_models", + "client", + "client_session", + "close_litellm_async_clients", + "cloudflare_api_key", + "cloudflare_models", + "codestral_models", + "codestral_text_completions", + "cohere_chat_models", + "cohere_embed", + "cohere_embedding_models", + "cohere_key", + "cohere_models", + "cold_storage_custom_logger", + "cometapi_key", + "cometapi_models", + "common_cloud_provider_auth_params", + "compact_responses", + "completion", + "completion_extras", + "completion_with_fallbacks", + "completion_with_retries", + "compress", + "compression", + "config_completion", + "config_path", + "constants", + "containers", + "content_policy_fallbacks", + "context_window_fallbacks", + "contextmanager", + "contextvars", + "convert_file_document_to_url_document", + "convert_model_response_to_streaming", + "convert_to_model_response_object", + "cost_calculator", + "cost_discount_config", + "cost_margin_config", + "create_agent", + "create_assistants", + "create_batch", + "create_container", + "create_eval", + "create_file", + "create_fine_tuning_job", + "create_pretrained_tokenizer", + "create_run", + "create_skill", + "create_thread", + "create_tokenizer", + "credential_list", + "custom_batch_logger", + "custom_chat_llm_router", + "custom_guardrail", + "custom_logger", + "custom_prometheus_metadata_labels", + "custom_prometheus_tags", + "custom_prompt", + "custom_prompt_dict", + "custom_prompt_management", + "custom_provider_map", + "darkbloom_models", + "dashscope_models", + "databricks_embedding", + "databricks_key", + "databricks_models", + "dataclass", + "datadog", + "datadog_llm_observability_params", + "datadog_params", + "datadog_use_v1", + "datarobot_key", + "datarobot_models", + "datetime", + "decode_video_id_with_provider", + "declared_authenticating_provider", + "deepcopy", + "deepeval", + "deepgram_models", + "deepinfra_models", + "deepseek_models", + "default_fallbacks", + "default_in_memory_ttl", + "default_internal_user_params", + "default_key_generate_params", + "default_key_max_budget_alert_emails", + "default_max_internal_user_budget", + "default_redis_batch_cache_expiry", + "default_redis_ttl", + "default_soft_budget", + "default_team_params", + "default_team_settings", + "delete_agent", + "delete_assistant", + "delete_container", + "delete_eval", + "delete_responses", + "delete_run", + "delete_skill", + "disable_add_prefix_to_prompt", + "disable_add_transform_inline_image_block", + "disable_add_user_agent_to_request_tags", + "disable_aiohttp_transport", + "disable_aiohttp_trust_env", + "disable_anthropic_gemini_context_caching_transform", + "disable_cache", + "disable_copilot_system_to_assistant", + "disable_end_user_cost_tracking", + "disable_end_user_cost_tracking_prometheus_only", + "disable_hf_tokenizer_download", + "disable_stop_sequence_limit", + "disable_streaming_logging", + "disable_token_counter", + "disable_vertex_batch_output_transformation", + "docker_model_runner_models", + "dotenv", + "dotprompt", + "drop_params", + "dynamodb", + "dynamodb_table_name", + "elevenlabs_models", + "email", + "email_templates", + "embedding", + "empower_models", + "enable_anthropic_prompt_caching", + "enable_azure_ad_token_refresh", + "enable_cache", + "enable_caching_on_provider_specific_optional_params", + "enable_end_user_cost_tracking_prometheus_only", + "enable_gemini_default_thinking_level_low", + "enable_json_schema_validation", + "enable_key_alias_format_validation", + "enable_loadbalancing_on_batch_endpoints", + "enable_model_config_credential_overrides", + "enable_preview_features", + "enum", + "error_logs", + "evals", + "exception_type", + "exceptions", + "expose_router_debug_in_errors", + "extra_spend_tag_headers", + "failure_callback", + "fal_ai_models", + "fallbacks", + "featherless_ai_models", + "field_serializer", + "field_validator", + "file_content", + "file_content_streaming", + "file_delete", + "file_list", + "file_retrieve", + "files", + "filter_invalid_headers", + "filter_out_litellm_params", + "fine_tuning", + "fireworks_ai_embedding_models", + "fireworks_ai_models", + "flatten_form_field_values", + "flatten_unencrypted_web_search_results_in_anthropic_messages", + "force_ipv4", + "forward_traceparent_to_llm_provider", + "friendliai_models", + "function_call_prompt", + "futures", + "galadriel_models", + "galileo", + "gcs_bucket", + "gcs_pub_sub_use_v1", + "gcs_pubsub", + "gdc_api_base", + "gdc_key", + "gdc_transformation", + "gemini_live_defer_setup", + "gemini_models", + "generic_api", + "generic_api_use_v1", + "generic_logger_headers", + "get_agent", + "get_api_key_from_env", + "get_args", + "get_assistants", + "get_audio_file_for_health_check", + "get_azure_credentials", + "get_completion_messages", + "get_configured_request_timeout", + "get_content_from_model_response", + "get_eval", + "get_litellm_gateway_api_key", + "get_litellm_params", + "get_llm_provider", + "get_messages", + "get_messages_interceptors", + "get_mime_type", + "get_model_cost_map", + "get_model_info", + "get_non_default_completion_params", + "get_non_default_transcription_params", + "get_openai_credentials", + "get_optional_params", + "get_optional_params_add_message", + "get_optional_params_embeddings", + "get_optional_params_image_gen", + "get_optional_params_transcription", + "get_optional_rerank_params", + "get_requester_metadata", + "get_responses", + "get_run", + "get_secret", + "get_secret_bool", + "get_secret_str", + "get_skill", + "get_standard_openai_params", + "get_thread", + "get_type_hints", + "get_vertex_ai_model_route", + "gigachat_key", + "gigachat_models", + "github_copilot_models", + "global_bitbucket_config", + "global_disable_no_log_param", + "global_gitlab_config", + "google_batch_embeddings", + "google_genai", + "google_moderation_confidence_threshold", + "gradient_ai_api_key", + "gradient_ai_models", + "greenscale", + "groq_chat_completions", + "groq_key", + "groq_models", + "guardrail_name_config_map", + "headers", + "heapq", + "helicone", + "helicone_mock_client", + "heroku_key", + "heroku_models", + "heroku_transformation", + "httpx", + "huggingface_embed", + "huggingface_key", + "huggingface_models", + "humanloop", + "hyperbolic_models", + "identify", + "image_edit", + "image_generation", + "image_variation", + "images", + "importlib", + "in_memory_llm_clients_cache", + "inception_key", + "inception_models", + "include_cost_in_streaming_usage", + "infer_openai_data_residency", + "infinity_key", + "infinity_models", + "ingest", + "initialized_langfuse_clients", + "input_callback", + "inspect", + "integrations", + "interactions", + "internal_user_budget_duration", + "is_azure_document_intelligence_model", + "is_bedrock_pricing_only_model", + "is_openai_finetune_model", + "is_reasoning_auto_summary_enabled", + "jina_ai_models", + "json", + "json_logs", + "key_generation_settings", + "known_tokenizer_config", + "lago", + "lambda_ai_models", + "langfuse", + "langfuse_default_tags", + "langfuse_enable_update_trace_keys", + "langsmith", + "langsmith_batch_size", + "langsmith_mock_client", + "lemonade_key", + "lemonade_models", + "lemonade_transformation", + "list_agent_versions", + "list_agents", + "list_batches", + "list_container_files", + "list_containers", + "list_evals", + "list_fine_tuning_jobs", + "list_input_items", + "list_runs", + "list_skills", + "litellm", + "litellm_agent", + "litellm_completion_transformation_handler", + "litellm_core_utils", + "litellm_mode", + "literal_ai", + "llama_api_key", + "llama_models", + "llamagate_models", + "llamaguard_model_name", + "llamaguard_unsafe_content_categories", + "llm_guard_mode", + "llm_http_handler", + "llm_passthrough_route", + "llms", + "log_client_error_tracebacks", + "log_level", + "log_raw_request_response", + "logfire_logger", + "logged_real_time_event_types", + "logging", + "longer_context_model_fallback_dict", + "lunary", + "main", + "map_system_message_pt", + "maritalk_key", + "maritalk_models", + "max_budget", + "max_end_user_budget", + "max_end_user_budget_id", + "max_fallbacks", + "max_internal_user_budget", + "max_tokens", + "max_ui_session_budget", + "max_user_budget", + "maybe_run_chat_completion_agentic_loop", + "mcp_tool_search", + "mimetypes", + "minimax_models", + "mistral_chat_models", + "mlflow", + "mock_client_factory", + "mock_completion", + "mock_completion_streaming_obj", + "mock_embedding", + "mock_image_generation", + "mock_response", + "mock_responses_api_response", + "model_alias_map", + "model_cost", + "model_cost_map_url", + "model_fallbacks", + "model_group_settings", + "model_list", + "model_list_set", + "model_serializer", + "model_validator", + "models", + "models_by_provider", + "modelscope_models", + "moderation", + "modify_params", + "moonshot_models", + "morph_models", + "nebius_embedding_models", + "nebius_key", + "nebius_models", + "network_mock", + "newrelic", + "newrelic_params", + "nlp_cloud_chat_completion", + "nlp_cloud_key", + "nlp_cloud_models", + "novita_api_key", + "novita_models", + "nscale_models", + "num_retries", + "num_retries_per_request", + "nvidia_nim_models", + "nvidia_riva_audio_transcriptions", + "nvidia_riva_models", + "oci_models", + "oci_transformation", + "ocr", + "ollama", + "ollama_key", + "ollama_models", + "ollama_pt", + "oobabooga", + "open_ai_chat_completion_models", + "open_ai_embedding_models", + "open_ai_text_completion_models", + "openai", + "openai_assistants_api", + "openai_audio_transcriptions", + "openai_batches_instance", + "openai_chat_completions", + "openai_compatible_endpoints", + "openai_compatible_providers", + "openai_files_instance", + "openai_fine_tuning_apis_instance", + "openai_image_generation_models", + "openai_image_variations", + "openai_key", + "openai_like_chat_completion", + "openai_like_embedding", + "openai_like_key", + "openai_moderations_model_name", + "openai_text_completion_compatible_providers", + "openai_text_completions", + "openai_video_generation_models", + "openmeter", + "openrouter_key", + "openrouter_models", + "opentelemetry", + "opentelemetry_utils", + "opik", + "organization", + "os", + "otel", + "output_parse_pii", + "overload", + "override", + "overwrite_user_with_key_hash", + "ovhcloud_embedding_models", + "ovhcloud_key", + "ovhcloud_models", + "ovhcloud_transformation", + "palm", + "palm_models", + "parse_ocr_request_format", + "partial", + "passthrough", + "peek_reasoning_summary_aliases", + "perplexity_models", + "petals_handler", + "petals_models", + "post_call_rules", + "posthog", + "posthog_mock_client", + "pre_call_rules", + "pre_process_non_default_params", + "predibase_chat_completions", + "predibase_key", + "predibase_tenant_id", + "presidio_ad_hoc_recognizers", + "print_verbose", + "priority_reservation", + "project", + "prometheus_deployment_and_latency_caller_identity", + "prometheus_emit_rate_limit_labels", + "prometheus_emit_stream_label", + "prometheus_end_user_metrics_cleanup_interval_seconds", + "prometheus_end_user_metrics_max_series_per_metric", + "prometheus_end_user_metrics_ttl_seconds", + "prometheus_exclude_labels", + "prometheus_exclude_metrics", + "prometheus_initialize_budget_metrics", + "prometheus_latency_buckets", + "prometheus_metrics_config", + "prometheus_user_budget_label_include_email_alias", + "prompt_factory", + "prompt_layer", + "prompt_management_base", + "prompt_name_config_map", + "provider_url_destination_allowed_hosts", + "proxy", + "proxy_auth", + "public_agent_groups", + "public_mcp_hub_strict_whitelist", + "public_mcp_servers", + "public_model_groups", + "public_model_groups_links", + "publicai_models", + "query", + "qwen_ai_platform_models", + "qwencloud_models", + "rag", + "random", + "re", + "read_config_args", + "realtime_api", + "reasoning_auto_summary", + "recraft_models", + "redact_messages_in_exceptions", + "redact_user_api_key_info", + "reducto_models", + "replicate_chat_completion", + "replicate_key", + "replicate_models", + "repositories", + "request_correlation_in_logs", + "request_timeout", + "request_timeout_explicitly_set", + "require_auth_for_metrics_endpoint", + "require_managed_files", + "rerank", + "rerank_api", + "responses", + "responses_api_bridge_check", + "responses_with_retries", + "retrieve_batch", + "retrieve_container", + "retrieve_fine_tuning_job", + "retry", + "return_response_headers", + "route_all_chat_openai_to_responses", + "router", + "router_strategy", + "router_utils", + "run_async_function", + "run_server", + "run_thread", + "run_thread_stream", + "runtime_checkable", + "runwayml_models", + "rust", + "rust_bridge", + "rust_ocr_bridge", + "s3", + "s3_audit_callback_params", + "s3_callback_params", + "s3_v2", + "safe_deep_copy", + "safe_memory_mode", + "sagemaker_chat_completion", + "sagemaker_llm", + "sambanova_embedding_models", + "sambanova_models", + "sandbox", + "sanitize_tool_use_ids_in_anthropic_messages", + "sap_gen_ai_hub_chat_completions", + "sap_gen_ai_hub_emb", + "sap_service_key", + "scheduler", + "search", + "secret_manager_client", + "secret_managers", + "service_callback", + "set_global_bitbucket_config", + "set_global_gitlab_config", + "set_verbose", + "should_run_mock_completion", + "skills", + "skip_system_message_in_guardrail", + "skip_tool_message_in_guardrail", + "snowflake_key", + "snowflake_models", + "soniox_models", + "speech", + "sqs", + "sse_keepalive_ping_interval_seconds", + "ssl_certificate", + "ssl_ecdh_curve", + "ssl_security_level", + "ssl_verify", + "stability_models", + "standard_logging_payload_excluded_fields", + "store_audit_logs", + "stream_chunk_builder", + "stream_chunk_builder_text_completion", + "stringify_json_tool_call_content", + "strip_anthropic_total_tokens", + "strip_empty_content_blocks_from_anthropic_messages", + "strip_reasoning_summary_aliases_from_optional_params", + "success_callback", + "supabase", + "supports_httpx_timeout", + "suppress_debug_info", + "sys", + "tag_budget_config", + "telemetry", + "tencent_models", + "text_completion", + "text_completion_codestral_models", + "text_completion_inception_models", + "threading", + "tiktoken", + "time", + "together_ai_models", + "together_rerank", + "togetherai_api_key", + "token", + "token_counter", + "traceback", + "traceloop", + "tracer", + "transcription", + "turn_off_message_logging", + "types", + "updateDeployment", + "updateLiteLLMParams", + "update_cache", + "update_messages_with_model_file_ids", + "update_responses_input_with_model_file_ids", + "update_responses_tools_with_model_file_ids", + "upload_container_file", + "upperbound_key_generate_params", + "urlsplit", + "use_aiohttp_transport", + "use_chat_completions_url_for_anthropic_messages", + "use_client", + "use_legacy_interactions_schema", + "use_litellm_proxy", + "user_url_allowed_hosts", + "user_url_validation", + "utils", + "uuid", + "uuid_module", + "v0_models", + "validate_and_fix_openai_messages", + "validate_and_fix_openai_tools", + "validate_and_fix_thinking_param", + "validate_anthropic_api_metadata", + "validate_chat_completion_tool_choice", + "validate_end_user_id_in_db", + "validate_openai_optional_params", + "vector_store_file_content", + "vector_store_file_create", + "vector_store_file_delete", + "vector_store_file_list", + "vector_store_file_retrieve", + "vector_store_file_update", + "vector_store_files", + "vector_store_index_registry", + "vector_store_registry", + "vector_stores", + "verbose_logger", + "vercel_ai_gateway_key", + "vercel_ai_gateway_models", + "vertexAITextEmbeddingConfig", + "vertex_ai_ai21_models", + "vertex_ai_batches_instance", + "vertex_ai_files_instance", + "vertex_ai_image_models", + "vertex_ai_non_gemini", + "vertex_ai_safety_settings", + "vertex_ai_video_models", + "vertex_anthropic_models", + "vertex_chat_completion", + "vertex_chat_models", + "vertex_code_chat_models", + "vertex_code_text_models", + "vertex_deepseek_models", + "vertex_embedding", + "vertex_embedding_models", + "vertex_fine_tuning_apis_instance", + "vertex_gemma_chat_completion", + "vertex_image_generation", + "vertex_language_models", + "vertex_llama3_models", + "vertex_location", + "vertex_minimax_models", + "vertex_mistral_models", + "vertex_model_garden_chat_completion", + "vertex_moonshot_models", + "vertex_multimodal_embedding", + "vertex_openai_models", + "vertex_partner_models_chat_completion", + "vertex_project", + "vertex_text_models", + "vertex_vision_models", + "vertex_zai_models", + "video_content", + "video_create_character", + "video_edit", + "video_extension", + "video_generation", + "video_get_character", + "video_list", + "video_remix", + "video_status", + "videos", + "vllm_handler", + "volcengine_models", + "voyage_models", + "wait", + "wandb_key", + "wandb_models", + "warnings", + "watsonx_chat_completion", + "watsonx_models", + "xai_key", + "xai_models", + "zai_models", +) diff --git a/litellm/proxy/__init__.py b/litellm/proxy/__init__.py index b6e690fd591..dc819fbc85c 100644 --- a/litellm/proxy/__init__.py +++ b/litellm/proxy/__init__.py @@ -1 +1,11 @@ -from . import * +from types import ModuleType +from typing import Final + + +def __getattr__(name: str) -> ModuleType: + from litellm._lazy_imports import lazy_import_submodule + + submodule: Final = lazy_import_submodule(__name__, name) + if submodule is not None: + return submodule + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") diff --git a/tests/test_litellm/test_lazy_imports.py b/tests/test_litellm/test_lazy_imports.py index 2b16a812611..09d5e23e2e6 100644 --- a/tests/test_litellm/test_lazy_imports.py +++ b/tests/test_litellm/test_lazy_imports.py @@ -1,5 +1,8 @@ """Simple tests for lazy import functionality.""" +import importlib +import json +import subprocess import sys import pytest @@ -7,6 +10,10 @@ import pytest import litellm from litellm._lazy_imports import ( + _SDK_MODULE_ALIASES, + _SDK_SYMBOLS_IMPORT_MAP, + lazy_import_litellm_submodule, + _lazy_import_sdk_symbols, COST_CALCULATOR_NAMES, LITELLM_LOGGING_NAMES, UTILS_NAMES, @@ -346,3 +353,83 @@ def test_utils_module_lazy_imports(): assert name in utils_globals _verify_only_requested_name_imported_in_utils(name, UTILS_MODULE_NAMES) + + +def test_sdk_symbols_lazy_imports(): + """Every symbol previously imported eagerly in litellm/__init__.py resolves to the source module attribute.""" + for name, (module_path, attr_name) in _SDK_SYMBOLS_IMPORT_MAP.items(): + resolved = getattr(litellm, name) + expected = getattr(importlib.import_module(module_path), attr_name) + assert resolved is expected, f"litellm.{name} is not {module_path}.{attr_name}" + + +def test_sdk_module_aliases(): + """Module-valued attributes (litellm.anthropic, litellm.httpx, ...) resolve to the aliased modules.""" + for name, module_path in _SDK_MODULE_ALIASES.items(): + assert getattr(litellm, name) is importlib.import_module(module_path) + + +def test_litellm_submodule_fallback(): + """litellm. attribute access resolves real submodules and returns None for unknown names.""" + assert lazy_import_litellm_submodule("budget_manager") is importlib.import_module("litellm.budget_manager") + assert litellm.utils is importlib.import_module("litellm.utils") + assert lazy_import_litellm_submodule("not_a_real_submodule") is None + with pytest.raises(AttributeError): + _ = litellm.not_a_real_attribute + + +def test_missing_attribute_stays_attribute_error_when_find_spec_lies(monkeypatch): + """getattr(litellm, name, default) must not leak ModuleNotFoundError when find_spec is patched to always succeed.""" + monkeypatch.setattr(importlib.util, "find_spec", lambda name: object()) + assert getattr(litellm, "not_a_real_submodule", None) is None + with pytest.raises(AttributeError): + _ = litellm.not_a_real_attribute + + +def test_proxy_private_submodule_resolves_in_fresh_process(): + """litellm.proxy._types resolves without an eager proxy import (used by documentation checks).""" + code = "import litellm\nprint(litellm.proxy._types.__name__)\n" + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "litellm.proxy._types" + + +def test_lazy_instances_are_singletons(): + """Lazily created instances are cached, so repeated access returns the same object.""" + assert litellm._key_management_settings is litellm._key_management_settings + assert litellm.vertexAITextEmbeddingConfig is litellm.vertexAITextEmbeddingConfig + from litellm.types.secret_managers.main import KeyManagementSettings + + assert isinstance(litellm._key_management_settings, KeyManagementSettings) + + +def test_star_import_exports_public_api(): + """`from litellm import *` keeps exporting the full public surface despite lazy loading.""" + code = ( + "from litellm import *\n" + "import litellm\n" + "missing = [n for n in litellm.__all__ if n not in dir()]\n" + "assert not missing, missing[:20]\n" + "assert callable(completion) and callable(Router)\n" + ) + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True) + assert result.returncode == 0, result.stderr + + +@pytest.mark.skipif(sys.platform != "linux", reason="reads /proc for RSS") +def test_import_litellm_stays_lightweight(): + """`import litellm` must not pull in the SDK/proxy heavyweights or blow up RSS (LIT-6607).""" + code = ( + "import json, re, sys\n" + "import litellm\n" + "heavy = [m for m in ('litellm.main', 'litellm.utils', 'litellm.router', 'litellm.proxy.proxy_cli',\n" + " 'tiktoken', 'fastapi', 'grpc', 'boto3') if m in sys.modules]\n" + "with open('/proc/self/status') as f:\n" + " rss_kb = int(re.search(r'VmRSS:\\s+(\\d+) kB', f.read()).group(1))\n" + "print(json.dumps({'total': len(sys.modules), 'heavy': heavy, 'rss_mb': rss_kb / 1024}))\n" + ) + result = subprocess.run([sys.executable, "-c", code], capture_output=True, text=True, check=True) + stats = json.loads(result.stdout) + assert stats["heavy"] == [], f"heavy modules imported eagerly: {stats['heavy']}" + assert stats["total"] < 800, f"import litellm loaded {stats['total']} modules" + assert stats["rss_mb"] < 75, f"import litellm used {stats['rss_mb']:.1f} MB RSS" From 948e5755eba9cb80e1239ecebfb717fcad9b2c36 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Sat, 5 Sep 2026 13:03:28 -0700 Subject: [PATCH 2/7] test(e2e): cover presidio post_call, tool_permission, and weave logging cells (#39279) * test(e2e): cover presidio post_call, tool_permission, and weave logging cells Five registry cells in Logging & Guardrails had no covering test. Each one now has a live scenario read back from the real destination: - guardrail.presidio.post_call.masks: an output-scoped Presidio guardrail anonymizes the PII the model repeats back. The prompt also asks for the address's local part, which Presidio does not mask, so one response proves the model saw the raw address (no pre-call masking) while the address itself comes back as - guardrail.tool_permission.pre_call.blocks / .allows: an allow-list of one tool. A request declaring an unlisted tool is rejected 400 naming it; a request declaring the permitted tool is served and carries x-litellm-applied-guardrails, so the allow half cannot pass by the guardrail never running - logging.niche_integrations.success.logs_spend / .failure.logs_spend: a key-scoped weave_otel callback delivers to the real Weave project, read back through Weave's query API. Success asserts exactly one call whose llm.response.cost equals the x-litellm-response-cost header; failure asserts one ERROR-status call naming the provider exception and carrying no cost Logging & Guardrails coverage goes 24/59 to 29/59. No registry rows are added. * test(e2e): make the tool-permission allow case deterministic and scope the Weave read-back Review follow-ups on the coverage PR. - the allow scenario forced the outcome to depend on whether the model felt like calling an optional tool, and checked for the tool name as a substring of the whole body, which a prose mention would satisfy. It now sends tool_choice="required" and asserts the parsed response carries exactly one tool call, for the permitted tool - the Weave read-back queried the newest 200 calls of a shared project and filtered client-side, so busy traffic could push the target out of the window and read as a delivery failure. The query now scopes server-side to the litellm_request op and to calls started after the request, and pages through the window with offset - the reader builds its results as tuples instead of accumulating into lists Also unblocks the lint gate: `basedpyright tests/e2e` runs only on PRs that touch tests/e2e, and it has been failing on staging for three FakeItem arguments in test_junit_properties.py. The stand-in now goes through one typed adapter that says why, so the gate is green without touching junit_properties.py itself. * test(e2e): scope the presidio post_call guardrail to email and phone Running the suite three times in a row caught a real flake: Presidio's broader recognizers sometimes claim the email's local part as an NRP entity, so the answer came back as `\n\n` and the assertion that the raw local part survives failed. That token is what tells output masking apart from input masking, so it has to survive. The post_call guardrail now registers pii_entities_config for EMAIL_ADDRESS and PHONE_NUMBER only, which is also the narrower thing the scenario means. Verified against the exact marker that failed, plus two others. * test(e2e): mark weave logging cells stage red * test(e2e): use per-test stage red skips for the weave logging cells --- tests/e2e/CONTRIBUTING.md | 4 +- tests/e2e/guardrails/guardrails_client.py | 71 ++++- .../guardrails/test_presidio_masking_e2e.py | 106 ++++++- .../test_tool_permission_guardrail_e2e.py | 167 +++++++++++ tests/e2e/logging/logging_client.py | 39 +++ tests/e2e/logging/test_weave_log_e2e.py | 192 ++++++++++++ tests/e2e/logging/weave_reader.py | 282 ++++++++++++++++++ tests/e2e/models.py | 2 + 8 files changed, 844 insertions(+), 19 deletions(-) create mode 100644 tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py create mode 100644 tests/e2e/logging/test_weave_log_e2e.py create mode 100644 tests/e2e/logging/weave_reader.py diff --git a/tests/e2e/CONTRIBUTING.md b/tests/e2e/CONTRIBUTING.md index 871d6b3904c..b270feb820e 100644 --- a/tests/e2e/CONTRIBUTING.md +++ b/tests/e2e/CONTRIBUTING.md @@ -50,7 +50,9 @@ The suites run against a live proxy, so bring one up first by running the litell They also need a proxy whose bundled UI contains the change under test, so run the proxy from your branch (an editable install serves the UI your checkout builds) -Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and point `PHOENIX_COLLECTOR_HTTP_ENDPOINT` at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy +Some suites need extra services the bare proxy does not start. The `logging/` OTEL trace-completeness tests read spans back from a jaeger query API at `http://localhost:16686` (override with `E2E_OTEL_QUERY_URL`); run a `jaegertracing/all-in-one` and point `PHOENIX_COLLECTOR_HTTP_ENDPOINT` at its OTLP ingest. The `mcp/` suite needs the deterministic upstream MCP server in `mcp_tests/mcp_e2e_upstream_server.py` reachable by the proxy. The presidio guardrail tests need a running Presidio analyzer and anonymizer the proxy can reach, addressed by `PRESIDIO_ANALYZER_API_BASE` / `PRESIDIO_ANONYMIZER_API_BASE` + +A couple of logging destinations are configured on the proxy rather than by the test. The Weave tests scope their callback to the key they create, but litellm builds the `weave_otel` logger from `WANDB_API_KEY` and `WANDB_PROJECT_ID` before it applies the per-key vars, so the proxy needs both in its own environment or the key-scoped callback never initializes and nothing ships ### Record and replay diff --git a/tests/e2e/guardrails/guardrails_client.py b/tests/e2e/guardrails/guardrails_client.py index f03e70df84a..1f55a0f9a56 100644 --- a/tests/e2e/guardrails/guardrails_client.py +++ b/tests/e2e/guardrails/guardrails_client.py @@ -18,6 +18,7 @@ from models import ( ChatBody, ChatMessage, ChatResponse, + ChatTool, KeyGenerateBody, LiteLLMParamsBody, TeamDeleteBody, @@ -31,6 +32,8 @@ from proxy_client import ProxyClient from pydantic import BaseModel GuardrailMode = Literal["pre_call", "post_call", "during_call", "logging_only"] +PiiEntity = Literal["EMAIL_ADDRESS", "PHONE_NUMBER", "PERSON", "CREDIT_CARD", "US_SSN"] +PiiAction = Literal["MASK", "BLOCK"] BlockedWordAction = Literal["BLOCK", "MASK"] @@ -81,6 +84,27 @@ class PresidioParamsBody(GuardrailParamsBase): presidio_filter_scope: Literal["input", "output", "both"] | None = None presidio_language: str | None = None output_parse_pii: bool | None = None + pii_entities_config: dict[PiiEntity, PiiAction] | None = None + + +class ToolPermissionRuleBody(BaseModel): + """One tool_permission rule: a decision for the tool named by `tool_name`.""" + + id: str + tool_name: str + decision: Literal["allow", "deny"] + + +class ToolPermissionParamsBody(GuardrailParamsBase): + """Tool-permission guardrail params. `default_action="deny"` makes the rules + an allow-list, and `on_disallowed_action="block"` turns a disallowed tool into + a 400 instead of rewriting the request; "rewrite" is a different product + promise and belongs to its own scenario.""" + + guardrail: Literal["tool_permission"] = "tool_permission" + rules: list[ToolPermissionRuleBody] + default_action: Literal["allow", "deny"] = "deny" + on_disallowed_action: Literal["block", "rewrite"] = "block" GuardrailParamsBody = ( @@ -89,6 +113,7 @@ GuardrailParamsBody = ( | OpenAIModerationParamsBody | BlockCodeExecutionParamsBody | PresidioParamsBody + | ToolPermissionParamsBody ) @@ -200,9 +225,7 @@ class GuardrailsClient: self.proxy.transport.post( "/guardrails", headers=self.proxy.transport.master, - json=GuardrailCreateBody( - guardrail=GuardrailSpecBody(guardrail_name=name, litellm_params=params) - ), + json=GuardrailCreateBody(guardrail=GuardrailSpecBody(guardrail_name=name, litellm_params=params)), response_type=GuardrailCreateResponse, ) ).guardrail_id @@ -241,9 +264,7 @@ class GuardrailsClient: ) def create_key_in_team(self, team_id: str) -> str: - return self.proxy.generate_key( - KeyGenerateBody(team_id=team_id, user_id="e2e-guardrails-user") - ) + return self.proxy.generate_key(KeyGenerateBody(team_id=team_id, user_id="e2e-guardrails-user")) def chat( self, @@ -253,6 +274,7 @@ class GuardrailsClient: *, guardrails: list[str] | None = None, max_tokens: int = 16, + tools: list[ChatTool] | None = None, ) -> Result[ChatResponse]: """Drive a chat call, optionally opting into named guardrails for this request only (the per-request `guardrails` selector). With `guardrails` @@ -266,6 +288,35 @@ class GuardrailsClient: messages=[ChatMessage(role="user", content=text)], max_tokens=max_tokens, guardrails=guardrails, + tools=tools, + ), + ) + + def chat_raw( + self, + key: str, + model: str, + text: str, + *, + guardrails: list[str] | None = None, + max_tokens: int = 16, + tools: list[ChatTool] | None = None, + tool_choice: str | None = None, + ) -> StreamingResponse: + """Drive /chat/completions returning the raw HTTP outcome, for the + assertions a typed body cannot carry: the `x-litellm-applied-guardrails` + response header, which is how an ALLOW scenario proves the guardrail ran + rather than being absent.""" + return self.proxy.transport.send( + "/chat/completions", + headers=self.proxy.transport.bearer(key), + json=ChatBody( + model=model, + messages=[ChatMessage(role="user", content=text)], + max_tokens=max_tokens, + guardrails=guardrails, + tools=tools, + tool_choice=tool_choice, ), ) @@ -323,9 +374,7 @@ class GuardrailsClient: return self.proxy.transport.send( "/v1/responses", headers=self.proxy.transport.bearer(key), - json=_ResponsesGuardrailBody( - model=model, input=text, guardrails=guardrails - ), + json=_ResponsesGuardrailBody(model=model, input=text, guardrails=guardrails), ) def apply_guardrail(self, key: str, *, name: str, text: str) -> Result[ApplyGuardrailResponse]: @@ -349,9 +398,7 @@ class GuardrailsClient: if isinstance(last, Success): return time.sleep(POLL_INTERVAL) - raise AssertionError( - f"team {team_id!r} was created but /team/info never returned it: {last}" - ) + raise AssertionError(f"team {team_id!r} was created but /team/info never returned it: {last}") def build_client(proxy: ProxyClient) -> GuardrailsClient: diff --git a/tests/e2e/guardrails/test_presidio_masking_e2e.py b/tests/e2e/guardrails/test_presidio_masking_e2e.py index 6d927292975..c6d87473c21 100644 --- a/tests/e2e/guardrails/test_presidio_masking_e2e.py +++ b/tests/e2e/guardrails/test_presidio_masking_e2e.py @@ -6,11 +6,19 @@ messages BEFORE the model runs, so the model only ever sees placeholders like must come back with the placeholders echoed and the raw PII absent, on /chat/completions and on /v1/messages (Anthropic format). +post_call: the mirror hook. The request reaches the model unmasked and the +MODEL OUTPUT is what gets anonymized, so the caller never receives raw PII the +model repeated back. The two hooks are told apart behaviorally rather than by +configuration: the post_call prompt asks for a value derived from the raw email +(its local part, which is not itself an entity Presidio masks) alongside the +address itself, so the answer proves the model saw the raw address while the +address in the same response comes back as . + The analyzer/anonymizer endpoints come from PRESIDIO_ANALYZER_API_BASE / PRESIDIO_ANONYMIZER_API_BASE; missing env is a hard failure, never a skip. -Each guardrail registers with presidio_filter_scope="input" so only the -configured hook's callback exists (the default "both" adds a second post_call -output masker), and is deleted on teardown. +Each guardrail registers with an explicit presidio_filter_scope so only the +configured hook's callback exists (the default "both" registers input masking +AND a post_call output masker), and is deleted on teardown. """ from __future__ import annotations @@ -18,13 +26,14 @@ from __future__ import annotations import os import time from collections.abc import Callable +from typing import Literal import pytest from pydantic import BaseModel from e2e_config import unique_marker from e2e_http import Result, Success -from guardrails_client import GuardrailsClient, PresidioParamsBody +from guardrails_client import GuardrailMode, GuardrailsClient, PiiAction, PiiEntity, PresidioParamsBody from lifecycle import ResourceManager from models import AnthropicMessagesResponse, ChatResponse @@ -65,16 +74,20 @@ def _register_presidio( resources: ResourceManager, *, name: str, + mode: GuardrailMode = "pre_call", + filter_scope: Literal["input", "output", "both"] = "input", + entities: dict[PiiEntity, PiiAction] | None = None, ) -> None: analyzer, anonymizer = _presidio_bases() guardrail_id = client.register( name, PresidioParamsBody( - mode="pre_call", + mode=mode, default_on=False, presidio_analyzer_api_base=analyzer, presidio_anonymizer_api_base=anonymizer, - presidio_filter_scope="input", + presidio_filter_scope=filter_scope, + pii_entities_config=entities, ), ) resources.defer(lambda: client.delete_guardrail(guardrail_id)) @@ -182,3 +195,84 @@ class TestPresidioPreCallMasking: _messages_text, email=email, ) + + +#: Room for the model's reasoning tokens plus the three-line answer; a lower cap +#: truncates the response before the address it is supposed to mask. +_POST_CALL_MAX_TOKENS = 512 + +#: The post_call scenario masks these two entities and nothing else. Left +#: unscoped, Presidio's broader recognizers claim the local part too (a random +#: marker reads as an NRP), which would erase the very token that tells output +#: masking apart from input masking. +_POST_CALL_ENTITIES: dict[PiiEntity, PiiAction] = {"EMAIL_ADDRESS": "MASK", "PHONE_NUMBER": "MASK"} + + +def _post_call_prompt(marker: str, local_part: str) -> str: + """Ask for the local part and the full address in one answer. Presidio masks + an EMAIL_ADDRESS entity and a bare local part is not one, so the two land + differently in the same response and pin the hook point behaviorally.""" + return ( + f"{marker} My email address is {local_part}@example.com and my phone number is {FAKE_PHONE}. " + "Reply with exactly three lines and nothing else. " + "Line 1: the part of the email address before the @ sign. " + "Line 2: the full email address. " + "Line 3: the phone number." + ) + + +class TestPresidioPostCallMasking: + @pytest.mark.covers( + "guardrail.presidio.post_call.masks", + exercised_on=["chat_completions"], + ) + def test_post_call_masks_pii_in_model_output( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + """A guardrail scoped to the output must anonymize the PII the model + repeats back, so a caller (or a downstream log of the response) never + receives it, while the request itself reaches the model untouched. + + Both facts are asserted from one response: the local part comes back raw, + which is only possible if the model saw the real address, and the address + itself comes back as in the same answer. + """ + name = f"e2e-presidio-post-chat-{unique_marker()}" + _register_presidio( + client, + resources, + name=name, + mode="post_call", + filter_scope="output", + entities=_POST_CALL_ENTITIES, + ) + + local_part = f"e2euser{unique_marker()}" + email = f"{local_part}@example.com" + prompt = _post_call_prompt(unique_marker(), local_part) + + deadline = time.monotonic() + GUARDRAIL_PROPAGATION_DEADLINE_SECONDS + last = "" + while True: + result = client.chat(scoped_key, MODEL, prompt, guardrails=[name], max_tokens=_POST_CALL_MAX_TOKENS) + match result: + case Success(data=data): + last = _first_content(data) + if MASKED_EMAIL_TOKEN in last and email not in last: + assert local_part in last, ( + "the model must have seen the RAW address (it is asked for the local " + "part, which Presidio does not mask); the local part is missing, so " + f"this response cannot tell post_call masking from pre_call: {last[:300]!r}" + ) + assert MASKED_PHONE_TOKEN in last and FAKE_PHONE not in last, ( + f"the phone number in the model's answer must be masked too, got: {last[:300]!r}" + ) + return + case _: + last = f"" + if time.monotonic() >= deadline: + pytest.fail( + f"presidio post_call guardrail never masked the model's output within " + f"{GUARDRAIL_PROPAGATION_DEADLINE_SECONDS}s; last observation: {last[:300]!r}" + ) + time.sleep(GUARDRAIL_PROPAGATION_POLL_INTERVAL_SECONDS) diff --git a/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py b/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py new file mode 100644 index 00000000000..9ef3650625c --- /dev/null +++ b/tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py @@ -0,0 +1,167 @@ +"""Live e2e: the tool_permission guardrail gates which tools a request may declare. + +The guardrail is registered `mode="pre_call"` with `default_action="deny"`, so its +rules are an allow-list applied to the tools the CALLER declares, before the model +runs. Two halves of one product promise: + +- blocks: a request declaring a tool outside the allow-list is rejected with a 400 + naming the denied tool, and never reaches the model +- allows: a request declaring only the permitted tool is served normally, comes + back with a real tool call for that tool, and carries an + `x-litellm-applied-guardrails` header naming the guardrail, which is what + separates "the guardrail ran and allowed it" from "the guardrail was never + attached". `tool_choice="required"` keeps the model from answering directly and + making the outcome depend on its mood + +No vendor API is involved: `tool_permission` is a built-in guardrail, so the +verdict comes from the proxy itself. +""" + +from __future__ import annotations + +from typing import Final + +import pytest + +from e2e_config import unique_marker +from e2e_http import StreamingResponse, UnknownApiError +from guardrails_client import ( + GuardrailsClient, + ToolPermissionParamsBody, + ToolPermissionRuleBody, + poll_until_blocked, +) +from lifecycle import ResourceManager +from models import ChatResponse, ChatTool, ChatToolFunction + +pytestmark = pytest.mark.e2e + +MODEL = "gemini-2.5-flash" + +#: The one tool the guardrail permits, and one it does not. Both are declared by +#: the caller in the request body; the guardrail reads them there. +ALLOWED_TOOL: Final = ChatTool( + function=ChatToolFunction( + name="get_weather", + description="Get the current weather for a city", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], + }, + ) +) +DENIED_TOOL: Final = ChatTool( + function=ChatToolFunction( + name="delete_customer_database", + description="Permanently delete the customer database", + parameters={"type": "object", "properties": {}}, + ) +) + +TOOL_PROMPT: Final = "What is the weather in Paris right now?" + + +def _register_tool_permission(client: GuardrailsClient, resources: ResourceManager, *, name: str) -> None: + """Allow-list exactly one tool: everything else falls to `default_action=deny` + and, with `on_disallowed_action=block`, is rejected outright.""" + guardrail_id = client.register( + name, + ToolPermissionParamsBody( + mode="pre_call", + default_on=False, + default_action="deny", + on_disallowed_action="block", + rules=[ + ToolPermissionRuleBody( + id="allow-get-weather", + tool_name=ALLOWED_TOOL.function.name, + decision="allow", + ) + ], + ), + ) + resources.defer(lambda: client.delete_guardrail(guardrail_id)) + + +def _applied_guardrails(outcome: StreamingResponse) -> str: + return outcome.headers.get("x-litellm-applied-guardrails", "") + + +def _tool_call_names(response: ChatResponse) -> tuple[str, ...]: + return tuple( + call.function.name + for choice in response.choices + if choice.message + for call in choice.message.tool_calls or () + if call.function.name + ) + + +class TestToolPermissionPreCall: + @pytest.mark.covers("guardrail.tool_permission.pre_call.blocks", exercised_on=["chat_completions"]) + def test_pre_call_blocks_tool_outside_the_allow_list( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + """A request declaring a tool the guardrail does not permit must be + rejected with a 400 that names the denied tool. An unauthorized tool that + merely reaches the model is the whole failure mode this guardrail exists + to prevent, so a 200 here is a hard failure.""" + name = f"e2e-toolperm-block-{unique_marker()}" + _register_tool_permission(client, resources, name=name) + + result = poll_until_blocked( + lambda: client.chat( + scoped_key, + MODEL, + TOOL_PROMPT, + guardrails=[name], + max_tokens=128, + tools=[DENIED_TOOL], + ) + ) + + match result: + case UnknownApiError(status_code=status, body=body): + assert status == 400, f"expected the guardrail block status 400, got {status}: {body[:400]}" + assert DENIED_TOOL.function.name in body, ( + f"the block must name the denied tool so the caller can fix the request; got: {body[:400]}" + ) + assert "guardrail" in body.lower(), ( + f"the block body should identify itself as a guardrail verdict; got: {body[:400]}" + ) + case _: + pytest.fail(f"tool_permission let a tool outside the allow-list through; got {result}") + + @pytest.mark.covers("guardrail.tool_permission.pre_call.allows", exercised_on=["chat_completions"]) + def test_pre_call_allows_permitted_tool( + self, client: GuardrailsClient, resources: ResourceManager, scoped_key: str + ) -> None: + """The mirror half: a request declaring only the permitted tool is served + and the model calls it. Without the header check a guardrail that never + attached would pass this test for the wrong reason, so the 200 alone is + not the contract.""" + name = f"e2e-toolperm-allow-{unique_marker()}" + _register_tool_permission(client, resources, name=name) + + outcome = client.chat_raw( + scoped_key, + MODEL, + TOOL_PROMPT, + guardrails=[name], + max_tokens=128, + tools=[ALLOWED_TOOL], + tool_choice="required", + ) + + assert outcome.ok, f"the permitted tool must be served, got {outcome.status_code}: {outcome.body[:400]}" + applied = _applied_guardrails(outcome) + assert name in applied, ( + "the allowed call must carry x-litellm-applied-guardrails naming the guardrail; " + f"without it the 200 only proves the guardrail never ran. Got {applied!r}" + ) + + called = _tool_call_names(ChatResponse.model_validate_json(outcome.body)) + assert called == (ALLOWED_TOOL.function.name,), ( + f"the served call must carry one tool call for the permitted tool, got {called!r}: {outcome.body[:400]}" + ) diff --git a/tests/e2e/logging/logging_client.py b/tests/e2e/logging/logging_client.py index f0f7ad7eaa4..66dfa233ec4 100644 --- a/tests/e2e/logging/logging_client.py +++ b/tests/e2e/logging/logging_client.py @@ -189,6 +189,45 @@ class LangfuseCreds: ) +@dataclass(frozen=True, slots=True) +class WeaveCreds: + """Weights & Biases Weave credentials for a key-scoped ``weave_otel`` callback. + + The proxy still needs WANDB_API_KEY / WANDB_PROJECT_ID in its own environment: + the weave_otel logger is constructed from those before the per-key vars are + applied, so a key-scoped callback on a proxy without them never initializes. + The per-key vars are what direct THIS key's spans at this project. + """ + + api_key: str + project_id: str + + def key_logging_metadata(self) -> KeyMetadata: + return KeyMetadata( + logging=[ + KeyLoggingCallback( + callback_name="weave_otel", + callback_type="success_and_failure", + callback_vars=KeyLoggingCallbackVars( + wandb_api_key=self.api_key, + weave_project_id=self.project_id, + ), + ) + ] + ) + + +def load_weave_creds() -> WeaveCreds: + api_key = os.getenv("WANDB_API_KEY") + project_id = (os.getenv("WEAVE_PROJECT_ID") or os.getenv("WANDB_PROJECT_ID") or "").strip() + if not (api_key and project_id): + pytest.fail( + "Weave e2e requires WANDB_API_KEY and WEAVE_PROJECT_ID (or WANDB_PROJECT_ID, " + "format /); missing credentials is a hard failure, not a skip" + ) + return WeaveCreds(api_key=api_key, project_id=project_id) + + def load_langfuse_creds() -> LangfuseCreds: public_key = os.getenv("LANGFUSE_PUBLIC_KEY") secret_key = os.getenv("LANGFUSE_SECRET_KEY") diff --git a/tests/e2e/logging/test_weave_log_e2e.py b/tests/e2e/logging/test_weave_log_e2e.py new file mode 100644 index 00000000000..dab5993c87a --- /dev/null +++ b/tests/e2e/logging/test_weave_log_e2e.py @@ -0,0 +1,192 @@ +"""Live e2e: key-scoped Weave (Weights & Biases) delivery, success and failure. + +Covers the two `logging.niche_integrations.*.logs_spend` cells with a real member +of that cohort. A key carrying a `weave_otel` callback in its logging metadata +must deliver its calls to the real Weave project, and each call must arrive +exactly once, carrying the same cost the response header reported: + +- success: one `litellm_request` call, OTEL status OK, `llm.response.cost` equal + to `x-litellm-response-cost`, and non-zero tokens +- failure: a provider-rejected call arrives too, as one call with OTEL status + ERROR naming the provider exception, and with no cost - a failed call that + silently never reaches the destination is an invisible outage, and a billed + one is worse + +Both halves assert the recorded state (the key's callback registration answers +success and the destination holds the call) and the enforced behavior (the +delivered payload's status and cost). Delivery is read back through Weave's own +query API; nothing is mocked. +""" + +from __future__ import annotations + +import time + +import pytest + +from e2e_config import CHEAP_ANTHROPIC_MODEL, unique_marker +from e2e_http import StreamingResponse +from lifecycle import ResourceManager +from logging_client import ( + INVALID_UPSTREAM_API_KEY, + LoggingClient, + WeaveCreds, + costs_agree, + first_ok, + load_weave_creds, +) +from models import LiteLLMParamsBody +from weave_reader import WeaveCall, WeaveReader, build_weave_reader + +pytestmark = pytest.mark.e2e + + +@pytest.fixture(scope="session") +def weave_creds() -> WeaveCreds: + return load_weave_creds() + + +@pytest.fixture(scope="session") +def weave_reader() -> WeaveReader: + return build_weave_reader() + + +#: How far before the request the Weave read-back window opens, to absorb clock +#: skew between this host and Weave. Without it a host running slightly fast +#: would filter out its own call. +_WINDOW_SKEW_SECONDS = 120.0 + + +def _window_start() -> float: + return time.time() - _WINDOW_SKEW_SECONDS + + +def _exactly_one(calls: tuple[WeaveCall, ...], *, marker: str, what: str) -> WeaveCall: + assert calls, f"no Weave call for the {what} (marker {marker}) reached the project within the deadline" + assert len(calls) == 1, ( + f"expected exactly ONE Weave call for the {what} (marker {marker}), got {len(calls)}: " + f"{[call.id for call in calls]} - more than one call for one request is the " + "duplicate-delivery bug" + ) + return calls[0] + + +WEAVE_STAGE_RED_REASON = ( + "stage red: product gap, key-scoped weave_otel spans are not delivered when the OTEL v2 callback is active" +) + + +class TestWeaveLogDelivery: + @pytest.mark.skip(reason=WEAVE_STAGE_RED_REASON) + @pytest.mark.covers("logging.niche_integrations.success.logs_spend", exercised_on=["chat_completions"]) + def test_chat_completions_delivers_one_call_with_spend( + self, + client: LoggingClient, + weave_creds: WeaveCreds, + weave_reader: WeaveReader, + resources: ResourceManager, + ) -> None: + alias = f"weave-key-{unique_marker()}" + key = client.key_with_alias( + alias, + models=[CHEAP_ANTHROPIC_MODEL], + metadata=weave_creds.key_logging_metadata(), + ) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + since = _window_start() + outcome = first_ok( + client, + lambda: client.chat_raw(key, CHEAP_ANTHROPIC_MODEL, f"reply with one word {marker}", max_tokens=64), + ) + assert outcome.response_cost is not None and outcome.response_cost > 0, ( + f"the response must report x-litellm-response-cost, got {outcome.response_cost!r}" + ) + + call = _exactly_one( + weave_reader.poll_calls_matching(marker, since=since), marker=marker, what="successful call" + ) + + assert call.status_code == "OK", f"a successful call must land at OK span status, got {call.status_code!r}" + cost = call.response_cost + assert cost is not None and costs_agree(outcome.response_cost, cost), ( + f"the Weave call's llm.response.cost {cost!r} must agree with the header cost " + f"{outcome.response_cost} - a delivered span with the wrong cost is a silent " + "billing-attribution bug" + ) + assert call.total_tokens is not None and call.total_tokens > 0, ( + f"the delivered call must carry token usage, got {call.total_tokens!r}" + ) + + @pytest.mark.skip(reason=WEAVE_STAGE_RED_REASON) + @pytest.mark.covers("logging.niche_integrations.failure.logs_spend", exercised_on=["chat_completions"]) + def test_failed_chat_completions_delivers_one_error_call( + self, + client: LoggingClient, + weave_creds: WeaveCreds, + weave_reader: WeaveReader, + resources: ResourceManager, + ) -> None: + """A deployment with an invalid upstream key passes proxy auth and fails + at the provider, so exactly one provider failure exists for it. Proxy-side + 401s during key propagation never reach the provider and ship no payload, + which is what the retry loop below relies on.""" + model_name = f"weave-err-{unique_marker()}" + model_id = client.create_model( + model_name, + LiteLLMParamsBody(model="anthropic/claude-haiku-4-5", api_key=INVALID_UPSTREAM_API_KEY), + ) + resources.defer(lambda: client.delete_model(model_id)) + key = client.key_with_alias( + f"weave-err-key-{unique_marker()}", + models=[model_name], + metadata=weave_creds.key_logging_metadata(), + ) + resources.defer(lambda: client.delete_key(key)) + + marker = unique_marker() + since = _window_start() + outcome = _provoke_provider_failure(client, key, model_name, marker) + + call = _exactly_one(weave_reader.poll_calls_matching(marker, since=since), marker=marker, what="failed call") + + assert call.status_code == "ERROR", ( + f"a failed call must land at ERROR span status, got {call.status_code!r} - " + "Weave's own summary.weave.status reads success either way, which is exactly " + "why the span status is what this asserts on" + ) + error = call.error + assert error is not None and error.message is not None and "AnthropicException" in error.message, ( + f"the delivered call must carry the provider error, got {error!r}" + ) + assert not call.response_cost, f"a failed call must not be billed, got llm.response.cost={call.response_cost!r}" + assert outcome.status_code == 401, ( + f"an upstream auth failure must map to 401, got {outcome.status_code}: {outcome.body[:200]}" + ) + + +def _provoke_provider_failure(client: LoggingClient, key: str, model_name: str, marker: str) -> StreamingResponse: + """Send until the provider (not the proxy) is the one rejecting the call. + + A network failure between the test and the proxy is NOT retried: the request + may have been served, and a retry would double-log the failure payload and + falsely trip the exactly-one assertion. + """ + deadline = time.monotonic() + client.proxy.poll_timeout + while True: + outcome = client.chat_raw(key, model_name, f"trigger an upstream auth failure {marker}", max_tokens=16) + assert not outcome.ok, "the call must fail; the deployment's upstream key is invalid" + assert outcome.status_code != -1, ( + "network failure between the test and the proxy while provoking the provider failure; " + "retrying now could double-log the failure payload and falsely trip the exactly-one " + f"assertion - fix the rig connectivity first: {outcome.body[:200]}" + ) + if "AnthropicException" in outcome.body or time.monotonic() >= deadline: + break + time.sleep(client.proxy.poll_interval) + assert "AnthropicException" in outcome.body, ( + "never saw the upstream provider failure before the deadline; the key may still be " + f"propagating - last outcome {outcome.status_code}: {outcome.body[:200]}" + ) + return outcome diff --git a/tests/e2e/logging/weave_reader.py b/tests/e2e/logging/weave_reader.py new file mode 100644 index 00000000000..2f8f759d299 --- /dev/null +++ b/tests/e2e/logging/weave_reader.py @@ -0,0 +1,282 @@ +"""Read-back for the Weave (Weights & Biases) logging tests against the real +Weave project. + +The proxy ships OTEL spans to https://trace.wandb.ai/otel/v1/traces with the +``weave_otel`` callback, and the tests read the ingested calls back through +Weave's own query API (``POST /calls/stream_query``), which answers JSON Lines: +one JSON object per call, so the body is parsed line by line rather than as one +document. + +The project is shared with other traffic, so the read never relies on the target +being among the newest N calls: the query is scoped server-side to the +``litellm_request`` op and to calls that started after the test's own request, +and pages with ``offset`` until the window is exhausted. + +Weave's own ``summary.weave.status`` is a rollup that reads "success" even for a +span the exporter marked failed, so status comes from the OTEL span itself +(``attributes.otel_span.status.code``), and the shipped cost from +``attributes.otel_span.attributes.llm.response.cost`` - the StandardLogging +``response_cost``, which is what makes this a spend assertion rather than a +delivery ping. + +Missing configuration is a hard failure, never a skip. +""" + +from __future__ import annotations + +import base64 +import json +import os +import time +from dataclasses import dataclass +from itertools import count, takewhile +from typing import Final + +import pytest +from pydantic import BaseModel, ConfigDict, Field + +from e2e_config import POLL_INTERVAL, POLL_TIMEOUT +from e2e_http import URL, AuthHeaders, send + +_WEAVE_TRACE_API: Final = "https://trace.wandb.ai" + +#: The op every litellm LLM call lands under. The proxy also exports a root +#: server span ("Received Proxy Server Request") and management spans; only the +#: LLM call carries the usage and cost this suite asserts on. +LITELLM_REQUEST_OP: Final = "litellm_request" + +#: How long to keep re-reading after the first matching call before trusting the +#: exactly-one assertion. The OTEL batch exporter flushes on its own schedule, so +#: a duplicate export can surface well after the first one, and a duplicate IS +#: the bug being guarded against. +WEAVE_SETTLE_SECONDS: Final = 45.0 + +#: Rows per page. The query is already scoped to this run's time window, so this +#: only bounds one round trip, not what the read can see. +_PAGE_SIZE: Final = 500 + + +class _WeaveSortBy(BaseModel): + field: str + direction: str + + +class _WeaveOpFilter(BaseModel): + op_names: list[str] + + +class _WeaveGetField(BaseModel): + get_field: str = Field(serialization_alias="$getField") + + +class _WeaveLiteral(BaseModel): + literal: float = Field(serialization_alias="$literal") + + +class _WeaveGreaterThan(BaseModel): + gt: tuple[_WeaveGetField, _WeaveLiteral] = Field(serialization_alias="$gt") + + +class _WeaveQuery(BaseModel): + expr: _WeaveGreaterThan = Field(serialization_alias="$expr") + + +class _WeaveQueryBody(BaseModel): + project_id: str + filter: _WeaveOpFilter + query: _WeaveQuery + limit: int = _PAGE_SIZE + offset: int = 0 + sort_by: list[_WeaveSortBy] = [_WeaveSortBy(field="started_at", direction="asc")] + + +class _OtelStatus(BaseModel): + model_config = ConfigDict(extra="ignore") + + code: str | None = None + message: str | None = None + + +class _OtelError(BaseModel): + model_config = ConfigDict(extra="ignore") + + code: str | None = None + type: str | None = None + message: str | None = None + + +class _LlmResponse(BaseModel): + model_config = ConfigDict(extra="ignore") + + cost: float | None = None + + +class _LlmAttributes(BaseModel): + model_config = ConfigDict(extra="ignore") + + response: _LlmResponse | None = None + + +class _OtelSpanAttributes(BaseModel): + model_config = ConfigDict(extra="ignore") + + llm: _LlmAttributes | None = None + error: _OtelError | None = None + + +class _OtelSpan(BaseModel): + model_config = ConfigDict(extra="ignore") + + name: str | None = None + status: _OtelStatus | None = None + attributes: _OtelSpanAttributes | None = None + + +class _CallAttributes(BaseModel): + model_config = ConfigDict(extra="ignore") + + otel_span: _OtelSpan | None = None + + +class _Usage(BaseModel): + model_config = ConfigDict(extra="ignore") + + total_tokens: int | None = None + + +class _WeaveSummary(BaseModel): + model_config = ConfigDict(extra="ignore") + + usage: dict[str, _Usage] = {} + + +class WeaveCall(BaseModel): + """One ingested Weave call, reduced to what the scenarios assert on.""" + + model_config = ConfigDict(extra="ignore") + + id: str + op_name: str + started_at: str | None = None + inputs: dict[str, object] = {} + attributes: _CallAttributes | None = None + summary: _WeaveSummary | None = Field(default=None) + + @property + def op(self) -> str: + """The bare op name out of ``weave://///op/:``.""" + return self.op_name.split("/op/")[-1].split(":")[0] + + @property + def status_code(self) -> str | None: + """The OTEL span status, not Weave's own rollup (which reads "success" + even for a span the exporter marked ERROR).""" + span = self.attributes.otel_span if self.attributes else None + return span.status.code if span and span.status else None + + @property + def error(self) -> _OtelError | None: + span = self.attributes.otel_span if self.attributes else None + return span.attributes.error if span and span.attributes else None + + @property + def response_cost(self) -> float | None: + span = self.attributes.otel_span if self.attributes else None + llm = span.attributes.llm if span and span.attributes else None + return llm.response.cost if llm and llm.response else None + + @property + def total_tokens(self) -> int | None: + """Weave keys usage by model, so the total is summed across whatever + models the call reported.""" + if not self.summary or not self.summary.usage: + return None + totals = [usage.total_tokens for usage in self.summary.usage.values() if usage.total_tokens is not None] + return sum(totals) if totals else None + + def mentions(self, needle: str) -> bool: + return needle in json.dumps(self.inputs, default=str) + + +@dataclass(frozen=True, slots=True) +class WeaveReader: + project_id: str + api_key: str + + @property + def _headers(self) -> AuthHeaders: + """Weave authenticates with HTTP Basic as the fixed user ``api``.""" + token = base64.b64encode(f"api:{self.api_key}".encode()).decode() + return AuthHeaders(authorization=f"Basic {token}") + + def _query_body(self, *, since: float, offset: int, op: str) -> _WeaveQueryBody: + return _WeaveQueryBody( + project_id=self.project_id, + filter=_WeaveOpFilter(op_names=[f"weave:///{self.project_id}/op/{op}:*"]), + query=_WeaveQuery( + expr=_WeaveGreaterThan(gt=(_WeaveGetField(get_field="started_at"), _WeaveLiteral(literal=since))) + ), + offset=offset, + ) + + def _page(self, *, since: float, offset: int, op: str) -> tuple[WeaveCall, ...]: + outcome = send( + URL(f"{_WEAVE_TRACE_API}/calls/stream_query"), + headers=self._headers, + json=self._query_body(since=since, offset=offset, op=op), + ) + if not outcome.ok: + pytest.fail( + f"Weave calls query for project {self.project_id!r} failed " + f"({outcome.status_code}): {outcome.body[:300]}" + ) + return tuple(WeaveCall.model_validate_json(line) for line in outcome.body.splitlines() if line.strip()) + + def calls_matching(self, marker: str, *, since: float, op: str = LITELLM_REQUEST_OP) -> tuple[WeaveCall, ...]: + """Every call under ``op`` started after ``since`` whose inputs carry + ``marker``, paging until the window is exhausted. + + More than one is the duplicate-delivery bug, so this never collapses to a + single call. + """ + pages = tuple( + takewhile( + bool, + (self._page(since=since, offset=offset, op=op) for offset in count(0, _PAGE_SIZE)), + ) + ) + return tuple(call for page in pages for call in page if call.mentions(marker)) + + def poll_calls_matching(self, marker: str, *, since: float, op: str = LITELLM_REQUEST_OP) -> tuple[WeaveCall, ...]: + """Poll until the call is readable, then keep re-reading for + WEAVE_SETTLE_SECONDS so a duplicate exported by a later batch flush + cannot hide from the exactly-one assertion. A duplicate ends the settle + early, because more waiting cannot clear it.""" + deadline = time.monotonic() + POLL_TIMEOUT + while time.monotonic() < deadline: + calls = self.calls_matching(marker, since=since, op=op) + if calls: + return self._settled(marker, since=since, op=op, first=calls) + time.sleep(POLL_INTERVAL) + return () + + def _settled(self, marker: str, *, since: float, op: str, first: tuple[WeaveCall, ...]) -> tuple[WeaveCall, ...]: + """A transiently empty re-read never downgrades what was already seen.""" + settle_deadline = time.monotonic() + WEAVE_SETTLE_SECONDS + latest = first # rebind-ok: one settle window, re-read per poll interval + while time.monotonic() < settle_deadline and len(latest) <= 1: + time.sleep(POLL_INTERVAL) + latest = self.calls_matching(marker, since=since, op=op) or latest + return latest + + +def build_weave_reader() -> WeaveReader: + project_id = (os.environ.get("WEAVE_PROJECT_ID") or os.environ.get("WANDB_PROJECT_ID") or "").strip() + api_key = os.environ.get("WANDB_API_KEY", "").strip() + if not project_id or not api_key: + pytest.fail( + "Weave e2e requires WANDB_API_KEY and WEAVE_PROJECT_ID (or WANDB_PROJECT_ID, " + "format /): the test reads the proxy's weave_otel delivery " + "back from the real Weave project; missing credentials is a hard failure, not a skip" + ) + return WeaveReader(project_id=project_id, api_key=api_key) diff --git a/tests/e2e/models.py b/tests/e2e/models.py index 5de49ead3ed..1379ecb4530 100644 --- a/tests/e2e/models.py +++ b/tests/e2e/models.py @@ -35,6 +35,8 @@ class KeyLoggingCallbackVars(BaseModel): langfuse_public_key: str | None = None langfuse_secret_key: str | None = None langfuse_host: str | None = None + wandb_api_key: str | None = None + weave_project_id: str | None = None class KeyLoggingCallback(BaseModel): From 8544faec91c398519b305dbf7850de87ec999486 Mon Sep 17 00:00:00 2001 From: "cursor[bot]" <206951365+cursor[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 13:04:31 -0700 Subject: [PATCH 3/7] fix(ci): grant pull_requests write for release wheel reporter (#39922) * fix(ci): grant pull_requests write for release wheel reporter The reporter posts a PR comment via github.rest.issues.createComment. GitHub requires both issues=write and pull_requests=write to comment on a PR issue, as returned in x-accepted-github-permissions. The workflow had pull-requests: read, so the POST failed with 403 'Resource not accessible by integration'. Bumping to pull-requests: write fixes the create path; the read-only pulls.get call still works. Same-repo scope is preserved by the existing head_repository.full_name check. Co-authored-by: Krrish Dholakia * fix(ci): scope release wheel reporter permissions to pull requests --------- Co-authored-by: Cursor Agent Co-authored-by: Krrish Dholakia Co-authored-by: Yujong Lee --- .github/workflows/report-rust-release-wheel.yml | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/.github/workflows/report-rust-release-wheel.yml b/.github/workflows/report-rust-release-wheel.yml index 1d93b56f77f..74e6be69604 100644 --- a/.github/workflows/report-rust-release-wheel.yml +++ b/.github/workflows/report-rust-release-wheel.yml @@ -24,8 +24,7 @@ jobs: runs-on: ubuntu-latest timeout-minutes: 5 permissions: - issues: write # PR comments use the issues API - pull-requests: read # Current-head validation rejects stale workflow runs + pull-requests: write steps: - name: Link release wheel report on PR From 5df0e12e0f2628ed847c8110759a589f3fa1c138 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 13:08:03 -0700 Subject: [PATCH 4/7] feat(guardrails): add non-blocking flag() verdict to custom code guardrails (#39728) Custom code guardrails could only allow(), block(reason) or modify(). This adds flag(reason, metadata={}) which lets the request or response through unchanged and records a guardrail_flagged entry carrying the guardrail name, configured mode, evaluated input_type (request or response), reason and structured metadata. The new status is threaded through the request-level guardrail_status aggregation, the Guardrails Monitor rollup (flagged_count), Request Logs (action=flagged, most severe phase wins when a guardrail runs pre and post call) and the Request Logs detail view in the dashboard, which now renders FLAGGED with warning styling instead of falling into FAILED. Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/litellm_core_utils/litellm_logging.py | 2 + .../custom_code/custom_code_guardrail.py | 27 ++++++ .../guardrail_hooks/custom_code/primitives.py | 28 +++++- litellm/proxy/guardrails/usage_endpoints.py | 18 ++-- litellm/proxy/guardrails/usage_tracking.py | 6 +- litellm/types/utils.py | 5 +- .../test_litellm_logging.py | 18 +++- .../guardrails/test_custom_code_security.py | 65 +++++++++++++ .../proxy/guardrails/test_usage_endpoints.py | 97 +++++++++++++++++++ .../proxy/guardrails/test_usage_tracking.py | 21 ++++ .../custom_code/CustomCodeModal.tsx | 1 + .../GuardrailViewer/GuardrailViewer.test.tsx | 16 +++ .../GuardrailViewer/GuardrailViewer.tsx | 95 ++++++++++++------ .../LogDetailContent.test.tsx | 16 ++- .../LogDetailsDrawer/LogDetailContent.tsx | 31 +++--- 15 files changed, 386 insertions(+), 60 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 83e0b4d84f1..c31c4323157 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5881,6 +5881,7 @@ def _get_status_fields( # Mapping for legacy guardrail status values to new GuardrailStatus values GUARDRAIL_STATUS_MAP: Final[dict[str, GuardrailStatus]] = { "success": "success", + "guardrail_flagged": "guardrail_flagged", "blocked": "guardrail_intervened", # legacy "guardrail_intervened": "guardrail_intervened", # direct "failure": "guardrail_failed_to_respond", # legacy @@ -5902,6 +5903,7 @@ def _get_status_fields( GUARDRAIL_STATUS_SEVERITY: Final[tuple[GuardrailStatus, ...]] = ( "not_run", "success", + "guardrail_flagged", "guardrail_failed_to_respond", "guardrail_intervened", ) diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py index 830dec8d80d..d5ef1e949b8 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/custom_code_guardrail.py @@ -36,6 +36,7 @@ Example: block when response rejects the user (input_type response only): import asyncio import threading +import time from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast @@ -93,6 +94,7 @@ class CustomCodeGuardrail(CustomGuardrail): that returns one of: - allow() - let the request/response through - block(reason) - reject with a message + - flag(reason) - let it through but log a non-blocking violation - modify(texts=...) - transform the content Example: @@ -227,6 +229,7 @@ class CustomCodeGuardrail(CustomGuardrail): raise CustomCodeExecutionError(f"Custom code guardrail not compiled: {self._compile_error}") raise CustomCodeExecutionError("Custom code guardrail not compiled") + start_time: Final = time.time() try: # Prepare inputs dict for the function @@ -245,6 +248,7 @@ class CustomCodeGuardrail(CustomGuardrail): inputs=inputs, request_data=request_data, input_type=input_type, + start_time=start_time, ) except HTTPException: @@ -290,6 +294,7 @@ class CustomCodeGuardrail(CustomGuardrail): inputs: GenericGuardrailAPIInputs, request_data: dict[str, object], input_type: Literal["request", "response"], + start_time: float, ) -> GenericGuardrailAPIInputs: """ Process the result from the custom code function. @@ -299,6 +304,7 @@ class CustomCodeGuardrail(CustomGuardrail): inputs: The original inputs request_data: The request data input_type: "request" or "response" + start_time: Unix timestamp of when the guardrail started running, used for the flagged log entry Returns: GenericGuardrailAPIInputs - possibly modified @@ -348,6 +354,27 @@ class CustomCodeGuardrail(CustomGuardrail): }, ) + elif action == "flag": + flag_reason: Final = result.get("reason", "Flagged by custom code guardrail") + verbose_proxy_logger.info( + "Custom code guardrail '%s': Flagging %s - %s", self.guardrail_name, input_type, flag_reason + ) + end_time: Final = time.time() + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={ # mutable-ok: logging helper requires a dict + "action": "flag", + "reason": flag_reason, + "input_type": input_type, + "metadata": result.get("metadata") or {}, # mutable-ok: logging helper requires a dict + }, + request_data=request_data, + guardrail_status="guardrail_flagged", + start_time=start_time, + end_time=end_time, + duration=end_time - start_time, + ) + return inputs + elif action == "modify": verbose_proxy_logger.debug("Custom code guardrail '%s': Modifying %s", self.guardrail_name, input_type) diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py index 24801aa2df1..d5dbfaeb84b 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_code/primitives.py @@ -8,7 +8,7 @@ and provide safe, sandboxed functionality for common guardrail operations. import json import re from collections.abc import Mapping, Sequence -from typing import Final +from typing import Final, Literal from urllib.parse import urlparse import httpx @@ -51,6 +51,31 @@ def block(reason: str, detection_info: Mapping[str, object] | None = None) -> di return result +class FlagResult(TypedDict): + action: ReadOnly[Literal["flag"]] + reason: ReadOnly[str] + metadata: ReadOnly[Mapping[str, object]] + + +def flag(reason: str, metadata: Mapping[str, object] | None = None) -> FlagResult: + """ + Let the request/response proceed unchanged but record a non-blocking violation. + + Args: + reason: Human-readable reason for flagging + metadata: Optional structured metadata stored alongside the reason + + Returns: + Dict indicating the request should be flagged but allowed + """ + result: Final[FlagResult] = { + "action": "flag", + "reason": reason, + "metadata": metadata if metadata is not None else {}, + } + return result + + def modify( texts: Sequence[str] | None = None, images: Sequence[object] | None = None, @@ -787,6 +812,7 @@ def get_custom_code_primitives() -> dict[str, object]: # Result types "allow": allow, "block": block, + "flag": flag, "modify": modify, # Regex "regex_match": regex_match, diff --git a/litellm/proxy/guardrails/usage_endpoints.py b/litellm/proxy/guardrails/usage_endpoints.py index 62145b9ede9..014ba3d1472 100644 --- a/litellm/proxy/guardrails/usage_endpoints.py +++ b/litellm/proxy/guardrails/usage_endpoints.py @@ -17,6 +17,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.guardrails.usage_tracking import guardrail_status_to_action from litellm.repositories.prisma_protocols import TableActions from litellm.repositories.table_repositories import ( DailyGuardrailMetricsRepository, @@ -41,6 +42,7 @@ if TYPE_CHECKING: router: Final = APIRouter() _EMPTY_UNITS: Final[Mapping[str, int]] = MappingProxyType({}) +_ACTION_SEVERITY: Final[Mapping[str, int]] = MappingProxyType({"passed": 0, "flagged": 1, "blocked": 2}) _T = TypeVar("_T") @@ -759,21 +761,17 @@ def _usage_log_entry_from_row( except Exception: meta = {} guardrail_info_list: Final[Sequence[_GuardrailRunInfo]] = (meta or {}).get("guardrail_information") or [] - entry_for_guardrail: _GuardrailRunInfo | None = None - for gi in guardrail_info_list: - if (gi.get("guardrail_id") or gi.get("guardrail_name")) == r.guardrail_id: - entry_for_guardrail = gi - break + entry_for_guardrail: Final[_GuardrailRunInfo | None] = max( + (gi for gi in guardrail_info_list if (gi.get("guardrail_id") or gi.get("guardrail_name")) == r.guardrail_id), + key=lambda gi: _ACTION_SEVERITY[guardrail_status_to_action(gi.get("guardrail_status"))], + default=None, + ) action_val = "passed" score_val = None latency_val = None reason_val = None if entry_for_guardrail: - st: Final = (entry_for_guardrail.get("guardrail_status") or "").lower() - if "intervened" in st or "block" in st: - action_val = "blocked" - elif "fail" in st or "error" in st: - action_val = "flagged" + action_val = guardrail_status_to_action(entry_for_guardrail.get("guardrail_status")) duration: Final = entry_for_guardrail.get("duration") if duration is not None: latency_val = round(float(duration) * 1000, 0) diff --git a/litellm/proxy/guardrails/usage_tracking.py b/litellm/proxy/guardrails/usage_tracking.py index a20ad3935e5..df967058cf0 100644 --- a/litellm/proxy/guardrails/usage_tracking.py +++ b/litellm/proxy/guardrails/usage_tracking.py @@ -190,14 +190,14 @@ async def _upsert_rows_with_retry( return await _upsert_rows_with_retry(retryable, upsert_row, label, sleep, retries_left - 1) -def _guardrail_status_to_action(status: str | None) -> str: +def guardrail_status_to_action(status: str | None) -> str: """Map StandardLogging guardrail_status to blocked/passed/flagged.""" if not status: return "passed" s: Final = (status or "").lower() if "intervened" in s or "block" in s: return "blocked" - if "fail" in s or "error" in s: + if "flagged" in s or "fail" in s or "error" in s: return "flagged" return "passed" @@ -367,7 +367,7 @@ async def process_spend_logs_guardrail_usage( continue key = _MetricsKey(guardrail_id, date_key) daily_guardrail[key]["requests_evaluated"] += 1 - action = _guardrail_status_to_action(entry.get("guardrail_status")) + action = guardrail_status_to_action(entry.get("guardrail_status")) if action == "passed": daily_guardrail[key]["passed_count"] += 1 elif action == "blocked": diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 2118fe77aad..61c2fc8c5a5 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3078,7 +3078,9 @@ class GuardrailMode(TypedDict, total=False): default: str | list[str] | None -GuardrailStatus = Literal["success", "guardrail_intervened", "guardrail_failed_to_respond", "not_run"] +GuardrailStatus = Literal[ + "success", "guardrail_flagged", "guardrail_intervened", "guardrail_failed_to_respond", "not_run" +] # Fields on a guardrail record whose values can quote the caller's prompt: the payload sent to the # guardrail, the provider response that echoes it back, and the two first-party hooks that inline @@ -3320,6 +3322,7 @@ class StandardLoggingPayloadStatusFields(TypedDict, total=False): """ Status of guardrail execution: - 'success': Guardrail ran and allowed content through + - 'guardrail_flagged': Guardrail allowed content through but recorded a non-blocking violation - 'guardrail_intervened': Guardrail blocked or modified content - 'guardrail_failed_to_respond': Guardrail had technical failure - 'not_run': No guardrail was run diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 1991170707d..16a99713a06 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py @@ -16,7 +16,10 @@ from litellm._logging import session_id_var, trace_id_var from litellm.constants import SENTRY_DENYLIST, SENTRY_PII_DENYLIST from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging -from litellm.litellm_core_utils.litellm_logging import set_callbacks +from litellm.litellm_core_utils.litellm_logging import ( + _get_status_fields, + set_callbacks, +) from litellm.types.utils import ModelResponse, TextCompletionResponse @@ -6441,3 +6444,16 @@ def test_passthrough_embeddings_result_swapped_for_callbacks(): assert isinstance(swapped_result, EmbeddingResponse) assert swapped_result.data[0]["embedding"] == [0.1, 0.2, 0.3] + + +def test_get_status_fields_ranks_guardrail_flagged_between_success_and_intervened(): + """LIT-6894: a non-blocking flagged verdict must outrank success in the + request-level guardrail_status but never mask an intervention.""" + flagged = {"guardrail_status": "guardrail_flagged"} + + assert _get_status_fields( + "success", [{"guardrail_status": "success"}, flagged], None + )["guardrail_status"] == "guardrail_flagged" + assert _get_status_fields( + "success", [flagged, {"guardrail_status": "guardrail_intervened"}], None + )["guardrail_status"] == "guardrail_intervened" diff --git a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py index f93ecfc3010..7971cf62c9a 100644 --- a/tests/test_litellm/proxy/guardrails/test_custom_code_security.py +++ b/tests/test_litellm/proxy/guardrails/test_custom_code_security.py @@ -197,6 +197,71 @@ async def test_custom_code_post_call_block_raises_http_400(): } +FLAG_CODE = ( + "def apply_guardrail(inputs, request_data, input_type):\n" + ' return flag("audit hit", metadata={"category": "topic"})\n' +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("input_type", ["request", "response"]) +async def test_custom_code_flag_passes_content_through_and_records_flagged_entry(input_type): + """LIT-6894: flag() must not raise, must return the content unchanged and must log + exactly one guardrail_flagged entry (the decorator must not add a second "success").""" + guardrail = CustomCodeGuardrail(custom_code=FLAG_CODE, guardrail_name="t", event_hook=["pre_call", "post_call"]) + request_data = {"model": "test-model", "litellm_metadata": {}} + + result = await guardrail.apply_guardrail( + inputs={"texts": ["hello"]}, + request_data=request_data, + input_type=input_type, + ) + + assert result == {"texts": ["hello"]} + entries = request_data["litellm_metadata"]["standard_logging_guardrail_information"] + assert len(entries) == 1 + entry = entries[0] + assert entry["guardrail_status"] == "guardrail_flagged" + assert entry["guardrail_name"] == "t" + assert entry["guardrail_mode"] == ["pre_call", "post_call"] + assert entry["guardrail_response"] == { + "action": "flag", + "reason": "audit hit", + "input_type": input_type, + "metadata": {"category": "topic"}, + } + assert entry["duration"] is not None and entry["duration"] >= 0 + + +@pytest.mark.asyncio +async def test_custom_code_flag_default_reason_and_empty_metadata(): + code = "def apply_guardrail(inputs, request_data, input_type):\n return flag('just a note')\n" + guardrail = _compile(code) + request_data = {"model": "m", "litellm_metadata": {}} + + await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data, input_type="request") + + entry = request_data["litellm_metadata"]["standard_logging_guardrail_information"][0] + assert entry["guardrail_response"] == { + "action": "flag", + "reason": "just a note", + "input_type": "request", + "metadata": {}, + } + + +@pytest.mark.asyncio +async def test_custom_code_allow_still_records_success_not_flagged(): + code = "def apply_guardrail(inputs, request_data, input_type):\n return allow()\n" + guardrail = _compile(code) + request_data = {"model": "m", "litellm_metadata": {}} + + await guardrail.apply_guardrail(inputs={"texts": ["x"]}, request_data=request_data, input_type="request") + + entries = request_data["litellm_metadata"]["standard_logging_guardrail_information"] + assert [e["guardrail_status"] for e in entries] == ["success"] + + def test_typical_sync_guardrail_still_works(): code = ( "def apply_guardrail(inputs, request_data, input_type):\n" diff --git a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py index ebb2be6edc2..4e5a7ad4b2b 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py @@ -477,6 +477,103 @@ async def test_logs_resolves_config_guardrail_logical_name(): assert where["guardrail_id"] == {"in": ["yaml-uuid", "yaml-pii"]} +def _index_row(request_id: str, guardrail_id: str = "cc-flag") -> Any: + r = MagicMock(spec=["request_id", "guardrail_id", "policy_id", "start_time"]) + r.request_id = request_id + r.guardrail_id = guardrail_id + return r + + +def _spend_log(request_id: str, *guardrail_statuses: str, guardrail_id: str = "cc-flag") -> Any: + sl = MagicMock(spec=["request_id", "metadata", "startTime", "model", "messages", "response"]) + sl.request_id = request_id + sl.startTime = datetime(2026, 4, 25, 12, 0) + sl.model = "gpt-4o-mini" + sl.messages = [{"role": "user", "content": "hi"}] + sl.response = "ok" + sl.metadata = { + "guardrail_information": [ + { + "guardrail_name": guardrail_id, + "guardrail_status": status, + "guardrail_response": ( + {"action": "flag", "reason": "audit hit"} if status == "guardrail_flagged" else "allow" + ), + "duration": 0.002, + } + for status in guardrail_statuses + ] + } + return sl + + +@pytest.mark.asyncio +async def test_logs_reports_flagged_action_for_guardrail_flagged_status(): + """LIT-6894: Request Logs surface a custom code flag() verdict as flagged with its reason.""" + prisma = _prisma(index_find_many=[_index_row("r-flag"), _index_row("r-pass"), _index_row("r-block")]) + prisma.db.litellm_spendlogs.find_many = AsyncMock( + return_value=[ + _spend_log("r-flag", "guardrail_flagged"), + _spend_log("r-pass", "success"), + _spend_log("r-block", "guardrail_intervened"), + ] + ) + p1, p2 = _patches(prisma, _config_handler()) + with p1, p2: + resp = await guardrails_usage_logs( + guardrail_id="cc-flag", + policy_id=None, + page=1, + page_size=50, + action=None, + start_date=START, + end_date=END, + user_api_key_dict=ADMIN, + ) + flagged_only = await guardrails_usage_logs( + guardrail_id="cc-flag", + policy_id=None, + page=1, + page_size=50, + action="flagged", + start_date=START, + end_date=END, + user_api_key_dict=ADMIN, + ) + assert [(log.id, log.action) for log in resp.logs] == [ + ("r-flag", "flagged"), + ("r-pass", "passed"), + ("r-block", "blocked"), + ] + assert resp.logs[0].reason == "{'action': 'flag', 'reason': 'audit hit'}" + assert [log.id for log in flagged_only.logs] == ["r-flag"] + + +@pytest.mark.asyncio +async def test_logs_reports_post_call_flag_when_pre_call_allowed(): + """LIT-6894: a guardrail on mode [pre_call, post_call] that allows the request but flags the response + shows as flagged, not hidden behind the pre_call allow entry.""" + prisma = _prisma(index_find_many=[_index_row("r-post-flag")]) + prisma.db.litellm_spendlogs.find_many = AsyncMock( + return_value=[_spend_log("r-post-flag", "success", "guardrail_flagged")] + ) + p1, p2 = _patches(prisma, _config_handler()) + with p1, p2: + resp = await guardrails_usage_logs( + guardrail_id="cc-flag", + policy_id=None, + page=1, + page_size=50, + action=None, + start_date=START, + end_date=END, + user_api_key_dict=ADMIN, + ) + assert [(log.id, log.action, log.reason) for log in resp.logs] == [ + ("r-post-flag", "flagged", "{'action': 'flag', 'reason': 'audit hit'}") + ] + + # ---- date window cap (LIT-5762) --------------------------------------------- diff --git a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py index ae360b281cb..110de7dbe70 100644 --- a/tests/test_litellm/proxy/guardrails/test_usage_tracking.py +++ b/tests/test_litellm/proxy/guardrails/test_usage_tracking.py @@ -105,6 +105,27 @@ async def test_usage_units_rolled_up_by_guardrail_team_key_and_date(): } +@pytest.mark.asyncio +async def test_flagged_status_counts_as_flagged_not_passed_or_blocked(): + """LIT-6894: a custom code flag() verdict lands in flagged_count on the Monitor rollup.""" + prisma = _prisma() + logs = [ + _payload("r1", guardrail_status="success"), + _payload("r2", guardrail_status="guardrail_flagged"), + _payload("r3", guardrail_status="guardrail_intervened"), + ] + + await process_spend_logs_guardrail_usage(prisma, logs) + + create = prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"] + assert (create["requests_evaluated"], create["passed_count"], create["flagged_count"], create["blocked_count"]) == ( + 3, + 1, + 1, + 1, + ) + + def _fake_sleep() -> tuple[AsyncMock, list[float]]: delays: list[float] = [] sleep = AsyncMock(side_effect=lambda delay: delays.append(delay)) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx index a69824f32d3..05a48598859 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails/_components/custom_code/CustomCodeModal.tsx @@ -112,6 +112,7 @@ const PRIMITIVES = { "Return Values": [ { name: "allow()", desc: "Let request/response through" }, { name: "block(reason)", desc: "Reject with message" }, + { name: "flag(reason, metadata={})", desc: "Let through, record a non-blocking violation" }, { name: "modify(texts=[], images=[], tool_calls=[])", desc: "Transform content" }, ], "HTTP Requests (async)": [ diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx index b5e04c72440..aabac50a661 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.test.tsx @@ -33,6 +33,22 @@ describe("GuardrailViewer", () => { expect(screen.getByText("1235ms")).toBeInTheDocument(); }); + it("renders guardrail_flagged as FLAGGED (warning), not FAILED", () => { + const data = makeGuardrailInformation({ + guardrail_name: "cc-flag", + guardrail_status: "guardrail_flagged", + guardrail_provider: "custom_code", + }); + renderWithProviders(); + + expect(screen.getByText(/0 Passed/)).toBeInTheDocument(); + expect(screen.getByText(/1 Flagged/)).toBeInTheDocument(); + const badges = screen.getAllByText("FLAGGED"); + expect(badges.length).toBeGreaterThan(0); + expect(badges[0]).toHaveClass("text-warning"); + expect(screen.queryByText("FAILED")).not.toBeInTheDocument(); + }); + it("calculates and displays masked entity totals", async () => { const user = userEvent.setup(); const data = makeGuardrailInformation({ diff --git a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx index 863f4117510..271b8f6ce05 100644 --- a/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/GuardrailViewer/GuardrailViewer.tsx @@ -133,8 +133,27 @@ const getTotalMasked = (entry: GuardrailInformation): number => { ); }; -const isEntrySuccess = (entry: GuardrailInformation): boolean => { - return (entry.guardrail_status ?? "").toLowerCase() === "success"; +type EntryOutcome = "passed" | "flagged" | "failed"; + +const getEntryOutcome = (entry: GuardrailInformation): EntryOutcome => { + const status = (entry.guardrail_status ?? "").toLowerCase(); + if (status === "success") return "passed"; + if (status === "guardrail_flagged") return "flagged"; + return "failed"; +}; + +const isEntrySuccess = (entry: GuardrailInformation): boolean => getEntryOutcome(entry) === "passed"; + +const OUTCOME_LABEL: Record = { + passed: "PASSED", + flagged: "FLAGGED", + failed: "FAILED", +}; + +const OUTCOME_BADGE_CLASS: Record = { + passed: "bg-success/15 text-success border border-success/20", + flagged: "bg-warning/15 text-warning border border-warning/20", + failed: "bg-destructive/15 text-destructive border border-destructive/20", }; const getRiskColor = (score: number): string => { @@ -202,6 +221,19 @@ const FailCircleIcon = ({ className }: { className?: string }) => ( ); +const FlagCircleIcon = ({ className }: { className?: string }) => ( + + + + +); + +const OutcomeIcon = ({ outcome }: { outcome: EntryOutcome }) => { + if (outcome === "passed") return ; + if (outcome === "flagged") return ; + return ; +}; + const PlayCircleIcon = () => ( @@ -318,8 +350,7 @@ interface TimelineEntry { type: "request" | "guardrail" | "llm" | "response"; label: string; offsetMs: number; - status?: string; - isSuccess?: boolean; + outcome?: EntryOutcome; } const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { @@ -348,8 +379,7 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { type: "guardrail", label: `Pre-call guardrail: ${getDisplayName(e)}`, offsetMs, - status: isEntrySuccess(e) ? "PASSED" : "FAILED", - isSuccess: isEntrySuccess(e), + outcome: getEntryOutcome(e), }); } @@ -372,8 +402,7 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { type: "guardrail", label: `During-call guardrail: ${getDisplayName(e)}`, offsetMs, - status: isEntrySuccess(e) ? "PASSED" : "FAILED", - isSuccess: isEntrySuccess(e), + outcome: getEntryOutcome(e), }); } @@ -384,8 +413,7 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { type: "guardrail", label: `Post-call guardrail: ${getDisplayName(e)}`, offsetMs, - status: isEntrySuccess(e) ? "PASSED" : "FAILED", - isSuccess: isEntrySuccess(e), + outcome: getEntryOutcome(e), }); } @@ -410,10 +438,8 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { ) : item.type === "llm" ? ( - ) : item.isSuccess ? ( - ) : ( - + )} {idx < timeline.length - 1 &&
} @@ -425,13 +451,11 @@ const RequestLifecycle = ({ entries }: { entries: GuardrailInformation[] }) => { {item.label} - {item.status && ( + {item.outcome && ( - {item.status} + {OUTCOME_LABEL[item.outcome]} )} T+{item.offsetMs}ms @@ -455,7 +479,7 @@ const formatGuardrailCost = (cost: number): string => { const EvaluationCard = ({ entry }: { entry: GuardrailInformation }) => { const [expanded, setExpanded] = useState(false); - const success = isEntrySuccess(entry); + const outcome = getEntryOutcome(entry); const totalMasked = getTotalMasked(entry); const displayName = getDisplayName(entry); const durationStr = formatDurationMs(entry.duration); @@ -490,7 +514,9 @@ const EvaluationCard = ({ entry }: { entry: GuardrailInformation }) => { onClick={() => setExpanded(!expanded)} > {/* Status icon */} -
{success ? : }
+
+ +
{/* Name + badges */}
@@ -501,13 +527,9 @@ const EvaluationCard = ({ entry }: { entry: GuardrailInformation }) => { - {success ? "PASSED" : "FAILED"} + {OUTCOME_LABEL[outcome]} {matchCountStr && ( @@ -528,7 +550,7 @@ const EvaluationCard = ({ entry }: { entry: GuardrailInformation }) => { )} - {riskScore != null && success && ( + {riskScore != null && outcome === "passed" && ( getEntryOutcome(e) === "flagged").length; const allPassed = passedCount === guardrailEntries.length; + const headerOutcome: EntryOutcome = allPassed + ? "passed" + : passedCount + flaggedCount === guardrailEntries.length + ? "flagged" + : "failed"; const totalOverheadMs = useMemo(() => { return Math.round(guardrailEntries.reduce((sum, e) => sum + (e.duration ?? 0), 0) * 1000); @@ -709,11 +737,7 @@ const GuardrailViewer = ({ data, accessToken, logEntry }: GuardrailViewerProps) | {allPassed ? ( @@ -728,6 +752,13 @@ const GuardrailViewer = ({ data, accessToken, logEntry }: GuardrailViewerProps) ) : null} {passedCount} Passed + {flaggedCount > 0 && ( + + {flaggedCount} Flagged + + )}
diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx index a679dc49427..2e9bce5048f 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.test.tsx @@ -1,7 +1,7 @@ import { render, screen, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { describe, expect, it, vi } from "vitest"; -import { LogDetailContent } from "./LogDetailContent"; +import { GuardrailJumpLink, LogDetailContent } from "./LogDetailContent"; import type { LogEntry } from "../columns"; vi.mock("../GuardrailViewer/GuardrailViewer", () => ({ @@ -489,3 +489,17 @@ describe("LogDetailContent", () => { expect(within(descriptions).getByText("-")).toBeInTheDocument(); }); }); + +describe("GuardrailJumpLink", () => { + it.each([ + [["success", "success"], "text-success", "\u2713"], + [["success", "guardrail_flagged"], "text-warning", "\u26A0"], + [["guardrail_flagged", "guardrail_intervened"], "text-destructive", "\u2717"], + ])("styles %j as %s", (statuses, expectedClass, glyph) => { + render( ({ guardrail_status: s }))} />); + + const pill = screen.getByText(/2 guardrails evaluated/); + expect(pill).toHaveClass(expectedClass); + expect(pill).toHaveTextContent(glyph); + }); +}); diff --git a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx index 4c5c7b7b43f..052f1ec8802 100644 --- a/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx +++ b/ui/litellm-dashboard/src/components/view_logs/LogDetailsDrawer/LogDetailContent.tsx @@ -635,11 +635,24 @@ function RequestResponseSection({ ); } +const GUARDRAIL_JUMP_LINK_STYLE = { + passed: { className: "border border-success/20 bg-success/10 text-success", glyph: "\u2713" }, + flagged: { className: "border border-warning/20 bg-warning/10 text-warning", glyph: "\u26A0" }, + failed: { className: "border border-destructive/20 bg-destructive/10 text-destructive", glyph: "\u2717" }, +} as const; + +const isPassedStatus = (status: unknown) => status === "pass" || status === "passed" || status === "success"; +const isFlaggedStatus = (status: unknown) => status === "flagged" || status === "guardrail_flagged"; + +const guardrailJumpLinkOutcome = (statuses: unknown[]): keyof typeof GUARDRAIL_JUMP_LINK_STYLE => { + if (statuses.every(isPassedStatus)) return "passed"; + if (statuses.every((s) => isPassedStatus(s) || isFlaggedStatus(s))) return "flagged"; + return "failed"; +}; + export function GuardrailJumpLink({ guardrailEntries }: { guardrailEntries: any[] }) { - const allPassed = guardrailEntries.every((e) => { - const status = e?.guardrail_status || e?.status; - return status === "pass" || status === "passed" || status === "success"; - }); + const outcome = guardrailJumpLinkOutcome(guardrailEntries.map((e) => e?.guardrail_status || e?.status)); + const { className, glyph } = GUARDRAIL_JUMP_LINK_STYLE[outcome]; const handleClick = () => { const el = document.getElementById("guardrail-section"); @@ -650,11 +663,7 @@ export function GuardrailJumpLink({ guardrailEntries }: { guardrailEntries: any[
- {allPassed ? "\u2713" : "\u2717"} {guardrailEntries.length} guardrail{guardrailEntries.length !== 1 ? "s" : ""}{" "} - evaluated + {glyph} {guardrailEntries.length} guardrail + {guardrailEntries.length !== 1 ? "s" : ""} evaluated {"\u2193"}
From 80839bb33c318851af40b125c54c262bbe5dc90f Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 13:26:09 -0700 Subject: [PATCH 5/7] feat(proxy): serve Prometheus /metrics from a separate process via --prometheus_metrics_port (#39889) * feat(proxy): serve Prometheus /metrics from a separate process via --prometheus_metrics_port Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): ruff format prometheus_metrics_server Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): fail fast when the separate metrics server cannot start and force the multiproc dir whenever it is enabled - wait for the child's /health before starting uvicorn; raise a ClickException if it exits first (port in use) - create PROMETHEUS_MULTIPROC_DIR whenever --prometheus_metrics_port is set, so DB-configured prometheus callbacks work - honour lowercase prometheus_multiproc_dir; validate the port before spawning - cover main() entry point, readiness, bind failure and wildcard-host probing in tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): pin metrics-server readiness to the child pid so another service on the port cannot pass the health check Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): probe metrics-server readiness through the shared HTTPHandler instead of bare httpx.get Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(proxy): serve only /metrics on the prometheus metrics port Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): validate metrics server CLI args with pydantic instead of typing.cast Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): satisfy metrics server lint gate Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- basedpyright-code-budget.json | 4 +- litellm/proxy/prometheus_metrics_server.py | 167 +++++++++++ litellm/proxy/proxy_cli.py | 90 ++++-- .../proxy/test_prometheus_cleanup.py | 40 +++ .../proxy/test_prometheus_metrics_server.py | 259 ++++++++++++++++++ tests/test_litellm/proxy/test_proxy_cli.py | 118 ++++++++ type-discipline-budget.json | 4 +- 7 files changed, 649 insertions(+), 33 deletions(-) create mode 100644 litellm/proxy/prometheus_metrics_server.py create mode 100644 tests/test_litellm/proxy/test_prometheus_metrics_server.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index b876cf2d69a..0b0a61192e6 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -105,13 +105,13 @@ "limit": 109 }, "reportUnknownMemberType": { - "limit": 38283 + "limit": 38271 }, "reportUnknownParameterType": { "limit": 19584 }, "reportUnknownVariableType": { - "limit": 29829 + "limit": 29814 }, "reportUnnecessaryCast": { "limit": 110 diff --git a/litellm/proxy/prometheus_metrics_server.py b/litellm/proxy/prometheus_metrics_server.py new file mode 100644 index 00000000000..4a9651d62e1 --- /dev/null +++ b/litellm/proxy/prometheus_metrics_server.py @@ -0,0 +1,167 @@ +"""Serve Prometheus `/metrics` from its own process so a scrape never runs on an inference worker. + +Workers write their samples to `PROMETHEUS_MULTIPROC_DIR`; this process reads them back with a +``MultiProcessCollector`` and serves the aggregated output on a separate port. The proxy CLI starts +it with ``--prometheus_metrics_port``. It can also run as a sidecar sharing the same directory: +``python -m litellm.proxy.prometheus_metrics_server --host 0.0.0.0 --port 4001``. +""" + +from __future__ import annotations + +import argparse +import atexit +import os +import subprocess +import sys +import threading +import time +from collections.abc import Sequence +from contextlib import closing +from types import MappingProxyType +from typing import Final + +import httpx +from fastapi import FastAPI +from prometheus_client import CollectorRegistry, multiprocess +from pydantic import BaseModel, ConfigDict +from starlette.types import ASGIApp, Message, Receive, Scope, Send + +from litellm.integrations.prometheus_metrics_endpoint import make_metrics_asgi_app +from litellm.llms.custom_httpx.http_handler import HTTPHandler + +METRICS_PATH: Final = "/metrics" +PID_HEADER: Final = "x-litellm-metrics-pid" +_PARENT_POLL_INTERVAL_SECONDS: Final = 1.0 +_STARTUP_TIMEOUT_SECONDS: Final = 30.0 +_STARTUP_POLL_INTERVAL_SECONDS: Final = 0.1 +_STARTUP_PROBE_TIMEOUT_SECONDS: Final = 1.0 +_WILDCARD_TO_LOOPBACK: Final = MappingProxyType({"0.0.0.0": "127.0.0.1", "::": "::1"}) + + +class _CliArgs(BaseModel): + model_config = ConfigDict(frozen=True) + + host: str + port: int + multiproc_dir: str | None + + +class MetricsServerStartupError(RuntimeError): + """The metrics process died or never answered on its port before the proxy started serving.""" + + +def _add_pid_header(app: ASGIApp) -> ASGIApp: + async def app_with_pid(scope: Scope, receive: Receive, send: Send) -> None: + async def send_with_pid(message: Message) -> None: + if message["type"] == "http.response.start": + await send( + { + **message, + "headers": [ + *message["headers"], + (PID_HEADER.encode(), str(os.getpid()).encode()), + ], + } + ) + return + await send(message) + + await app(scope, receive, send_with_pid) + + return app_with_pid + + +def build_metrics_app(multiproc_dir: str) -> FastAPI: + registry: Final = CollectorRegistry() + multiprocess.MultiProcessCollector(registry, path=multiproc_dir) + app: Final = FastAPI(title="LiteLLM Prometheus metrics", docs_url=None, redoc_url=None, openapi_url=None) + app.mount(METRICS_PATH, _add_pid_header(make_metrics_asgi_app(registry))) + + return app + + +def _exit_when_parent_dies(parent_pid: int) -> None: + def watch() -> None: + while os.getppid() == parent_pid: + time.sleep(_PARENT_POLL_INTERVAL_SECONDS) + os._exit(0) + + threading.Thread(target=watch, name="litellm-metrics-parent-watchdog", daemon=True).start() + + +def run_metrics_server(host: str, port: int, multiproc_dir: str) -> None: + import uvicorn + + _exit_when_parent_dies(os.getppid()) + uvicorn.run(build_metrics_app(multiproc_dir), host=host, port=port, log_level="warning", access_log=False) + + +def metrics_url(host: str, port: int) -> str: + probe_host: Final = _WILDCARD_TO_LOOPBACK.get(host, host) + netloc: Final = f"[{probe_host}]" if ":" in probe_host else probe_host + return f"http://{netloc}:{port}{METRICS_PATH}" + + +def _answered_by(http: HTTPHandler, url: str, pid: int) -> bool: + """True only when the metrics response comes from our child, not from whatever else holds the port.""" + try: + response: Final = http.get(url) # pyright: ignore[reportUnknownMemberType] # HTTPHandler.get exposes untyped optional mappings + return response.status_code == 200 and response.headers.get(PID_HEADER) == str(pid) + except httpx.TransportError: + return False + + +def _wait_until_serving(process: subprocess.Popen[bytes], host: str, port: int) -> None: + url: Final = metrics_url(host, port) + deadline: Final = time.monotonic() + _STARTUP_TIMEOUT_SECONDS + with closing(HTTPHandler(timeout=_STARTUP_PROBE_TIMEOUT_SECONDS)) as http: + while time.monotonic() < deadline: + if (returncode := process.poll()) is not None: + raise MetricsServerStartupError( + f"Prometheus metrics server exited with code {returncode} before serving {host}:{port}; " + "is the port already in use?" + ) + if _answered_by(http, url, process.pid): + return + time.sleep(_STARTUP_POLL_INTERVAL_SECONDS) + process.terminate() + raise MetricsServerStartupError( + f"Prometheus metrics server did not answer {url} within {_STARTUP_TIMEOUT_SECONDS:.0f}s" + ) + + +def start_metrics_server_process(host: str, port: int, multiproc_dir: str) -> subprocess.Popen[bytes]: + """Spawn the metrics server next to the proxy and block until it answers on its port.""" + process: Final = subprocess.Popen( + ( + sys.executable, + "-m", + "litellm.proxy.prometheus_metrics_server", + "--host", + host, + "--port", + str(port), + "--multiproc_dir", + multiproc_dir, + ) + ) + atexit.register(process.terminate) + _wait_until_serving(process, host, port) + return process + + +def main(argv: Sequence[str] | None = None) -> None: + parser: Final = argparse.ArgumentParser( + description="Serve LiteLLM Prometheus metrics from PROMETHEUS_MULTIPROC_DIR" + ) + parser.add_argument("--host", default="0.0.0.0") + parser.add_argument("--port", type=int, required=True) + parser.add_argument("--multiproc_dir", default=os.environ.get("PROMETHEUS_MULTIPROC_DIR")) + args: Final = _CliArgs.model_validate(vars(parser.parse_args(argv))) + if not args.multiproc_dir: + parser.error("--multiproc_dir or PROMETHEUS_MULTIPROC_DIR is required") + run_metrics_server(host=args.host, port=args.port, multiproc_dir=args.multiproc_dir) + + +if __name__ == "__main__": + main() diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index e780beb4410..e245367b1b4 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -7,7 +7,7 @@ import re import subprocess import sys import urllib.parse as urlparse -from collections.abc import Iterable +from collections.abc import Iterable, Mapping, Sequence from pathlib import Path from typing import TYPE_CHECKING, Any, Final @@ -610,48 +610,49 @@ class ProxyInitializationHelpers: return None # Let uvicorn choose the default loop on Windows return "uvloop" + @staticmethod + def _prometheus_callback_configured(litellm_settings: Mapping[str, object] | None) -> bool: + if litellm_settings is None: + return False + configured: Final = tuple( + litellm_settings.get(key) for key in ("callbacks", "success_callback", "failure_callback") + ) + return any( + setting == "prometheus" + if isinstance(setting, str) + else isinstance(setting, Sequence) and "prometheus" in setting + for setting in configured + ) + @staticmethod def _maybe_setup_prometheus_multiproc_dir( num_workers: int, litellm_settings: dict | None, - ) -> None: + prometheus_metrics_port: int | None = None, + ) -> str | None: """ - Auto-create PROMETHEUS_MULTIPROC_DIR when running with multiple workers - and prometheus is configured as a callback. + Auto-create PROMETHEUS_MULTIPROC_DIR when another process needs to read the samples: extra workers + with prometheus configured as a callback in config.yaml, or the separate metrics server (always, since + callbacks may also be enabled from the DB after startup). """ import tempfile - if num_workers <= 1 or litellm_settings is None: - return - - # Check if prometheus is in any callback list - # Each setting can be a list or a single string; normalize to list - callbacks = litellm_settings.get("callbacks") or [] - success_callbacks = litellm_settings.get("success_callback") or [] - failure_callbacks = litellm_settings.get("failure_callback") or [] - if isinstance(callbacks, str): - callbacks = [callbacks] - if isinstance(success_callbacks, str): - success_callbacks = [success_callbacks] - if isinstance(failure_callbacks, str): - failure_callbacks = [failure_callbacks] - all_callbacks: Final = callbacks + success_callbacks + failure_callbacks - if "prometheus" not in all_callbacks: - return + if prometheus_metrics_port is None and ( + num_workers <= 1 or not ProxyInitializationHelpers._prometheus_callback_configured(litellm_settings) + ): + return None from litellm.proxy.prometheus_cleanup import wipe_directory - multiproc_dir = os.environ.get("PROMETHEUS_MULTIPROC_DIR") or os.environ.get("prometheus_multiproc_dir") - - auto_created: Final = not multiproc_dir - if not multiproc_dir: - multiproc_dir = os.path.join(tempfile.gettempdir(), "litellm_prometheus_multiproc") - os.environ["PROMETHEUS_MULTIPROC_DIR"] = multiproc_dir + configured_dir: Final = os.environ.get("PROMETHEUS_MULTIPROC_DIR") or os.environ.get("prometheus_multiproc_dir") + multiproc_dir: Final = configured_dir or os.path.join(tempfile.gettempdir(), "litellm_prometheus_multiproc") + os.environ["PROMETHEUS_MULTIPROC_DIR"] = multiproc_dir os.makedirs(multiproc_dir, exist_ok=True) wipe_directory(multiproc_dir) - action: Final = "Auto-created" if auto_created else "Using existing" + action: Final = "Using existing" if configured_dir else "Auto-created" print(f"LiteLLM: {action} PROMETHEUS_MULTIPROC_DIR={multiproc_dir}") + return multiproc_dir @click.command() @@ -930,6 +931,19 @@ class ProxyInitializationHelpers: default=False, help="Enable uvicorn hot reload (dev only). Also reloads when the --config YAML file changes. Incompatible with --num_workers>1, --run_gunicorn, and --run_hypercorn.", ) +@click.option( + "--prometheus_metrics_port", + default=None, + type=click.IntRange(min=1, max=65535), + help=( + "Serve Prometheus /metrics from a separate process on this port (bound to --host) so scraping and " + "multi-worker aggregation never run on an inference worker's event loop. Samples appear once the " + "`prometheus` callback is enabled (config.yaml or DB). /metrics stays mounted on the main port as well; " + "the separate port has no virtual-key auth, so keep it off public ingress. Startup fails if the metrics " + "server cannot bind." + ), + envvar="PROMETHEUS_METRICS_PORT", +) def run_server( cli_args, host, @@ -980,6 +994,7 @@ def run_server( enforce_prisma_migration_check: bool, use_v2_migration_resolver: bool, reload: bool, + prometheus_metrics_port: int | None, ): if cli_args: if cli_args == ("xai-oauth", "login"): @@ -1364,6 +1379,8 @@ def run_server( ) if port == 4000 and ProxyInitializationHelpers._is_port_in_use(port): port = random.randint(1024, 49152) + if prometheus_metrics_port == port: + raise click.UsageError("--prometheus_metrics_port must differ from --port") import litellm @@ -1374,9 +1391,10 @@ def run_server( from litellm.proxy.proxy_server import app # Auto-create PROMETHEUS_MULTIPROC_DIR for multi-worker setups - ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir( + prometheus_multiproc_dir: Final = ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir( num_workers=num_workers, litellm_settings=litellm_settings if config else None, + prometheus_metrics_port=prometheus_metrics_port, ) # Skip server startup if requested (after all setup is done) @@ -1384,6 +1402,20 @@ def run_server( print("LiteLLM: Setup complete. Skipping server startup as requested.") return + if prometheus_metrics_port is not None and prometheus_multiproc_dir is not None: + from litellm.proxy.prometheus_metrics_server import MetricsServerStartupError, start_metrics_server_process + + try: + metrics_process: Final = start_metrics_server_process( + host=host, port=prometheus_metrics_port, multiproc_dir=prometheus_multiproc_dir + ) + except MetricsServerStartupError as error: + raise click.ClickException(str(error)) from error + print( + f"\033[1;32mLiteLLM: Serving Prometheus metrics on {host}:{prometheus_metrics_port}/metrics " + f"(pid {metrics_process.pid})\033[0m" + ) + running_uvicorn: Final = run_gunicorn is False and run_hypercorn is False uvicorn_args: Final = ProxyInitializationHelpers._get_default_unvicorn_init_args( host=host, diff --git a/tests/test_litellm/proxy/test_prometheus_cleanup.py b/tests/test_litellm/proxy/test_prometheus_cleanup.py index ca5476d6af9..93b9b694c2c 100644 --- a/tests/test_litellm/proxy/test_prometheus_cleanup.py +++ b/tests/test_litellm/proxy/test_prometheus_cleanup.py @@ -131,3 +131,43 @@ class TestMaybeSetupPrometheusMultiprocDir: # Cleanup os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + + @pytest.mark.parametrize( + "litellm_settings", + [ + {"callbacks": ["prometheus"]}, + {"callbacks": ["langfuse"]}, + None, + ], + ) + def test_separate_metrics_port_forces_dir_for_single_worker(self, litellm_settings): + """The separate metrics process reads the samples, so one worker still needs the shared dir, even when + prometheus is not in config.yaml (callbacks can be turned on from the DB after startup).""" + with patch.dict(os.environ, {}, clear=False): + os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + os.environ.pop("prometheus_multiproc_dir", None) + + result_dir = ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir( + num_workers=1, + litellm_settings=litellm_settings, + prometheus_metrics_port=4001, + ) + + assert result_dir is not None + assert os.environ.get("PROMETHEUS_MULTIPROC_DIR") == result_dir + assert os.path.isdir(result_dir) + + os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + + def test_lowercase_env_var_is_reused_and_exported_uppercase(self, tmp_path): + """prometheus_client honours both spellings; the metrics server only reads the uppercase one.""" + with patch.dict(os.environ, {"prometheus_multiproc_dir": str(tmp_path)}, clear=False): + os.environ.pop("PROMETHEUS_MULTIPROC_DIR", None) + + result_dir = ProxyInitializationHelpers._maybe_setup_prometheus_multiproc_dir( + num_workers=4, + litellm_settings={"callbacks": "prometheus"}, + ) + + assert result_dir == str(tmp_path) + assert os.environ["PROMETHEUS_MULTIPROC_DIR"] == str(tmp_path) diff --git a/tests/test_litellm/proxy/test_prometheus_metrics_server.py b/tests/test_litellm/proxy/test_prometheus_metrics_server.py new file mode 100644 index 00000000000..fc1fa381fa4 --- /dev/null +++ b/tests/test_litellm/proxy/test_prometheus_metrics_server.py @@ -0,0 +1,259 @@ +"""The separate metrics server must aggregate PROMETHEUS_MULTIPROC_DIR, expose only /metrics, and follow its +parent's lifetime. + +Everything here runs on loopback against a child of this test process; no LLM keys or external network. +""" + +from __future__ import annotations + +import os +import socket +import subprocess +import sys +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Final +from unittest.mock import patch + +import httpx +import pytest +from fastapi.testclient import TestClient +from prometheus_client import values + +from litellm.proxy.prometheus_metrics_server import ( + PID_HEADER, + MetricsServerStartupError, + build_metrics_app, + main, + metrics_url, + start_metrics_server_process, +) + +_STARTUP_TIMEOUT_SECONDS: Final = 60.0 +_SHUTDOWN_TIMEOUT_SECONDS: Final = 15.0 + + +def _write_worker_sample(pid: int, value: float) -> None: + """Write one counter sample into PROMETHEUS_MULTIPROC_DIR the way a proxy worker would.""" + counter: Final = values.MultiProcessValue(process_identifier=lambda: pid)( + "counter", + "litellm_requests_metric_total", + "litellm_requests_metric_total", + ("model",), + ("gpt-5",), + "Total number of LLM calls", + ) + counter.inc(value) + + +def _free_port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _wait_for_metrics(port: int, pid: int) -> httpx.Response: + deadline: Final = time.monotonic() + _STARTUP_TIMEOUT_SECONDS + while time.monotonic() < deadline: + try: + response: Final = httpx.get(f"http://127.0.0.1:{port}/metrics", follow_redirects=True, timeout=1.0) + if response.status_code == 200 and response.headers.get(PID_HEADER) == str(pid): + return response + except httpx.TransportError: + pass + time.sleep(0.2) + raise AssertionError(f"metrics server on port {port} never served metrics") + + +def _wait_until_down(port: int) -> None: + deadline: Final = time.monotonic() + _SHUTDOWN_TIMEOUT_SECONDS + while time.monotonic() < deadline: + try: + httpx.get(f"http://127.0.0.1:{port}/metrics", timeout=1.0) + except httpx.TransportError: + return + time.sleep(0.2) + raise AssertionError(f"metrics server on port {port} kept running after its parent died") + + +def test_metrics_app_aggregates_multiproc_dir_and_reports_pid(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) + _write_worker_sample(pid=1001, value=2) + _write_worker_sample(pid=1002, value=3) + other_dir: Final = tmp_path / "other" + other_dir.mkdir() + + client: Final = TestClient(build_metrics_app(str(tmp_path))) + metrics: Final = client.get("/metrics") + assert metrics.status_code == 200 + assert metrics.headers[PID_HEADER] == str(os.getpid()) + assert 'litellm_requests_metric_total{model="gpt-5"} 5.0' in metrics.text + + assert client.get("/health").status_code == 404 + + empty: Final = TestClient(build_metrics_app(str(other_dir))).get("/metrics") + assert empty.status_code == 200 + assert "litellm_requests_metric_total" not in empty.text + + +@pytest.mark.parametrize( + ("host", "expected"), + ( + ("0.0.0.0", "http://127.0.0.1:4001/metrics"), + ("::", "http://[::1]:4001/metrics"), + ("10.1.2.3", "http://10.1.2.3:4001/metrics"), + ("metrics.internal", "http://metrics.internal:4001/metrics"), + ), +) +def test_metrics_url_probes_loopback_for_wildcard_binds(host: str, expected: str): + assert metrics_url(host, 4001) == expected + + +def test_main_serves_the_app_for_the_given_dir_with_uvicorn(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) + _write_worker_sample(pid=2001, value=6) + with patch("uvicorn.run") as run: + main(["--host", "10.1.2.3", "--port", "4001", "--multiproc_dir", str(tmp_path)]) + + run.assert_called_once() + assert run.call_args.kwargs["host"] == "10.1.2.3" + assert run.call_args.kwargs["port"] == 4001 + client: Final = TestClient(run.call_args.args[0]) + metrics: Final = client.get("/metrics") + assert metrics.status_code == 200 + assert metrics.headers[PID_HEADER] == str(os.getpid()) + assert 'litellm_requests_metric_total{model="gpt-5"} 6.0' in client.get("/metrics").text + + +def test_main_falls_back_to_env_multiproc_dir(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) + with patch("uvicorn.run") as run: + main(["--port", "4001"]) + + (app,), served_on = run.call_args + assert served_on["host"] == "0.0.0.0" + metrics: Final = TestClient(app).get("/metrics") + assert metrics.status_code == 200 + assert metrics.headers[PID_HEADER] == str(os.getpid()) + + +def test_main_rejects_missing_multiproc_dir(monkeypatch: pytest.MonkeyPatch): + monkeypatch.delenv("PROMETHEUS_MULTIPROC_DIR", raising=False) + with patch("uvicorn.run") as run, pytest.raises(SystemExit) as exit_info: + main(["--port", "4001"]) + + assert exit_info.value.code == 2 + run.assert_not_called() + + +def test_start_metrics_server_process_returns_only_once_child_serves(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) + _write_worker_sample(pid=3001, value=4) + port: Final = _free_port() + with patch("atexit.register") as register: + process: Final = start_metrics_server_process(host="127.0.0.1", port=port, multiproc_dir=str(tmp_path)) + try: + register.assert_called_once_with(process.terminate) + assert process.poll() is None + startup_metrics: Final = httpx.get(f"http://127.0.0.1:{port}/metrics", follow_redirects=True, timeout=5.0) + assert startup_metrics.status_code == 200 + assert startup_metrics.headers[PID_HEADER] == str(process.pid) + metrics: Final = httpx.get(f"http://127.0.0.1:{port}/metrics", follow_redirects=True, timeout=10.0) + assert 'litellm_requests_metric_total{model="gpt-5"} 4.0' in metrics.text + finally: + process.kill() + process.wait(timeout=10) + + +def test_start_metrics_server_process_fails_when_port_is_taken(tmp_path: Path): + with socket.socket() as occupied: + occupied.bind(("127.0.0.1", 0)) + occupied.listen() + port: Final = occupied.getsockname()[1] + with ( + patch("atexit.register"), + pytest.raises( + MetricsServerStartupError, match=rf"exited with code [1-9]\d* before serving 127.0.0.1:{port}" + ), + ): + start_metrics_server_process(host="127.0.0.1", port=port, multiproc_dir=str(tmp_path)) + + +class _ImpostorMetrics(BaseHTTPRequestHandler): + """An unrelated service already on the port that answers /metrics with 200 and plausible metrics.""" + + def do_GET(self) -> None: + body: Final = b"# HELP impostor_metric A plausible metric\n# TYPE impostor_metric counter\nimpostor_metric 1\n" + self.send_response(200) + self.send_header("Content-Type", "text/plain") + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + return + + +def test_start_metrics_server_process_rejects_metrics_from_another_service_on_the_port(tmp_path: Path): + with ThreadingHTTPServer(("127.0.0.1", 0), _ImpostorMetrics) as impostor: + threading.Thread(target=impostor.serve_forever, daemon=True).start() + port: Final = impostor.server_address[1] + impostor_response: Final = httpx.get(f"http://127.0.0.1:{port}/metrics") + assert impostor_response.status_code == 200 + assert "# HELP impostor_metric" in impostor_response.text + assert PID_HEADER not in impostor_response.headers + with ( + patch("atexit.register"), + pytest.raises(MetricsServerStartupError, match=rf"exited with code [1-9]\d* before serving 127.0.0.1:{port}"), + ): + start_metrics_server_process(host="127.0.0.1", port=port, multiproc_dir=str(tmp_path)) + impostor.shutdown() + + +def test_metrics_server_process_serves_and_exits_with_parent(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("PROMETHEUS_MULTIPROC_DIR", str(tmp_path)) + _write_worker_sample(pid=2001, value=7) + port: Final = _free_port() + server_argv: Final = ( + sys.executable, + "-m", + "litellm.proxy.prometheus_metrics_server", + "--host", + "127.0.0.1", + "--port", + str(port), + "--multiproc_dir", + str(tmp_path), + ) + parent: Final = subprocess.Popen( + ( + sys.executable, + "-c", + "import subprocess, sys, time; p = subprocess.Popen(sys.argv[1:]); print(p.pid, flush=True); time.sleep(600)", + *server_argv, + ), + stdout=subprocess.PIPE, + text=True, + ) + assert parent.stdout is not None + server_pid: Final = int(parent.stdout.readline()) + try: + metrics: Final = _wait_for_metrics(port, server_pid) + assert metrics.status_code == 200 + assert metrics.headers[PID_HEADER] == str(server_pid) + + scrape: Final = httpx.get(f"http://127.0.0.1:{port}/metrics", follow_redirects=True, timeout=10.0) + assert scrape.status_code == 200 + assert scrape.headers[PID_HEADER] == str(server_pid) + assert 'litellm_requests_metric_total{model="gpt-5"} 7.0' in scrape.text + + parent.kill() + parent.wait(timeout=10) + _wait_until_down(port) + finally: + parent.kill() + try: + os.kill(server_pid, 9) + except ProcessLookupError: + pass diff --git a/tests/test_litellm/proxy/test_proxy_cli.py b/tests/test_litellm/proxy/test_proxy_cli.py index 9256706d340..0c20d5e0ff0 100644 --- a/tests/test_litellm/proxy/test_proxy_cli.py +++ b/tests/test_litellm/proxy/test_proxy_cli.py @@ -662,6 +662,124 @@ class TestProxyInitializationHelpers: assert "Invalid value for '--limit_concurrency'" in result.output mock_uvicorn_run.assert_not_called() + @patch("uvicorn.run") + @patch("httpx.HTTPTransport.handle_request") + @patch("atexit.register") + @patch("subprocess.Popen") + @patch("litellm.proxy.db.prisma_client.PrismaManager.setup_database") # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + @patch( # test-quality-ok: run_server always wires the DB; same isolation as the sibling CLI tests above + "litellm.proxy.db.prisma_client.should_update_prisma_schema", return_value=False + ) + def test_prometheus_metrics_port_starts_separate_metrics_process( + self, + mock_should_update, + mock_setup_db, + mock_popen, + mock_atexit_register, + mock_handle_request, + mock_uvicorn_run, + tmp_path, + ): + """--prometheus_metrics_port must spawn `python -m litellm.proxy.prometheus_metrics_server` on --host + with the shared multiproc dir, wait for its /metrics response, and only then start uvicorn. It must stay off by + default, refuse to share --port, and abort the proxy when the child dies before serving.""" + import httpx + from click.testing import CliRunner + + from litellm.proxy.proxy_cli import run_server + + runner = CliRunner() + mock_popen.return_value = MagicMock(pid=4242, **{"poll.return_value": None}) + probed_urls: list[str] = [] + + def child_metrics(request: httpx.Request) -> httpx.Response: + probed_urls.append(str(request.url)) + return httpx.Response(200, headers={"x-litellm-metrics-pid": "4242"}, content=b"") + + mock_handle_request.side_effect = child_metrics + mock_proxy_module = MagicMock( + app=MagicMock(), + ProxyConfig=MagicMock(), + KeyManagementSettings=MagicMock(), + save_worker_config=MagicMock(), + ) + clean_env = { + k: v + for k, v in os.environ.items() + if k not in ("DATABASE_URL", "DIRECT_URL", "PROMETHEUS_METRICS_PORT") + } + clean_env["PROMETHEUS_MULTIPROC_DIR"] = str(tmp_path) + with ( + patch.dict(os.environ, clean_env, clear=True), + patch.dict( + "sys.modules", + { + "proxy_server": mock_proxy_module, + "litellm.proxy.proxy_server": mock_proxy_module, + }, + ), + patch( # test-quality-ok: same isolation as the sibling CLI tests above + "litellm.proxy.proxy_cli.ProxyInitializationHelpers._get_default_unvicorn_init_args" + ) as mock_get_args, + ): + mock_get_args.side_effect = lambda *a, **k: { + "app": "litellm.proxy.proxy_server:app", + "host": "localhost", + "port": 8000, + } + + result = runner.invoke( + run_server, + ["--local", "--host", "127.0.0.1", "--port", "4000", "--prometheus_metrics_port", "4001"], + ) + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + mock_popen.assert_called_once() + spawned = list(mock_popen.call_args.args[0]) + assert spawned[1:3] == ["-m", "litellm.proxy.prometheus_metrics_server"] + assert spawned[3:] == ["--host", "127.0.0.1", "--port", "4001", "--multiproc_dir", str(tmp_path)] + assert probed_urls == ["http://127.0.0.1:4001/metrics"] + assert "Serving Prometheus metrics on 127.0.0.1:4001/metrics (pid 4242)" in result.output + mock_uvicorn_run.assert_called_once() + + mock_popen.reset_mock() + mock_uvicorn_run.reset_mock() + mock_popen.return_value = MagicMock(pid=4243, **{"poll.return_value": 1}) + result = runner.invoke( + run_server, + ["--local", "--port", "4000", "--prometheus_metrics_port", "4001"], + ) + assert result.exit_code == 1, f"exit_code={result.exit_code}, output={result.output}" + assert "Prometheus metrics server exited with code 1 before serving 0.0.0.0:4001" in result.output + mock_uvicorn_run.assert_not_called() + + mock_popen.reset_mock() + mock_uvicorn_run.reset_mock() + result = runner.invoke(run_server, ["--local"]) + assert ( + result.exit_code == 0 + ), f"exit_code={result.exit_code}, output={result.output}" + mock_popen.assert_not_called() + mock_uvicorn_run.assert_called_once() + + mock_uvicorn_run.reset_mock() + result = runner.invoke( + run_server, + ["--local", "--port", "4000", "--prometheus_metrics_port", "4000"], + ) + assert result.exit_code == 2 + assert "--prometheus_metrics_port must differ from --port" in result.output + mock_popen.assert_not_called() + mock_uvicorn_run.assert_not_called() + + result = runner.invoke( + run_server, ["--local", "--prometheus_metrics_port", "0"] + ) + assert result.exit_code == 2 + assert "Invalid value for '--prometheus_metrics_port'" in result.output + mock_popen.assert_not_called() + @pytest.mark.parametrize( "timeout_config,expected_timeout", [ diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 225134b4e2b..e35e470c979 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -3,7 +3,7 @@ "limit": 22180 }, "LIT002": { - "limit": 26745 + "limit": 26729 }, "LIT003": { "limit": 261 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16462 + "limit": 16430 }, "LIT011": { "limit": 5506 From 9832d6e4a6e3cc832e4425f759243f97926e6d46 Mon Sep 17 00:00:00 2001 From: Yassin Kortam Date: Sat, 5 Sep 2026 13:51:25 -0700 Subject: [PATCH 6/7] fix(mcp): scan and mask MCP tool call arguments in unified guardrails (#35142) * fix(mcp): scan and mask MCP tool call arguments in unified guardrails A guardrail configured with mode pre_mcp_call was handed only a synthetic tool definition (name plus an empty parameters schema), so it never saw the argument values it was configured to inspect, and any rewrite it returned was discarded. Detection could not fire and masking could not take effect, while the applied-guardrails metadata still reported the guardrail as having run. Pass every string leaf of the tool call arguments as texts, and fold the guardrail's rewritten leaves back into modified_arguments, which is the channel the MCP call path reads to decide what to send upstream. The leaf walk reuses the json_string_leaves / with_json_string_leaves helpers the tool result path already uses, so both directions share one bounded traversal. Two guardrails running concurrently under run_in_parallel scan the same payload snapshot, so each returns a full replacement derived from the original leaf. Rewrites of the same leaf to different values are rejected rather than silently losing one redaction; a leaf that already holds this guardrail's own replacement is convergent and still masks, which is what the bundled content filter does when it rewrites the arguments itself as well as through texts. * fix(mcp): annotate guardrail argument rewrites Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(tests): isolate MCP guardrail callback state Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore: ratchet LIT010 budget after merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(tests): remove duplicate Bedrock hook parameter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(mcp): fail closed when guardrail rewrites cannot be mapped to MCP arguments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(tests): patch the guardrail translation mappings cache where staging now keeps it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../guardrail_translation/handler.py | 137 ++++-- .../test_mcp_guardrail_handler.py | 429 +++++++++++++++++- type-discipline-budget.json | 2 +- 3 files changed, 530 insertions(+), 38 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py index 4918229c2b8..c0235077ecd 100644 --- a/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py +++ b/litellm/proxy/_experimental/mcp_server/guardrail_translation/handler.py @@ -1,16 +1,19 @@ """ MCP Guardrail Handler for Unified Guardrails. -Converts an MCP call_tool (name + arguments) into a single OpenAI-compatible -tool_call and passes it to apply_guardrail. Works with the synthetic payload -from ProxyLogging._convert_mcp_to_llm_format. +Converts an MCP call_tool (name + arguments) into the OpenAI-compatible shape +apply_guardrail expects: the tool as a single-entry ``tools`` definition, and +every string leaf of the call arguments as ``texts`` so text guardrails can +detect and mask sensitive values in the payload. Works with the synthetic +request from ProxyLogging._convert_mcp_to_llm_format. Note: For MCP tool definitions (schema) -> OpenAI tools=[], see litellm.experimental_mcp_client.tools.transform_mcp_tool_to_openai_tool when you have a full MCP Tool from list_tools. Here we only have the call -payload (name + arguments) so we just build the tool_call. +payload (name + arguments) so we just build the tool definition. """ +from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any, Final from fastapi import HTTPException @@ -20,6 +23,8 @@ from litellm._logging import verbose_proxy_logger from litellm.experimental_mcp_client.tools import transform_mcp_tool_to_openai_tool from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.proxy._experimental.mcp_server.utils import ( + MAX_STRUCTURED_CONTENT_SCAN_DEPTH, + JSONLeafPath, json_string_leaves, json_unrewritable_labels, mcp_content_item_text, @@ -42,6 +47,72 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +def _blocked(reason: str) -> HTTPException: + return HTTPException(status_code=400, detail={"error": f"Content blocked: {reason}"}) + + +def _too_deeply_nested() -> HTTPException: + return _blocked( + f"MCP tool call arguments exceed the maximum nesting depth of {MAX_STRUCTURED_CONTENT_SCAN_DEPTH} " + "and cannot be scanned by the configured guardrail" + ) + + +def _argument_replacements( + argument_leaves: tuple[tuple[JSONLeafPath, str], ...], + masked_texts: Sequence[str] | None, +) -> Mapping[JSONLeafPath, str]: + """Positionally pair the guardrail's returned texts with the leaves they came from. + + Only leaves the guardrail actually rewrote are returned, so a guardrail that + detects nothing leaves the outbound tool call byte-identical. A guardrail that + returns the wrong number of texts fails closed, because a positional write-back + would scramble the arguments rather than mask them. + """ + if masked_texts is not None and len(masked_texts) != len(argument_leaves): + raise _blocked( + f"guardrail returned {len(masked_texts)} texts for {len(argument_leaves)} MCP tool call argument strings, " + "so the redaction cannot be mapped back to the arguments" + ) + return {path: masked for (path, original), masked in zip(argument_leaves, masked_texts or ()) if masked != original} + + +def _conflicting_rewrite_paths( + scanned_leaves: tuple[tuple[JSONLeafPath, str], ...], + current_leaves: tuple[tuple[JSONLeafPath, str], ...], + replacements: Mapping[JSONLeafPath, str], +) -> tuple[JSONLeafPath, ...]: + """Paths another guardrail already rewrote differently from what this one wants. + + Guardrails opted into ``run_in_parallel`` all scan the same payload snapshot, so + each one returns a full replacement string derived from the *original* leaf. Two + of them rewriting one leaf to different values cannot be merged: writing either + result discards the other guardrail's redaction. A leaf still holding the text + this guardrail was handed, or already holding this guardrail's own replacement, + is safe to write; the latter is how a guardrail that masks the arguments itself + as well as through ``texts`` gets there first. Anything else fails closed, + including a payload reshaped so the leaves no longer line up, because the + write-back is positional and would land a redaction on the wrong value. + """ + if tuple(path for path, _ in scanned_leaves) != tuple(path for path, _ in current_leaves): + return tuple(replacements) + return tuple( + path + for (path, scanned), (_, current) in zip(scanned_leaves, current_leaves) + if path in replacements and current not in (scanned, replacements[path]) + ) + + +def _conflicting_rewrite(paths: tuple[JSONLeafPath, ...]) -> HTTPException: + return _blocked( + "two guardrails running concurrently rewrote the same MCP tool call " + f"argument{'s' if len(paths) > 1 else ''} " + f"({', '.join('.'.join(str(part) for part in path) for path in paths)}); " + "their redactions cannot be merged. Remove run_in_parallel from one of them so they " + "run in sequence." + ) + + class MCPGuardrailTranslationHandler(BaseTranslation): """Guardrail translation handler for MCP tool calls (passes a single tool_call to guardrail).""" @@ -52,10 +123,8 @@ class MCPGuardrailTranslationHandler(BaseTranslation): litellm_logging_obj: "LiteLLMLoggingObj | None" = None, ) -> dict[str, Any]: mcp_tool_name: Final = data.get("mcp_tool_name") or data.get("name") - mcp_arguments = data.get("mcp_arguments") or data.get("arguments") + mcp_arguments: Final[object] = data.get("mcp_arguments") or data.get("arguments") mcp_tool_description: Final = data.get("mcp_tool_description") or data.get("description") - if mcp_arguments is None or not isinstance(mcp_arguments, dict): - mcp_arguments = {} if not mcp_tool_name: verbose_proxy_logger.debug("MCP Guardrail: mcp_tool_name missing") @@ -84,16 +153,37 @@ class MCPGuardrailTranslationHandler(BaseTranslation): strict=fn.get("strict", False) or False, # Default to False if None ), } + argument_leaves: Final = json_string_leaves(mcp_arguments) + if argument_leaves is None: + raise _too_deeply_nested() inputs: Final[GenericGuardrailAPIInputs] = GenericGuardrailAPIInputs( tools=[tool_def], + texts=[text for _, text in argument_leaves], ) - await guardrail_to_apply.apply_guardrail( + guarded: Final = await guardrail_to_apply.apply_guardrail( inputs=inputs, request_data=data, input_type="request", logging_obj=litellm_logging_obj, ) + replacements: Final = _argument_replacements( + argument_leaves=argument_leaves, + masked_texts=guarded.get("texts") if guarded else None, + ) + if not replacements: + return data + + current_arguments: Final[object] = data.get("mcp_arguments") or data.get("arguments") + current_leaves: Final = json_string_leaves(current_arguments) + if current_leaves is None: + raise _too_deeply_nested() + conflicting: Final = _conflicting_rewrite_paths(argument_leaves, current_leaves, replacements) + if conflicting: + raise _conflicting_rewrite(conflicting) + masked_arguments: Final = with_json_string_leaves(current_arguments, replacements) + data["mcp_arguments"] = masked_arguments # rebind-ok: preserve the mask for the outbound MCP call + data["modified_arguments"] = masked_arguments # rebind-ok: expose the applied mask to the caller return data async def process_output_response( @@ -131,14 +221,8 @@ class MCPGuardrailTranslationHandler(BaseTranslation): structured_leaves: Final = json_string_leaves(structured) if structured is not None else () structured_labels: Final = json_unrewritable_labels(structured) if structured is not None else () if structured_leaves is None or structured_labels is None: - raise HTTPException( - status_code=400, - detail={ - "error": ( - "Content blocked: MCP tool result structuredContent is nested too deeply to be scanned " - "by the configured guardrail" - ) - }, + raise _blocked( + "MCP tool result structuredContent is nested too deeply to be scanned by the configured guardrail" ) if not text_blocks and not structured_leaves and not structured_labels: @@ -158,12 +242,10 @@ class MCPGuardrailTranslationHandler(BaseTranslation): if masked_texts is None: return response if len(masked_texts) != len(originals): - verbose_proxy_logger.warning( - "MCP Guardrail: guardrail returned %d texts for %d tool result texts; leaving the result unmasked", - len(masked_texts), - len(originals), + raise _blocked( + f"guardrail returned {len(masked_texts)} texts for {len(originals)} MCP tool result texts, " + "so the redaction cannot be mapped back to the result" ) - return response split: Final = len(text_blocks) if content is not None: @@ -173,15 +255,10 @@ class MCPGuardrailTranslationHandler(BaseTranslation): label_start: Final = split + len(structured_leaves) if any(masked != original for original, masked in zip(structured_labels, masked_texts[label_start:])): - raise HTTPException( - status_code=400, - detail={ - "error": ( - "Content blocked: MCP tool result matched a masking rule on a non-rewritable field " - "(a structuredContent key or numeric value), which cannot be redacted without changing " - "the payload contract" - ) - }, + raise _blocked( + "MCP tool result matched a masking rule on a non-rewritable field " + "(a structuredContent key or numeric value), which cannot be redacted without changing " + "the payload contract" ) structured_replacements: Final = { diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py b/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py index 2e286a237c4..28959054195 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/guardrail_translation/test_mcp_guardrail_handler.py @@ -1,13 +1,22 @@ """Tests for the MCP guardrail translation handler.""" +import asyncio + import pytest +from fastapi import HTTPException from mcp.types import CallToolResult, ImageContent, TextContent +import litellm +import litellm.llms as litellm_llms +from litellm.caching.caching import DualCache from litellm.exceptions import BlockedPiiEntityError from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.proxy._experimental.mcp_server.guardrail_translation.handler import ( MCPGuardrailTranslationHandler, ) +from litellm.proxy._experimental.mcp_server.utils import MAX_STRUCTURED_CONTENT_SCAN_DEPTH +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.utils import ProxyLogging from litellm.types.utils import GenericGuardrailAPIInputs @@ -24,12 +33,11 @@ class MockGuardrail(CustomGuardrail): self.call_count += 1 self.last_inputs = inputs self.last_request_data = request_data - return None # Guardrail doesn't modify for MCP tools @pytest.mark.asyncio async def test_process_input_messages_updates_content(): - """Handler should pass tool definition to guardrail when mcp_tool_name is present.""" + """Handler should pass the tool definition and the argument strings to the guardrail.""" handler = MCPGuardrailTranslationHandler() guardrail = MockGuardrail() @@ -45,7 +53,7 @@ async def test_process_input_messages_updates_content(): assert result == data # Guardrail was called assert guardrail.call_count == 1 - # Guardrail received tools (not texts) with tool definition + # Guardrail received tools with the tool definition assert guardrail.last_inputs is not None tools = guardrail.last_inputs.get("tools", []) assert len(tools) == 1 @@ -85,6 +93,412 @@ async def test_process_input_messages_handles_minimal_data(): assert tools[0]["function"]["name"] == "simple_tool" +class ArgumentMaskingGuardrail(CustomGuardrail): + """Unified guardrail that rewrites every text it is handed, like presidio does.""" + + def __init__( + self, + secret: str = "jane.doe@example.com", + replacement: str = "", + texts_override: list[str] | None = None, + **kwargs, + ): + kwargs.setdefault("guardrail_name", "argument-masking-mcp-guardrail") + super().__init__(**kwargs) + self.secret = secret + self.replacement = replacement + self.texts_override = texts_override + self.seen_texts: list[str] | None = None + + def _mask(self, text: str) -> str: + return text.replace(self.secret, self.replacement) + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.seen_texts = list(inputs.get("texts") or []) + if self.texts_override is not None: + inputs["texts"] = self.texts_override + else: + inputs["texts"] = [self._mask(text) for text in self.seen_texts] + return inputs + + +@pytest.fixture +def restore_callbacks(monkeypatch): + """Restore the process-wide state driving pre_call_hook through unified_guardrail. + + litellm.llms memoizes the guardrail translation mappings in a module global, and + ProxyLogging caches callback capabilities keyed on id()s of litellm.callbacks, + so leaving either populated leaks into unrelated tests in the same worker. + """ + monkeypatch.setattr(litellm, "callbacks", litellm.callbacks) + monkeypatch.setattr( + litellm_llms, + "endpoint_guardrail_translation_mappings", + litellm_llms.endpoint_guardrail_translation_mappings, + ) + yield + ProxyLogging._callback_capabilities_cache.clear() + + +@pytest.mark.asyncio +async def test_argument_strings_are_handed_to_the_guardrail(): + """A guardrail must see the argument values, not just the tool definition. + + Without this the guardrail is handed a name and an empty schema, so no + sensitive-data detection can ever fire on an MCP tool call. + """ + handler = MCPGuardrailTranslationHandler() + guardrail = MockGuardrail() + + data = { + "mcp_tool_name": "search", + "mcp_arguments": {"query": "contact jane.doe@example.com about the invoice"}, + } + + await handler.process_input_messages(data, guardrail) + + assert guardrail.last_inputs is not None + assert guardrail.last_inputs.get("texts") == ["contact jane.doe@example.com about the invoice"] + + +@pytest.mark.asyncio +async def test_masked_arguments_are_written_back_for_the_call_path(): + """A mask only takes effect once it lands in modified_arguments.""" + handler = MCPGuardrailTranslationHandler() + guardrail = ArgumentMaskingGuardrail() + + data = { + "mcp_tool_name": "search", + "mcp_arguments": {"query": "contact jane.doe@example.com about the invoice"}, + } + + result = await handler.process_input_messages(data, guardrail) + + masked = {"query": "contact about the invoice"} + assert result["modified_arguments"] == masked + assert result["mcp_arguments"] == masked + + +@pytest.mark.asyncio +async def test_nested_arguments_keep_their_shape_when_masked(): + """Masking rewrites string leaves in place and preserves non-string values.""" + handler = MCPGuardrailTranslationHandler() + guardrail = ArgumentMaskingGuardrail() + + arguments = { + "recipients": ["jane.doe@example.com", "ops@example.net"], + "envelope": {"reply_to": "jane.doe@example.com", "retries": 3, "urgent": True, "cc": None}, + "count": 2, + } + data = {"mcp_tool_name": "send_email", "mcp_arguments": arguments} + + result = await handler.process_input_messages(data, guardrail) + + assert guardrail.seen_texts == [ + "jane.doe@example.com", + "ops@example.net", + "jane.doe@example.com", + ] + assert result["modified_arguments"] == { + "recipients": ["", "ops@example.net"], + "envelope": {"reply_to": "", "retries": 3, "urgent": True, "cc": None}, + "count": 2, + } + + +@pytest.mark.asyncio +async def test_clean_arguments_are_not_overridden(): + """A guardrail that changes nothing must not set modified_arguments.""" + handler = MCPGuardrailTranslationHandler() + guardrail = ArgumentMaskingGuardrail() + + data = {"mcp_tool_name": "search", "mcp_arguments": {"query": "quarterly revenue"}} + + result = await handler.process_input_messages(data, guardrail) + + assert "modified_arguments" not in result + assert result["mcp_arguments"] == {"query": "quarterly revenue"} + + +@pytest.mark.asyncio +async def test_guardrail_returning_wrong_text_count_blocks_the_call(): + """Write-back is positional, so a length mismatch must block the call.""" + handler = MCPGuardrailTranslationHandler() + guardrail = ArgumentMaskingGuardrail(texts_override=["only", "two", "texts"]) + + arguments = {"query": "contact jane.doe@example.com about the invoice"} + data = {"mcp_tool_name": "search", "mcp_arguments": arguments} + + with pytest.raises(HTTPException) as exc_info: + await handler.process_input_messages(data, guardrail) + + assert exc_info.value.status_code == 400 + assert "modified_arguments" not in data + + +@pytest.mark.asyncio +async def test_deeply_nested_arguments_are_blocked_rather_than_skipped(): + """Arguments too deep to walk must block instead of passing unscanned.""" + handler = MCPGuardrailTranslationHandler() + guardrail = ArgumentMaskingGuardrail() + + nested: dict = {"leaf": "jane.doe@example.com"} + for _ in range(MAX_STRUCTURED_CONTENT_SCAN_DEPTH + 1): + nested = {"next": nested} + + data = {"mcp_tool_name": "search", "mcp_arguments": nested} + + with pytest.raises(HTTPException) as exc_info: + await handler.process_input_messages(data, guardrail) + + assert exc_info.value.status_code == 400 + + +class SelfWritingMaskingGuardrail(ArgumentMaskingGuardrail): + """Masks through ``texts`` and writes the masked arguments itself. + + The shape the bundled content filter guardrail already has: it rewrites + ``request_data["mcp_arguments"]`` from inside ``apply_guardrail`` as well as + returning masked texts. + """ + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + returned = await super().apply_guardrail(inputs, request_data, input_type, **kwargs) + arguments = request_data.get("mcp_arguments") or {} + masked = {key: self._mask(value) if isinstance(value, str) else value for key, value in arguments.items()} + request_data["mcp_arguments"] = masked + request_data["modified_arguments"] = masked + return returned + + +@pytest.mark.asyncio +async def test_guardrail_that_masks_the_arguments_itself_is_not_treated_as_a_conflict(): + """Converging on the same replacement is not an unmergeable rewrite. + + A guardrail that both returns masked texts and rewrites the arguments in + request_data must still mask, not be rejected as if a second guardrail had + clobbered the leaf. + """ + handler = MCPGuardrailTranslationHandler() + guardrail = SelfWritingMaskingGuardrail() + + data = { + "mcp_tool_name": "search", + "mcp_arguments": {"query": "contact jane.doe@example.com about the invoice"}, + } + + result = await handler.process_input_messages(data, guardrail) + + assert result["modified_arguments"] == {"query": "contact about the invoice"} + + +class ReshapingGuardrail(ArgumentMaskingGuardrail): + """Masks through ``texts`` while moving the secret to a different path.""" + + def __init__(self, reshaped: dict, **kwargs): + super().__init__(**kwargs) + self.reshaped = reshaped + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + returned = await super().apply_guardrail(inputs, request_data, input_type, **kwargs) + request_data["mcp_arguments"] = self.reshaped + return returned + + +@pytest.mark.asyncio +async def test_arguments_reshaped_under_the_guardrail_fail_closed(): + """A payload that no longer lines up leaf for leaf must block, not be written blind. + + Write-back pairs masked texts to leaves positionally, so a tree another guardrail + reshaped would take the redaction on the wrong value. + """ + handler = MCPGuardrailTranslationHandler() + guardrail = ReshapingGuardrail({"query": "contact jane.doe@example.com", "note": "added"}) + + data = { + "mcp_tool_name": "search", + "mcp_arguments": {"query": "contact jane.doe@example.com about the invoice"}, + } + + with pytest.raises(HTTPException) as exc_info: + await handler.process_input_messages(data, guardrail) + + assert exc_info.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_arguments_shortened_under_the_guardrail_fail_closed(): + handler = MCPGuardrailTranslationHandler() + guardrail = ReshapingGuardrail({"padding": "jane.doe@example.com"}) + + data = { + "mcp_tool_name": "search", + "mcp_arguments": {"padding": "x", "secret": "jane.doe@example.com"}, + } + + with pytest.raises(HTTPException) as exc_info: + await handler.process_input_messages(data, guardrail) + + assert exc_info.value.status_code == 400 + assert "jane.doe@example.com" not in str(data.get("modified_arguments")) + + +@pytest.mark.asyncio +async def test_a_renamed_argument_key_blocks_rather_than_dropping_the_mask(): + """The leak this closes: same text, new path, so the write-back would find nothing. + + Matching purely on position would see an unchanged value and write the mask to a + path that no longer exists, shipping the secret while reporting a clean scan. + """ + handler = MCPGuardrailTranslationHandler() + guardrail = ReshapingGuardrail({"renamed": "jane.doe@example.com", "other": "kept"}) + + data = { + "mcp_tool_name": "search", + "mcp_arguments": {"query": "jane.doe@example.com", "other": "kept"}, + } + + with pytest.raises(HTTPException) as exc_info: + await handler.process_input_messages(data, guardrail) + + assert exc_info.value.status_code == 400 + assert "jane.doe@example.com" not in str(data.get("modified_arguments")) + + +@pytest.mark.parametrize("run_in_parallel", [False, True]) +@pytest.mark.asyncio +async def test_masked_arguments_reach_the_outbound_mcp_call(restore_callbacks, monkeypatch, run_in_parallel): + """End to end over the real MCP pre-call path, not just the handler. + + Drives the same sequence mcp_server_manager.call_tool uses: + synthetic payload -> pre_call_hook -> arguments sent upstream. + + Covers run_in_parallel both ways: that path shares one payload snapshot and + discards whatever a guardrail returns, so the mask has to land on the caller's + dict rather than on a copy of it. + """ + guardrail = ArgumentMaskingGuardrail( + event_hook="pre_mcp_call", + default_on=True, + run_in_parallel=run_in_parallel, + ) + monkeypatch.setattr(litellm, "callbacks", [guardrail]) + + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + arguments = {"query": "contact jane.doe@example.com about the invoice"} + pre_hook_kwargs = { + "name": "search", + "arguments": arguments, + "server_name": "test-server", + "user_api_key_auth": UserAPIKeyAuth(api_key="sk-test", user_id="test-user"), + } + + request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs(pre_hook_kwargs) + synthetic_data = proxy_logging_obj._convert_mcp_to_llm_format(request_obj, pre_hook_kwargs) + + modified_data = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=pre_hook_kwargs["user_api_key_auth"], + data=synthetic_data, + call_type="call_mcp_tool", + ) + modified_kwargs = proxy_logging_obj._convert_mcp_hook_response_to_kwargs(modified_data, pre_hook_kwargs) + + assert modified_kwargs["arguments"] == {"query": "contact about the invoice"} + + +class SlowSubstitutionGuardrail(CustomGuardrail): + """Rewrites one substring, after a delay, so two instances genuinely interleave.""" + + def __init__(self, needle: str, replacement: str, delay: float, **kwargs): + super().__init__(**kwargs) + self.needle = needle + self.replacement = replacement + self.delay = delay + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + await asyncio.sleep(self.delay) + inputs["texts"] = [text.replace(self.needle, self.replacement) for text in (inputs.get("texts") or [])] + return inputs + + +def _two_interleaving_maskers(run_in_parallel: bool): + return [ + SlowSubstitutionGuardrail( + "jane.doe@example.com", + "", + 0.02, + guardrail_name="mask-email", + event_hook="pre_mcp_call", + default_on=True, + run_in_parallel=run_in_parallel, + ), + SlowSubstitutionGuardrail( + "415-555-0132", + "", + 0.04, + guardrail_name="mask-phone", + event_hook="pre_mcp_call", + default_on=True, + run_in_parallel=run_in_parallel, + ), + ] + + +async def _arguments_sent_upstream(arguments: dict): + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + pre_hook_kwargs = { + "name": "search", + "arguments": arguments, + "server_name": "test-server", + "user_api_key_auth": UserAPIKeyAuth(api_key="sk-test", user_id="test-user"), + } + request_obj = proxy_logging_obj._create_mcp_request_object_from_kwargs(pre_hook_kwargs) + modified_data = await proxy_logging_obj.pre_call_hook( + user_api_key_dict=pre_hook_kwargs["user_api_key_auth"], + data=proxy_logging_obj._convert_mcp_to_llm_format(request_obj, pre_hook_kwargs), + call_type="call_mcp_tool", + ) + return proxy_logging_obj._convert_mcp_hook_response_to_kwargs(modified_data, pre_hook_kwargs)["arguments"] + + +@pytest.mark.asyncio +async def test_two_sequential_guardrails_both_masks_survive(restore_callbacks, monkeypatch): + """The recommended config: each guardrail sees the previous one's output.""" + monkeypatch.setattr(litellm, "callbacks", _two_interleaving_maskers(run_in_parallel=False)) + + sent = await _arguments_sent_upstream({"note": "mail jane.doe@example.com or call 415-555-0132"}) + + assert sent == {"note": "mail or call "} + + +@pytest.mark.asyncio +async def test_two_parallel_guardrails_on_separate_arguments_both_masks_survive(restore_callbacks, monkeypatch): + """Concurrent rewrites of different leaves compose; neither is lost.""" + monkeypatch.setattr(litellm, "callbacks", _two_interleaving_maskers(run_in_parallel=True)) + + sent = await _arguments_sent_upstream({"email": "jane.doe@example.com", "phone": "415-555-0132"}) + + assert sent == {"email": "", "phone": ""} + + +@pytest.mark.asyncio +async def test_two_parallel_guardrails_on_one_argument_block_instead_of_losing_a_mask(restore_callbacks, monkeypatch): + """Unmergeable concurrent rewrites must fail closed, not ship one redaction. + + Both guardrails derive a full replacement string from the same snapshot, so + writing either result would silently discard the other's redaction and leak + the value it was configured to mask. + """ + monkeypatch.setattr(litellm, "callbacks", _two_interleaving_maskers(run_in_parallel=True)) + original = "mail jane.doe@example.com or call 415-555-0132" + + with pytest.raises(HTTPException) as exc_info: + await _arguments_sent_upstream({"note": original}) + + assert exc_info.value.status_code == 400 + assert "note" in str(exc_info.value.detail) + + class MaskingGuardrail(CustomGuardrail): """Guardrail that rewrites every scanned text, recording what it saw.""" @@ -190,8 +604,8 @@ async def test_process_output_response_handles_result_without_content(): @pytest.mark.asyncio -async def test_process_output_response_leaves_result_unmasked_on_text_count_mismatch(): - """A guardrail returning the wrong number of texts must not shuffle content.""" +async def test_process_output_response_blocks_on_text_count_mismatch(): + """A guardrail returning the wrong number of texts must block the result.""" handler = MCPGuardrailTranslationHandler() guardrail = MaskingGuardrail(masked_texts=[""]) result = CallToolResult( @@ -202,9 +616,10 @@ async def test_process_output_response_leaves_result_unmasked_on_text_count_mism isError=False, ) - returned = await handler.process_output_response(response=result, guardrail_to_apply=guardrail) + with pytest.raises(HTTPException) as exc_info: + await handler.process_output_response(response=result, guardrail_to_apply=guardrail) - assert [item.text for item in returned.content] == ["jane@example.com", "415-555-0132"] + assert exc_info.value.status_code == 400 class SubstitutingGuardrail(CustomGuardrail): diff --git a/type-discipline-budget.json b/type-discipline-budget.json index e35e470c979..e7186dfe186 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16430 + "limit": 16426 }, "LIT011": { "limit": 5506 From a46a076b2abd46b88f65d6d21d7afd9c052bb826 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 5 Sep 2026 14:22:26 -0700 Subject: [PATCH 7/7] fix(proxy): reject ambiguous name or alias keys in mcp_tool_permissions on write (#39947) Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../organization_endpoints.py | 6 + .../object_permission_utils.py | 71 ++++++++++- .../test_internal_user_endpoints.py | 2 + .../test_key_management_endpoints.py | 1 + .../test_organization_endpoints.py | 31 +++++ .../test_team_endpoints.py | 1 + .../test_object_permission_utils.py | 119 ++++++++++++++++++ 7 files changed, 229 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 0af9f816318..1e711b036d2 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -46,6 +46,7 @@ from litellm.proxy.management_endpoints.common_utils import ( from litellm.proxy.management_helpers.object_permission_utils import ( handle_update_object_permission_common, prepare_object_permission_upsert, + reject_ambiguous_mcp_tool_permission_keys, ) from litellm.proxy.management_helpers.utils import ( get_new_internal_user_defaults, @@ -606,6 +607,11 @@ async def _set_object_permission( return None if data.object_permission is not None: + await reject_ambiguous_mcp_tool_permission_keys( + new_mcp_tool_permissions=data.object_permission.mcp_tool_permissions, + existing_mcp_tool_permissions=None, + prisma_client=prisma_client, + ) created_object_permission: Final = await _table(ObjectPermissionRepository(prisma_client)).create( data=data.object_permission.model_dump(exclude_none=True), ) diff --git a/litellm/proxy/management_helpers/object_permission_utils.py b/litellm/proxy/management_helpers/object_permission_utils.py index a2fbf80422c..daab38d3662 100644 --- a/litellm/proxy/management_helpers/object_permission_utils.py +++ b/litellm/proxy/management_helpers/object_permission_utils.py @@ -5,10 +5,13 @@ organizations, teams, and keys. import json from collections.abc import Mapping, Sequence +from collections.abc import Set as AbstractSet from dataclasses import dataclass +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Optional from fastapi import HTTPException, status +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid @@ -103,6 +106,11 @@ async def prepare_object_permission_upsert( if existing_object_permission is not None else {} ) + await reject_ambiguous_mcp_tool_permission_keys( + new_mcp_tool_permissions=new_object_permission.get("mcp_tool_permissions"), + existing_mcp_tool_permissions=existing_fields.get("mcp_tool_permissions"), + prisma_client=prisma_client, + ) merged: Final[dict[str, object]] = { **existing_fields, **new_object_permission, @@ -194,6 +202,12 @@ async def _set_object_permission( k: v for k, v in permission_data.items() if v is not None and k != "object_permission_id" } + await reject_ambiguous_mcp_tool_permission_keys( + new_mcp_tool_permissions=clean_data.get("mcp_tool_permissions"), + existing_mcp_tool_permissions=None, + prisma_client=prisma_client, + ) + # Serialize mcp_tool_permissions to JSON string for GraphQL compatibility if "mcp_tool_permissions" in clean_data: clean_data["mcp_tool_permissions"] = safe_dumps(clean_data["mcp_tool_permissions"]) @@ -226,7 +240,7 @@ def _mcp_server_identifier_matches(server: Any, identifier: str) -> bool: async def _get_db_mcp_servers_by_identifiers( - identifiers: set[str], + identifiers: AbstractSet[str], prisma_client: PrismaClient | None, ) -> "Sequence[prisma_models.LiteLLM_MCPServerTable]": if prisma_client is None or not identifiers: @@ -245,7 +259,7 @@ async def _get_db_mcp_servers_by_identifiers( async def _resolve_mcp_server_identifiers_to_ids( - identifiers: set[str], + identifiers: AbstractSet[str], prisma_client: PrismaClient | None, ) -> dict[str, set[str]]: """ @@ -286,6 +300,59 @@ async def _resolve_mcp_server_identifiers_to_ids( return resolved +_MCP_TOOL_PERMISSIONS_ADAPTER: Final = TypeAdapter(dict[str, list[str] | None]) + + +def _mcp_tool_permission_entries(raw: object) -> Mapping[str, frozenset[str]]: + parsed: Final[Mapping[str, Sequence[str] | None]] = ( + _MCP_TOOL_PERMISSIONS_ADAPTER.validate_json(raw) + if isinstance(raw, str) + else _MCP_TOOL_PERMISSIONS_ADAPTER.validate_python(raw) + if isinstance(raw, Mapping) + else MappingProxyType({}) + ) + return MappingProxyType({identifier: frozenset(tools or ()) for identifier, tools in parsed.items()}) + + +async def reject_ambiguous_mcp_tool_permission_keys( + new_mcp_tool_permissions: object, + existing_mcp_tool_permissions: object, + prisma_client: PrismaClient | None, +) -> None: + """ + A name or alias shared by several MCP servers cannot key ``mcp_tool_permissions``: + the read path unions the entry into every match, so no edit can narrow one of + those servers without also changing the other. An exact server_id is never + ambiguous, even when another server uses that string as its alias. Entries the + row already stores with the same tool list are left alone, so unrelated edits + to such an entity still succeed. + + Raises HTTPException(400) naming the colliding servers. + """ + requested: Final = _mcp_tool_permission_entries(new_mcp_tool_permissions) + stored: Final = _mcp_tool_permission_entries(existing_mcp_tool_permissions) + resolved: Final = await _resolve_mcp_server_identifiers_to_ids( + identifiers=frozenset(identifier for identifier, tools in requested.items() if stored.get(identifier) != tools), + prisma_client=prisma_client, + ) + collisions: Final = "; ".join( + f"'{identifier}' matches MCP servers {sorted(server_ids)}" + for identifier, server_ids in sorted(resolved.items()) + if identifier not in server_ids and len(server_ids) > 1 + ) + if not collisions: + return + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ # mutable-ok: HTTPException.detail has no immutable form; same shape as the sibling errors here + "error": ( + f"Ambiguous mcp_tool_permissions key: {collisions}. " + "Key tool permissions by server_id when servers share a name or alias." + ) + }, + ) + + def _drop_stale_object_permission_mcp_servers( object_permission: ObjectPermissionDict, identifier_to_server_ids: dict[str, set[str]], diff --git a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py index 0d3ea5863a2..d1d669cae38 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_internal_user_endpoints.py @@ -3917,6 +3917,7 @@ def _object_permission_mocks(mocker, existing_object_permission_id=None): mock_prisma_client.db.litellm_objectpermissiontable.upsert = mocker.AsyncMock( return_value=SimpleNamespace(object_permission_id="perm-new") ) + mock_prisma_client.db.litellm_mcpservertable.find_many = mocker.AsyncMock(return_value=[]) mock_prisma_client.update_data = mocker.AsyncMock( return_value={"user_id": "target-user"} ) @@ -4146,6 +4147,7 @@ async def test_new_user_persists_the_requested_mcp_entitlement(mocker): mock_prisma_client.db.litellm_objectpermissiontable.create = mocker.AsyncMock( return_value=SimpleNamespace(object_permission_id="perm-created") ) + mock_prisma_client.db.litellm_mcpservertable.find_many = mocker.AsyncMock(return_value=[]) mock_prisma_client.db.litellm_usertable.find_first = mocker.AsyncMock( return_value=None ) diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 47571497f74..8766b1a1868 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -963,6 +963,7 @@ async def test_key_generation_with_mcp_tool_permissions(monkeypatch): mock_prisma_client.db = MagicMock() mock_prisma_client.db.litellm_objectpermissiontable = MagicMock() mock_prisma_client.db.litellm_objectpermissiontable.create = mock_create + mock_prisma_client.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) async def _insert_data_side_effect(*args, **kwargs): table_name = kwargs.get("table_name") diff --git a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py index da68492e3d7..4d13e054e46 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_organization_endpoints.py @@ -1183,6 +1183,37 @@ async def test_find_member_if_email_missing_row_raises_documented_400(): } +@pytest.mark.asyncio +async def test_new_organization_rejects_shared_alias_tool_permission_key(): + """/organization/new creates its permission row through its own helper, so the + ambiguous mcp_tool_permissions key check (LIT-4982) has to run there too.""" + from litellm.proxy._types import LiteLLM_ObjectPermissionBase, NewOrganizationRequest + from litellm.proxy.management_endpoints.organization_endpoints import ( + _set_object_permission, + ) + + prisma_client = MagicMock() + prisma_client.db.litellm_mcpservertable.find_many = AsyncMock( + return_value=[ + MagicMock(server_id="wiki-a-id", alias="wiki", server_name="wiki_a"), + MagicMock(server_id="wiki-b-id", alias="wiki", server_name="wiki_b"), + ] + ) + prisma_client.db.litellm_objectpermissiontable.create = AsyncMock() + data = NewOrganizationRequest( + organization_alias="org", + object_permission=LiteLLM_ObjectPermissionBase(mcp_tool_permissions={"wiki": ["ask_question"]}), + ) + + with pytest.raises(HTTPException) as exc_info: + await _set_object_permission(data=data, prisma_client=prisma_client) + + assert exc_info.value.status_code == 400 + assert "wiki-a-id" in str(exc_info.value.detail) + assert "wiki-b-id" in str(exc_info.value.detail) + prisma_client.db.litellm_objectpermissiontable.create.assert_not_called() + + def test_v2_update_organization_is_in_openapi_schema(): """PATCH /v2/organization/{organization_id} is documented in the generated OpenAPI spec.""" from fastapi import FastAPI diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 46678c8ff6a..051e6bed4fd 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -651,6 +651,7 @@ async def test_new_team_with_mcp_tool_permissions(mock_db_client, mock_admin_aut mock_db_client.db.litellm_objectpermissiontable = MagicMock() mock_db_client.db.litellm_objectpermissiontable.create = mock_obj_perm_create + mock_db_client.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) # Mock model table mock_db_client.db.litellm_modeltable = MagicMock() diff --git a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py index f2b6b799271..d7ebb1f60bf 100644 --- a/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py +++ b/tests/test_litellm/proxy/management_helpers/test_object_permission_utils.py @@ -17,6 +17,7 @@ from litellm.proxy.management_helpers.object_permission_utils import ( _resolve_team_allowed_mcp_servers, _set_object_permission, enforce_all_proxy_mcp_servers_grant_is_admin_only, + prepare_object_permission_upsert, validate_key_mcp_servers_against_team, validate_key_search_tools_against_team, validate_key_vector_stores_against_team, @@ -41,6 +42,7 @@ async def test_set_object_permission(): mock_prisma_client.db.litellm_objectpermissiontable.create = AsyncMock( return_value=mock_created_permission ) + mock_prisma_client.db.litellm_mcpservertable.find_many = AsyncMock(return_value=[]) # Test data with object_permission data_json = { @@ -1349,6 +1351,123 @@ async def test_validate_key_update_sentinels_do_not_grandfather(monkeypatch): assert exc_info.value.status_code == 403 +# ---- Tests for rejecting ambiguous mcp_tool_permissions keys on write (LIT-4982) ---- + + +_SHARED_ALIAS_DB_SERVERS = ( + _make_mock_mcp_server("wiki-a-id", alias="wiki", server_name="wiki_a"), + _make_mock_mcp_server("wiki-b-id", alias="wiki", server_name="wiki_b"), + _make_mock_mcp_server("gh-a-id", alias="gh_a", server_name="github"), + _make_mock_mcp_server("gh-b-id", alias="gh_b", server_name="github"), + _make_mock_mcp_server("solo-id", alias="solo", server_name="Solo Server"), + _make_mock_mcp_server("shadow-id", alias="solo-id", server_name="shadow"), +) + + +def _make_ambiguity_prisma(existing_tool_permissions=None): + """Mock prisma client whose MCP server table holds _SHARED_ALIAS_DB_SERVERS and whose + object permission row (if any) stores the given mcp_tool_permissions JSON string.""" + mock_prisma = MagicMock() + mock_prisma.db.litellm_mcpservertable.find_many = AsyncMock(return_value=list(_SHARED_ALIAS_DB_SERVERS)) + mock_prisma.db.litellm_objectpermissiontable.create = AsyncMock( + return_value=MagicMock(object_permission_id="perm-id") + ) + existing_row = None + if existing_tool_permissions is not None: + existing_row = MagicMock() + existing_row.model_dump.return_value = { + "object_permission_id": "perm-id", + "mcp_tool_permissions": json.dumps(existing_tool_permissions), + } + mock_prisma.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=existing_row) + return mock_prisma + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "identifier, colliding_ids", + [("wiki", ("wiki-a-id", "wiki-b-id")), ("github", ("gh-a-id", "gh-b-id"))], +) +async def test_set_object_permission_rejects_shared_alias_or_name_tool_permission_key(identifier, colliding_ids): + """An alias or server_name two servers share cannot key mcp_tool_permissions on + create: the write is rejected with 400 naming both servers and nothing is persisted.""" + mock_prisma = _make_ambiguity_prisma() + data_json = {"object_permission": {"mcp_tool_permissions": {identifier: ["read_wiki_structure"]}}} + + with pytest.raises(HTTPException) as exc_info: + await _set_object_permission(data_json=data_json, prisma_client=mock_prisma) + + assert exc_info.value.status_code == 400 + assert all(server_id in str(exc_info.value.detail) for server_id in colliding_ids) + mock_prisma.db.litellm_objectpermissiontable.create.assert_not_called() + + +@pytest.mark.asyncio +async def test_prepare_object_permission_upsert_rejects_shared_alias_tool_permission_key(): + """The update seam shared by key/team/org/user/customer/agent rejects a new + shared-alias key when the existing row does not already hold it.""" + mock_prisma = _make_ambiguity_prisma(existing_tool_permissions={"solo-id": ["tool1"]}) + + with pytest.raises(HTTPException) as exc_info: + await prepare_object_permission_upsert( + new_object_permission={"mcp_tool_permissions": {"wiki": ["ask_question"]}}, + existing_object_permission_id="perm-id", + prisma_client=mock_prisma, + ) + + assert exc_info.value.status_code == 400 + assert "'wiki'" in str(exc_info.value.detail) + + +@pytest.mark.asyncio +async def test_unambiguous_tool_permission_keys_persist_verbatim(): + """Exact ids (even when another server uses that id string as its alias), + unique aliases, and an id plus alias pointing at one server all still write.""" + mock_prisma = _make_ambiguity_prisma() + tool_permissions = { + "wiki-a-id": ["ask_question"], + "wiki-b-id": ["read_wiki_structure"], + "solo-id": ["tool1"], + "solo": ["tool2"], + "Solo Server": ["tool3"], + } + + upsert = await prepare_object_permission_upsert( + new_object_permission={"mcp_tool_permissions": dict(tool_permissions)}, + existing_object_permission_id=None, + prisma_client=mock_prisma, + ) + + assert json.loads(upsert.record["mcp_tool_permissions"]) == tool_permissions + + +@pytest.mark.asyncio +async def test_stored_ambiguous_tool_permission_key_is_grandfathered_until_changed(): + """A shared-alias entry already on the row may be re-sent unchanged so unrelated + edits succeed, but changing its tool list is rejected.""" + mock_prisma = _make_ambiguity_prisma(existing_tool_permissions={"wiki": ["read_wiki_structure"]}) + + upsert = await prepare_object_permission_upsert( + new_object_permission={ + "mcp_tool_permissions": {"wiki": ["read_wiki_structure"], "solo-id": ["tool1"]}, + }, + existing_object_permission_id="perm-id", + prisma_client=mock_prisma, + ) + assert json.loads(upsert.record["mcp_tool_permissions"]) == { + "wiki": ["read_wiki_structure"], + "solo-id": ["tool1"], + } + + with pytest.raises(HTTPException) as exc_info: + await prepare_object_permission_upsert( + new_object_permission={"mcp_tool_permissions": {"wiki": ["ask_question"]}}, + existing_object_permission_id="perm-id", + prisma_client=mock_prisma, + ) + assert exc_info.value.status_code == 400 + + def test_object_permission_dict_mirrors_pydantic_model(): """ObjectPermissionDict must stay field-for-field aligned with LiteLLM_ObjectPermissionBase. If a new field is added to the Pydantic