mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
chore: merge litellm_internal_staging into litellm_anthropic_wif_backend
This commit is contained in:
mateo-berri
commit
3f36fe396a
14 changed files with 3940 additions and 224 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.<name> bound to the SDK function even after the import
|
||||
machinery binds the identically named litellm.<name> 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 <package>.<name> 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.<name> 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.
|
||||
|
|
|
|||
File diff suppressed because it is too large
Load diff
|
|
@ -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}")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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 <EMAIL_ADDRESS>.
|
||||
|
||||
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 <EMAIL_ADDRESS> 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 = "<no successful response yet>"
|
||||
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"<non-Success result: {result}>"
|
||||
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)
|
||||
|
|
|
|||
167
tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py
Normal file
167
tests/e2e/guardrails/test_tool_permission_guardrail_e2e.py
Normal file
|
|
@ -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]}"
|
||||
)
|
||||
|
|
@ -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 <entity>/<project>); 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")
|
||||
|
|
|
|||
192
tests/e2e/logging/test_weave_log_e2e.py
Normal file
192
tests/e2e/logging/test_weave_log_e2e.py
Normal file
|
|
@ -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
|
||||
282
tests/e2e/logging/weave_reader.py
Normal file
282
tests/e2e/logging/weave_reader.py
Normal file
|
|
@ -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:///<entity>/<project>/op/<op>:<digest>``."""
|
||||
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 <entity>/<project>): 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)
|
||||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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.<submodule> 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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue