{announcement.title}
+{announcement.description}
- Read the docs
+ {announcement.linkLabel}
- {hasUnread ? (
-
- ([^<]+)")
+_TRANSIENT_ERRORS: Final = (httpx.TransportError, LiteLLMTimeout)
+_UPLOAD_ATTEMPTS: Final = 3
+_UPLOAD_ATTEMPT_TIMEOUT: Final = httpx.Timeout(S3_UPLOAD_ATTEMPT_TIMEOUT_SECONDS)
+
+PutOutcome: TypeAlias = httpx.Response | Exception
@dataclass(frozen=True, slots=True)
@@ -105,6 +113,35 @@ def _s3_key_parent(s3_object_key: str) -> str:
return s3_object_key.rsplit("/", 1)[0] if "/" in s3_object_key else ""
+def _backoff_seconds(attempt: int) -> float:
+ return (1 << attempt) * (0.5 + random.random())
+
+
+def _log_retry(level: int, attempt: int, reason: str, wait: float, s3_object_key: str) -> None:
+ verbose_logger.log(
+ level,
+ "S3 upload attempt %s/%s failed with %s, retrying in %.2fs key=%s",
+ attempt + 1,
+ _UPLOAD_ATTEMPTS,
+ reason,
+ wait,
+ s3_object_key,
+ )
+
+
+def _response_or_raise(outcome: PutOutcome) -> httpx.Response:
+ if isinstance(outcome, Exception):
+ raise outcome
+ outcome.raise_for_status()
+ return outcome
+
+
+def _throttled(outcome: PutOutcome) -> bool:
+ if isinstance(outcome, Exception):
+ return True
+ return outcome.status_code in (429, 503) or _s3_error_code(outcome) == "SlowDown"
+
+
class S3BatchUploadError(Exception):
def __init__(self, failed: int, total: int) -> None:
self.failed = failed
@@ -566,7 +603,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
url: Final = self._build_object_url(batch_logging_element.s3_object_key)
- async def signed_put(prepared: _PreparedPut) -> httpx.Response:
+ async def signed_put(prepared: _PreparedPut) -> PutOutcome:
credentials: Final = await asyncified_get_credentials(
aws_access_key_id=self.s3_aws_access_key_id,
aws_secret_access_key=self.s3_aws_secret_access_key,
@@ -582,34 +619,16 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
self._sign_put, credentials, url, prepared.json_string, prepared.headers
)
try:
- return await self.async_httpx_client.put(url, data=prepared.json_string, headers=signed_headers)
+ return await self.async_httpx_client.put(
+ url, data=prepared.json_string, headers=signed_headers, timeout=_UPLOAD_ATTEMPT_TIMEOUT
+ )
except httpx.HTTPStatusError as error:
return error.response
+ except _TRANSIENT_ERRORS as error:
+ return error
- max_retries: Final = 3
prepared: Final = self._prepare_put(batch_logging_element)
- for attempt in range(max_retries):
- response = await self._recorded_put(partial(signed_put, prepared))
- if (
- response.status_code in _RETRYABLE_STATUSES
- and not (self.s3_drop_on_terminal_error and _is_terminal(response))
- and attempt < max_retries - 1
- ):
- wait_time = 1 << attempt # 1s, 2s
- verbose_logger.log(
- logging.DEBUG if _in_flush.get() else logging.WARNING,
- "S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s",
- response.status_code,
- wait_time,
- attempt + 1,
- max_retries,
- batch_logging_element.s3_object_key,
- )
- self._flush_retries += 1
- await asyncio.sleep(wait_time)
- continue
- response.raise_for_status()
- break
+ await self._put_with_retries(partial(signed_put, prepared), batch_logging_element.s3_object_key)
except Exception as e:
verbose_logger.exception("Error uploading to s3: %s", e)
self.handle_callback_failure(callback_name="S3Logger")
@@ -678,9 +697,13 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
)
dropped: Final = len(failed) - len(requeued)
if dropped:
+ terminal: Final = sum(1 for _, outcome in failed if outcome == "dropped")
verbose_logger.warning(
- "s3 logging: %s uploads dropped (terminal or retrying longer than s3_max_retry_age_seconds=%s)",
+ "s3 logging: %s uploads dropped this flush: %s rejected as terminal by S3, "
+ "%s retrying longer than s3_max_retry_age_seconds=%s",
dropped,
+ terminal,
+ dropped - terminal,
self.s3_max_retry_age_seconds,
)
if not requeued:
@@ -728,22 +751,46 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
return "dropped"
return "retry"
- async def _recorded_put(self, signed_put: Callable[[], Awaitable[httpx.Response]]) -> httpx.Response:
+ async def _recorded_put(self, signed_put: Callable[[], Awaitable[PutOutcome]]) -> PutOutcome:
+ outcome: Final = await signed_put()
limiter: Final = self._upload_limiter
- adaptive: Final = limiter if isinstance(limiter, AdaptiveConcurrencyLimiter) else None
- try:
- response: Final = await signed_put()
- except Exception:
- if adaptive is not None:
- adaptive.record(PutSample(throttled=True))
- raise
- if adaptive is not None:
- adaptive.record(
- PutSample(
- throttled=response.status_code in (429, 503) or _s3_error_code(response) == "SlowDown",
- )
- )
- return response
+ if isinstance(limiter, AdaptiveConcurrencyLimiter):
+ limiter.record(PutSample(throttled=_throttled(outcome)))
+ return outcome
+
+ def _retry_reason(self, outcome: PutOutcome) -> str | None:
+ if isinstance(outcome, Exception):
+ return type(outcome).__name__
+ if outcome.status_code not in _RETRYABLE_STATUSES:
+ return None
+ if self.s3_drop_on_terminal_error and _is_terminal(outcome):
+ return None
+ return str(outcome.status_code)
+
+ async def _put_with_retries(
+ self, signed_put: Callable[[], Awaitable[PutOutcome]], s3_object_key: str
+ ) -> httpx.Response:
+ for attempt in range(_UPLOAD_ATTEMPTS - 1):
+ outcome = await self._recorded_put(signed_put)
+ reason = self._retry_reason(outcome)
+ if reason is None:
+ return _response_or_raise(outcome)
+ wait = _backoff_seconds(attempt)
+ _log_retry(logging.DEBUG if _in_flush.get() else logging.WARNING, attempt, reason, wait, s3_object_key)
+ self._flush_retries += 1
+ await asyncio.sleep(wait)
+ return _response_or_raise(await self._recorded_put(signed_put))
+
+ def _sync_put_with_retries(self, signed_put: Callable[[], PutOutcome], s3_object_key: str) -> httpx.Response:
+ for attempt in range(_UPLOAD_ATTEMPTS - 1):
+ outcome = signed_put()
+ reason = self._retry_reason(outcome)
+ if reason is None:
+ return _response_or_raise(outcome)
+ wait = _backoff_seconds(attempt)
+ _log_retry(logging.WARNING, attempt, reason, wait, s3_object_key)
+ time.sleep(wait)
+ return _response_or_raise(signed_put())
def _batch_file_elements(self, batch: tuple[s3BatchLoggingElement, ...]) -> tuple[s3BatchLoggingElement, ...]:
now: Final = datetime.now(timezone.utc)
@@ -879,7 +926,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
params=({"ssl_verify": self.s3_verify} if self.s3_verify is not None else None)
)
- def signed_put(prepared_put: _PreparedPut) -> httpx.Response:
+ def signed_put(prepared_put: _PreparedPut) -> PutOutcome:
credentials: Final = self.get_credentials(
aws_access_key_id=self.s3_aws_access_key_id,
aws_secret_access_key=self.s3_aws_secret_access_key,
@@ -887,29 +934,14 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM):
aws_region_name=self.s3_region_name,
)
signed_headers: Final = self._sign_put(credentials, url, prepared_put.json_string, prepared_put.headers)
- return httpx_client.put(url, data=prepared_put.json_string, headers=signed_headers)
-
- max_retries: Final = 3
- for attempt in range(max_retries):
- response = signed_put(prepared)
- if (
- response.status_code in _RETRYABLE_STATUSES
- and not (self.s3_drop_on_terminal_error and _is_terminal(response))
- and attempt < max_retries - 1
- ):
- wait_time = 1 << attempt # 1s, 2s
- verbose_logger.warning(
- "S3 upload returned %s, retrying in %ss (attempt %s/%s) key=%s",
- response.status_code,
- wait_time,
- attempt + 1,
- max_retries,
- batch_logging_element.s3_object_key,
+ try:
+ return httpx_client.put(
+ url, data=prepared_put.json_string, headers=signed_headers, timeout=_UPLOAD_ATTEMPT_TIMEOUT
)
- time.sleep(wait_time)
- continue
- response.raise_for_status()
- break
+ except _TRANSIENT_ERRORS as error:
+ return error
+
+ self._sync_put_with_retries(partial(signed_put, prepared), batch_logging_element.s3_object_key)
except Exception as e:
verbose_logger.exception("Error uploading to s3: %s", e)
self.handle_callback_failure(callback_name="S3Logger")
diff --git a/litellm/litellm_core_utils/env_utils.py b/litellm/litellm_core_utils/env_utils.py
index d641884b4cd..6883af98bbd 100644
--- a/litellm/litellm_core_utils/env_utils.py
+++ b/litellm/litellm_core_utils/env_utils.py
@@ -4,23 +4,65 @@ Utility helpers for reading and parsing environment variables.
import logging
import os
-from typing import Final
+from typing import Final, TypeVar
+
+_NumberT: Final = TypeVar("_NumberT", int, float)
+
+# Marked before logging: the LiteLLM loggers' own filters read env knobs through these helpers, so
+# warning first would re-enter the helper from inside the log call and recurse until RecursionError.
+_reported_unparsable_values: Final[set[tuple[str, str]]] = set() # mutable-ok: process-wide warn-once registry
+
+
+def _warn_unparsable_once(env_var: str, raw: str, number_type: type[int] | type[float]) -> None:
+ key: Final = (env_var, raw)
+ if key in _reported_unparsable_values:
+ return
+ _reported_unparsable_values.add(key)
+ logging.getLogger("LiteLLM").warning(
+ "%s=%r is not a valid %s. Treating it as unset.", env_var, raw, number_type.__name__
+ )
+
+
+def _get_env_number_or_none(env_var: str, number_type: type[_NumberT]) -> _NumberT | None:
+ raw: Final = os.getenv(env_var)
+ if raw is None or not raw.strip():
+ return None
+ try:
+ return number_type(raw.strip())
+ except ValueError:
+ _warn_unparsable_once(env_var, raw, number_type)
+ return None
+
+
+def get_env_int_or_none(env_var: str) -> int | None:
+ """Parse an environment variable as an integer, returning None when it is unset, blank or unparsable.
+
+ Use this instead of `get_env_int` when callers must distinguish "explicitly configured"
+ from "left at the default", for example when an override should take precedence over a
+ value resolved from somewhere else.
+ """
+ return _get_env_number_or_none(env_var, int)
+
+
+def get_env_float_or_none(env_var: str) -> float | None:
+ """Float counterpart of `get_env_int_or_none`."""
+ return _get_env_number_or_none(env_var, float)
def get_env_int(env_var: str, default: int) -> int:
- """Parse an environment variable as an integer, falling back to default on invalid values.
+ """Parse an environment variable as an integer, falling back to default when it is unset, blank or unparsable.
- Handles empty strings, whitespace, and non-numeric values gracefully
- so that misconfiguration doesn't crash the process at import time.
+ An unparsable value is warned about once, so misconfiguration neither crashes the process at
+ import time nor goes unnoticed.
"""
- raw = os.getenv(env_var)
- if raw is None:
- return default
- raw = raw.strip()
- try:
- return int(raw)
- except (ValueError, TypeError):
- return default
+ value: Final = get_env_int_or_none(env_var)
+ return default if value is None else value
+
+
+def get_env_float(env_var: str, default: float) -> float:
+ """Float counterpart of `get_env_int`."""
+ value: Final = get_env_float_or_none(env_var)
+ return default if value is None else value
def get_env_int_in_range(env_var: str, default: int, minimum: int, maximum: int) -> int:
@@ -41,19 +83,3 @@ def get_env_int_in_range(env_var: str, default: int, minimum: int, maximum: int)
default,
)
return default
-
-
-def get_env_int_or_none(env_var: str) -> int | None:
- """Parse an environment variable as an integer, returning None when it is unset or unusable.
-
- Use this instead of `get_env_int` when callers must distinguish "explicitly configured"
- from "left at the default", for example when an override should take precedence over a
- value resolved from somewhere else.
- """
- raw: Final = os.getenv(env_var)
- if raw is None:
- return None
- try:
- return int(raw.strip())
- except (ValueError, TypeError):
- return None
diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index 7b965dea28a..9b10b95956f 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -22,7 +22,6 @@ from httpx import Response
from pydantic import BaseModel, JsonValue
import litellm
-from litellm import _custom_logger_compatible_callbacks_literal
from litellm._internal_context import post_response_phase
from litellm._logging import (
is_debugging_on,
@@ -652,6 +651,10 @@ class Logging(LiteLLMLoggingBaseClass):
self.stream = stream
self.start_time = start_time # log the call start time
self.call_type = call_type
+ self.relays_llm_passthrough: Final = call_type in (
+ CallTypes.llm_passthrough_route.value,
+ CallTypes.allm_passthrough_route.value,
+ )
self.litellm_call_id = litellm_call_id
self.litellm_trace_id: str = litellm_trace_id if litellm_trace_id else str(uuid.uuid4())
@@ -2389,10 +2392,7 @@ class Logging(LiteLLMLoggingBaseClass):
service_tier=ws_service_tier,
)
- elif (
- self.call_type == CallTypes.llm_passthrough_route.value
- or self.call_type == CallTypes.allm_passthrough_route.value
- ) and isinstance(result, Response):
+ elif self.relays_llm_passthrough and isinstance(result, Response):
from litellm.utils import ProviderConfigManager
provider_config: Final = ProviderConfigManager.get_provider_passthrough_config(
@@ -2589,7 +2589,9 @@ class Logging(LiteLLMLoggingBaseClass):
logging_result: Final = self.normalize_logging_result(result=result)
- if isinstance(result, Response) and isinstance(logging_result, (ModelResponse, EmbeddingResponse)):
+ if isinstance(result, Response) and isinstance(
+ logging_result, (ModelResponse, EmbeddingResponse, TextCompletionResponse)
+ ):
result = logging_result
if standard_logging_object is None and result is not None and self.stream is not True:
@@ -4575,712 +4577,738 @@ def set_callbacks(
raise e
+def build_custom_logger_compatible_class(
+ logging_integration: str,
+ internal_usage_cache: DualCache | None,
+ llm_router: object,
+ custom_logger_init_args: dict[str, object] | None = None,
+) -> CustomLogger | None:
+ custom_logger_init_args_value: Final[dict[str, object]] = custom_logger_init_args or {}
+ if logging_integration == "agentops": # Add AgentOps initialization
+ _v2 = _maybe_construct_otel_v2("agentops", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2
+ for callback in _in_memory_loggers:
+ if isinstance(callback, AgentOps):
+ return callback
+
+ agentops_logger: Final = AgentOps()
+ _in_memory_loggers.append(agentops_logger)
+ return agentops_logger
+ elif logging_integration == "lago":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, LagoLogger):
+ return callback
+
+ lago_logger: Final = LagoLogger()
+ _in_memory_loggers.append(lago_logger)
+ return lago_logger
+ elif logging_integration == "openmeter":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, OpenMeterLogger):
+ return callback
+
+ _openmeter_logger: Final = OpenMeterLogger()
+ _in_memory_loggers.append(_openmeter_logger)
+ return _openmeter_logger
+ elif logging_integration == "posthog":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, PostHogLogger):
+ return callback
+
+ _posthog_logger: Final = PostHogLogger()
+ _in_memory_loggers.append(_posthog_logger)
+ return _posthog_logger
+ elif logging_integration == "braintrust":
+ from litellm.integrations.braintrust_logging import BraintrustLogger
+
+ for callback in _in_memory_loggers:
+ if isinstance(callback, BraintrustLogger):
+ return callback
+
+ braintrust_logger: Final = BraintrustLogger()
+ _in_memory_loggers.append(braintrust_logger)
+ return braintrust_logger
+ elif logging_integration == "langsmith":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, LangsmithLogger):
+ return callback
+
+ _langsmith_logger: Final = LangsmithLogger()
+ _in_memory_loggers.append(_langsmith_logger)
+ return _langsmith_logger
+ elif logging_integration == "argilla":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, ArgillaLogger):
+ return callback
+
+ _argilla_logger: Final = ArgillaLogger()
+ _in_memory_loggers.append(_argilla_logger)
+ return _argilla_logger
+ elif logging_integration == "literalai":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, LiteralAILogger):
+ return callback
+
+ _literalai_logger: Final = LiteralAILogger()
+ _in_memory_loggers.append(_literalai_logger)
+ return _literalai_logger
+ elif logging_integration == "litellm_agent":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, LiteLLMAgentModelResolver):
+ return callback
+
+ _litellm_agent_resolver: Final = LiteLLMAgentModelResolver()
+ _in_memory_loggers.append(_litellm_agent_resolver)
+ return _litellm_agent_resolver
+ elif logging_integration == "prometheus":
+ PrometheusLogger: Final = _get_cached_prometheus_logger()
+
+ for callback in _in_memory_loggers:
+ if isinstance(callback, PrometheusLogger):
+ return callback
+
+ _prometheus_logger: Final = PrometheusLogger()
+ _in_memory_loggers.append(_prometheus_logger)
+ return _prometheus_logger
+ elif logging_integration == "datadog":
+ # Check if team-scoped credentials are provided
+ _dd_api_key: Final = custom_logger_init_args_value.get("dd_api_key")
+ _dd_site: Final = custom_logger_init_args_value.get("dd_site")
+ _dd_agent_host: Final = custom_logger_init_args_value.get("dd_agent_host")
+ _dd_agent_port: Final = custom_logger_init_args_value.get("dd_agent_port")
+
+ if _dd_api_key or _dd_site or _dd_agent_host:
+ # Team-scoped credentials: use DynamicLoggingCache for per-credential isolation
+ from litellm.integrations.datadog.datadog_team_handler import (
+ DataDogHandler,
+ )
+
+ return DataDogHandler.get_datadog_logger_for_request(
+ standard_callback_dynamic_params=cast( # cast-ok: callback options come from dynamic config
+ StandardCallbackDynamicParams, custom_logger_init_args_value
+ ),
+ in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
+ )
+
+ # Global (env-var based): reuse cached instance
+ for callback in _in_memory_loggers:
+ if isinstance(callback, DataDogLogger):
+ return callback
+
+ _datadog_logger: Final = DataDogLogger()
+ _in_memory_loggers.append(_datadog_logger)
+ return _datadog_logger
+ elif logging_integration == "datadog_metrics":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, DatadogMetricsLogger):
+ return callback
+
+ _datadog_metrics_logger: Final = DatadogMetricsLogger()
+ _in_memory_loggers.append(_datadog_metrics_logger)
+ return _datadog_metrics_logger
+ elif logging_integration == "datadog_llm_observability":
+ _datadog_llm_obs_logger: Final = DataDogLLMObsLogger()
+ _in_memory_loggers.append(_datadog_llm_obs_logger)
+ return _datadog_llm_obs_logger
+ elif logging_integration == "azure_sentinel":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, AzureSentinelLogger):
+ return callback
+
+ _azure_sentinel_logger: Final = AzureSentinelLogger()
+ _in_memory_loggers.append(_azure_sentinel_logger)
+ return _azure_sentinel_logger
+ elif logging_integration == "gcs_bucket":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, GCSBucketLogger):
+ return callback
+
+ _gcs_bucket_logger: Final = GCSBucketLogger()
+ _in_memory_loggers.append(_gcs_bucket_logger)
+ return _gcs_bucket_logger
+ elif logging_integration == "s3_v2":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, S3V2Logger):
+ return callback
+
+ _s3_v2_logger: Final = S3V2Logger()
+ _in_memory_loggers.append(_s3_v2_logger)
+ return _s3_v2_logger
+ elif logging_integration == "clickhouse":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, ClickHouseSpendLogger):
+ return callback
+
+ _clickhouse_spend_logger: Final = ClickHouseSpendLogger()
+ _in_memory_loggers.append(_clickhouse_spend_logger)
+ return _clickhouse_spend_logger
+ elif logging_integration == "pointfive":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, PointFiveLogger):
+ return callback
+
+ _pointfive_logger: Final = PointFiveLogger()
+ _in_memory_loggers.append(_pointfive_logger)
+ return _pointfive_logger
+ elif logging_integration == "zerobus":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, ZerobusLogger):
+ return callback
+
+ _zerobus_logger: Final = ZerobusLogger()
+ _in_memory_loggers.append(_zerobus_logger)
+ return _zerobus_logger
+ elif logging_integration == "aws_sqs":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, SQSLogger):
+ return callback
+
+ _aws_sqs_logger: Final = SQSLogger()
+ _in_memory_loggers.append(_aws_sqs_logger)
+ return _aws_sqs_logger
+ elif logging_integration == "azure_storage":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, AzureBlobStorageLogger):
+ return callback
+
+ _azure_storage_logger: Final = AzureBlobStorageLogger()
+ _in_memory_loggers.append(_azure_storage_logger)
+ return _azure_storage_logger
+ elif logging_integration == "opik":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, OpikLogger):
+ return callback
+
+ _opik_logger: Final = OpikLogger()
+ _in_memory_loggers.append(_opik_logger)
+ return _opik_logger
+ elif logging_integration == "arize":
+ _v2 = _maybe_construct_otel_v2("arize", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2
+ from litellm.integrations.opentelemetry import (
+ OpenTelemetry,
+ OpenTelemetryConfig,
+ )
+
+ arize_config: Final = ArizeLogger.get_arize_config()
+ if arize_config.endpoint is None:
+ raise ValueError(
+ "No valid endpoint found for Arize, please set 'ARIZE_ENDPOINT' to your GRPC endpoint or 'ARIZE_HTTP_ENDPOINT' to your HTTP endpoint"
+ )
+ otel_config = OpenTelemetryConfig(
+ exporter=arize_config.protocol,
+ endpoint=arize_config.endpoint,
+ service_name=arize_config.project_name,
+ )
+
+ os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
+ f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}"
+ )
+ for callback in _in_memory_loggers:
+ if isinstance(callback, ArizeLogger) and callback.callback_name == "arize":
+ return callback
+ _arize_otel_logger: Final = ArizeLogger(config=otel_config, callback_name="arize")
+ _in_memory_loggers.append(_arize_otel_logger)
+ return _arize_otel_logger
+ elif logging_integration == "arize_phoenix":
+ _v2 = _maybe_construct_otel_v2("arize_phoenix", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2
+ from litellm.integrations.opentelemetry import (
+ OpenTelemetry,
+ OpenTelemetryConfig,
+ )
+
+ arize_phoenix_config: Final = ArizePhoenixLogger.get_arize_phoenix_config()
+ otel_config = OpenTelemetryConfig(
+ exporter=arize_phoenix_config.protocol,
+ endpoint=arize_phoenix_config.endpoint,
+ headers=arize_phoenix_config.otlp_auth_headers,
+ )
+
+ # auth can be disabled on local deployments of arize phoenix
+ if arize_phoenix_config.otlp_auth_headers is not None:
+ os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = arize_phoenix_config.otlp_auth_headers
+
+ for callback in _in_memory_loggers:
+ if isinstance(callback, ArizePhoenixLogger) and callback.callback_name == "arize_phoenix":
+ return callback
+ _arize_phoenix_otel_logger: Final = ArizePhoenixLogger(config=otel_config, callback_name="arize_phoenix")
+ _in_memory_loggers.append(_arize_phoenix_otel_logger)
+ return _arize_phoenix_otel_logger
+ elif logging_integration == "levo":
+ _v2 = _maybe_construct_otel_v2("levo", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2
+ from litellm.integrations.levo.levo import LevoLogger
+ from litellm.integrations.opentelemetry import (
+ OpenTelemetry,
+ OpenTelemetryConfig,
+ )
+
+ levo_config: Final = LevoLogger.get_levo_config()
+ otel_config = OpenTelemetryConfig(
+ exporter=levo_config.protocol,
+ endpoint=levo_config.endpoint,
+ headers=levo_config.otlp_auth_headers,
+ )
+
+ # Check if LevoLogger instance already exists
+ for callback in _in_memory_loggers:
+ if isinstance(callback, LevoLogger) and callback.callback_name == "levo":
+ return callback
+
+ _levo_otel_logger: Final = LevoLogger(config=otel_config, callback_name="levo")
+ _in_memory_loggers.append(_levo_otel_logger)
+ return _levo_otel_logger
+ elif logging_integration == "otel":
+ # Gate the new typed V2 adapter behind LITELLM_OTEL_V2. When off,
+ # the legacy 3,227-line god-class is used unchanged. The two are
+ # never registered simultaneously — the dedup loop below treats
+ # any module under ``litellm.integrations.otel`` or
+ # ``litellm.integrations.opentelemetry`` as "the OTel callback".
+ from litellm.integrations.otel.model.config import is_otel_v2_enabled
+
+ if is_otel_v2_enabled():
+ from litellm.integrations.otel.logger import OpenTelemetryV2, build_otel_v2_logger
+ from litellm.integrations.otel.model.config import OpenTelemetryV2Config
+
+ for callback in _in_memory_loggers:
+ if isinstance(callback, OpenTelemetryV2):
+ return callback
+ otel_settings: Final = _get_custom_logger_settings_from_proxy_server(callback_name=logging_integration)
+ otel_logger_v2: Final = build_otel_v2_logger(
+ config=OpenTelemetryV2Config(**otel_settings), settings=otel_settings
+ )
+ _in_memory_loggers.append(otel_logger_v2)
+ _maybe_auto_initialize_arize_phoenix(_in_memory_loggers)
+ return otel_logger_v2
+
+ from litellm.integrations.opentelemetry import OpenTelemetry
+
+ for callback in _in_memory_loggers:
+ if type(callback) is OpenTelemetry:
+ return callback
+ otel_logger: Final = OpenTelemetry(
+ **_get_custom_logger_settings_from_proxy_server(callback_name=logging_integration)
+ )
+ _in_memory_loggers.append(otel_logger)
+
+ # Auto-initialize Arize Phoenix if Phoenix env vars are configured
+ # This allows users to get nested traces in both OTEL and Phoenix
+ # by only specifying "otel" in callbacks
+ _maybe_auto_initialize_arize_phoenix(_in_memory_loggers)
+
+ return otel_logger
+
+ elif logging_integration == "galileo":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, GalileoObserve):
+ return callback
+
+ galileo_logger: Final = GalileoObserve()
+ _in_memory_loggers.append(galileo_logger)
+ return galileo_logger
+ elif logging_integration == "cloudzero":
+ from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
+
+ for callback in _in_memory_loggers:
+ if isinstance(callback, CloudZeroLogger):
+ return callback
+ cloudzero_logger: Final = CloudZeroLogger()
+ _in_memory_loggers.append(cloudzero_logger)
+ return cloudzero_logger
+ elif logging_integration == "focus":
+ from litellm.integrations.focus.focus_logger import FocusLogger
+
+ for callback in _in_memory_loggers:
+ if type(callback) is FocusLogger: # exact match; exclude subclasses like VantageLogger
+ return callback
+ focus_logger: Final = FocusLogger()
+ _in_memory_loggers.append(focus_logger)
+ return focus_logger
+ elif logging_integration == "mavvrik":
+ from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
+ MavvrikFocusLogger,
+ )
+
+ for callback in _in_memory_loggers:
+ if type(callback) is MavvrikFocusLogger:
+ return callback
+ mavvrik_focus_logger: Final = MavvrikFocusLogger()
+ _in_memory_loggers.append(mavvrik_focus_logger)
+ return mavvrik_focus_logger
+ elif logging_integration == "vantage":
+ from litellm.integrations.vantage.vantage_logger import VantageLogger
+
+ for callback in _in_memory_loggers:
+ if isinstance(callback, VantageLogger):
+ return callback
+ vantage_logger: Final = VantageLogger()
+ _in_memory_loggers.append(vantage_logger)
+ return vantage_logger
+ elif logging_integration == "ternary":
+ from litellm.integrations.ternary.ternary_logger import TernaryLogger
+
+ for callback in _in_memory_loggers:
+ if isinstance(callback, TernaryLogger):
+ return callback
+ ternary_logger: Final = TernaryLogger()
+ _in_memory_loggers.append(ternary_logger)
+ return ternary_logger
+ elif logging_integration == "deepeval":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, DeepEvalLogger):
+ return callback
+ deepeval_logger: Final = DeepEvalLogger()
+ _in_memory_loggers.append(deepeval_logger)
+ return deepeval_logger
+
+ elif logging_integration == "logfire":
+ if "LOGFIRE_TOKEN" not in os.environ:
+ raise ValueError("LOGFIRE_TOKEN not found in environment variables")
+ from litellm.integrations.opentelemetry import (
+ OpenTelemetry,
+ OpenTelemetryConfig,
+ )
+
+ logfire_base_url: Final = os.getenv("LOGFIRE_BASE_URL", "https://logfire-api.pydantic.dev")
+ otel_config = OpenTelemetryConfig(
+ exporter="otlp_http",
+ endpoint=f"{logfire_base_url.rstrip('/')}/v1/traces",
+ headers=f"Authorization={os.getenv('LOGFIRE_TOKEN')}",
+ )
+ for callback in _in_memory_loggers:
+ # Use exact type check to avoid matching ArizePhoenixLogger (subclass)
+ if type(callback) is OpenTelemetry:
+ return callback
+ _otel_logger = OpenTelemetry(config=otel_config)
+ _in_memory_loggers.append(_otel_logger)
+ return _otel_logger
+ elif logging_integration == "dynamic_rate_limiter":
+ from litellm.proxy.hooks.dynamic_rate_limiter import (
+ PROXY_DynamicRateLimitHandler,
+ )
+
+ for callback in _in_memory_loggers:
+ if isinstance(callback, PROXY_DynamicRateLimitHandler):
+ return callback
+
+ if internal_usage_cache is None:
+ raise Exception(f"Internal Error: Cache cannot be empty - internal_usage_cache={internal_usage_cache}")
+
+ dynamic_rate_limiter_obj: Final = PROXY_DynamicRateLimitHandler(internal_usage_cache=internal_usage_cache)
+
+ if llm_router is not None and isinstance(llm_router, litellm.Router):
+ dynamic_rate_limiter_obj.update_variables(llm_router=llm_router)
+ _in_memory_loggers.append(dynamic_rate_limiter_obj)
+ return dynamic_rate_limiter_obj
+ elif logging_integration == "dynamic_rate_limiter_v3":
+ from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
+ PROXY_DynamicRateLimitHandlerV3,
+ )
+
+ for callback in _in_memory_loggers:
+ if isinstance(callback, PROXY_DynamicRateLimitHandlerV3):
+ return callback
+
+ if internal_usage_cache is None:
+ raise Exception(f"Internal Error: Cache cannot be empty - internal_usage_cache={internal_usage_cache}")
+
+ dynamic_rate_limiter_obj_v3: Final = PROXY_DynamicRateLimitHandlerV3(internal_usage_cache=internal_usage_cache)
+
+ if llm_router is not None and isinstance(llm_router, litellm.Router):
+ dynamic_rate_limiter_obj_v3.update_variables(llm_router=llm_router)
+ _in_memory_loggers.append(dynamic_rate_limiter_obj_v3)
+ return dynamic_rate_limiter_obj_v3
+ elif logging_integration == "langtrace":
+ if "LANGTRACE_API_KEY" not in os.environ:
+ raise ValueError("LANGTRACE_API_KEY not found in environment variables")
+ _v2 = _maybe_construct_otel_v2("langtrace", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2
+
+ from litellm.integrations.opentelemetry import (
+ OpenTelemetry,
+ OpenTelemetryConfig,
+ )
+
+ otel_config = OpenTelemetryConfig(
+ exporter="otlp_http",
+ endpoint=langtrace_trace_endpoint(os.getenv("LANGTRACE_API_HOST")),
+ headers=f"x-api-key={os.environ['LANGTRACE_API_KEY']}",
+ )
+ for callback in _in_memory_loggers:
+ if isinstance(callback, OpenTelemetry) and callback.callback_name == "langtrace":
+ return callback
+ _otel_logger = OpenTelemetry(config=otel_config, callback_name="langtrace")
+ _in_memory_loggers.append(_otel_logger)
+ return _otel_logger
+
+ elif logging_integration == "signoz":
+ from litellm.integrations.otel.presets.signoz import (
+ SIGNOZ_INGESTION_ENDPOINT_ENV,
+ )
+
+ _signoz_endpoint: Final = os.getenv(SIGNOZ_INGESTION_ENDPOINT_ENV)
+ if not _signoz_endpoint:
+ raise ValueError(f"{SIGNOZ_INGESTION_ENDPOINT_ENV} not found in environment variables")
+
+ _signoz_v2: Final = _maybe_construct_otel_v2("signoz", _in_memory_loggers)
+ if _signoz_v2 is not None:
+ return _signoz_v2
+
+ from litellm.integrations.opentelemetry import (
+ OpenTelemetry,
+ OpenTelemetryConfig,
+ )
+
+ _signoz_base: Final = _signoz_endpoint.rstrip("/")
+ _signoz_key: Final = os.getenv("SIGNOZ_INGESTION_KEY")
+ _signoz_config: Final = OpenTelemetryConfig(
+ exporter="otlp_http",
+ endpoint=(_signoz_base if _signoz_base.endswith("/v1/traces") else f"{_signoz_base}/v1/traces"),
+ headers=(f"signoz-ingestion-key={_signoz_key}" if _signoz_key else None),
+ )
+ for callback in _in_memory_loggers:
+ if isinstance(callback, OpenTelemetry) and callback.callback_name == "signoz":
+ return callback
+ _signoz_logger: Final = OpenTelemetry(config=_signoz_config, callback_name="signoz")
+ _in_memory_loggers.append(_signoz_logger)
+ return _signoz_logger
+
+ elif logging_integration == "mlflow":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, MlflowLogger):
+ return callback
+
+ _mlflow_logger: Final = MlflowLogger()
+ _in_memory_loggers.append(_mlflow_logger)
+ return _mlflow_logger
+ elif logging_integration == "langfuse":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, LangfusePromptManagement):
+ return callback
+
+ langfuse_logger: Final = LangfusePromptManagement()
+ _in_memory_loggers.append(langfuse_logger)
+ return langfuse_logger
+ elif logging_integration == "langfuse_otel":
+ _v2 = _maybe_construct_otel_v2("langfuse_otel", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2
+ from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger
+
+ for callback in _in_memory_loggers:
+ if isinstance(callback, LangfuseOtelLogger) and callback.callback_name == "langfuse_otel":
+ return callback
+ # Allow LangfuseOtelLogger to initialize its own config safely
+ # This prevents startup crashes if LANGFUSE keys are not in env (e.g. for dynamic usage)
+ _otel_logger = LangfuseOtelLogger(config=None, callback_name="langfuse_otel")
+ _in_memory_loggers.append(_otel_logger)
+ return _otel_logger
+ elif logging_integration == "weave_otel":
+ _v2 = _maybe_construct_otel_v2("weave_otel", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2
+ from litellm.integrations.opentelemetry import OpenTelemetryConfig
+ from litellm.integrations.weave.weave_otel import (
+ WeaveOtelLogger,
+ get_weave_otel_config,
+ )
+
+ weave_otel_config: Final = get_weave_otel_config()
+
+ otel_config = OpenTelemetryConfig(
+ exporter=weave_otel_config.protocol,
+ endpoint=weave_otel_config.endpoint,
+ headers=weave_otel_config.otlp_auth_headers,
+ )
+
+ for callback in _in_memory_loggers:
+ if isinstance(callback, WeaveOtelLogger) and callback.callback_name == "weave_otel":
+ return callback
+ _otel_logger = WeaveOtelLogger(config=otel_config, callback_name="weave_otel")
+ _in_memory_loggers.append(_otel_logger)
+ return _otel_logger
+ elif logging_integration == "pagerduty":
+ pagerduty_loggers: Final = _enterprise_alerting_loggers()
+ for callback in _in_memory_loggers:
+ if isinstance(callback, pagerduty_loggers.pagerduty):
+ return callback
+ pagerduty_logger: Final = pagerduty_loggers.pagerduty(**custom_logger_init_args_value)
+ _in_memory_loggers.append(pagerduty_logger)
+ return pagerduty_logger
+ elif logging_integration == "anthropic_cache_control_hook":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, AnthropicCacheControlHook):
+ return callback
+ anthropic_cache_control_hook: Final = AnthropicCacheControlHook()
+ _in_memory_loggers.append(anthropic_cache_control_hook)
+ return anthropic_cache_control_hook
+ elif logging_integration == "vector_store_pre_call_hook":
+ from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
+ VectorStorePreCallHook,
+ )
+
+ for callback in _in_memory_loggers:
+ if isinstance(callback, VectorStorePreCallHook):
+ return callback
+ vector_store_pre_call_hook: Final = VectorStorePreCallHook()
+ _in_memory_loggers.append(vector_store_pre_call_hook)
+ return vector_store_pre_call_hook
+ elif logging_integration == "gcs_pubsub":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, GcsPubSubLogger):
+ return callback
+ _gcs_pubsub_logger: Final = GcsPubSubLogger()
+ _in_memory_loggers.append(_gcs_pubsub_logger)
+ return _gcs_pubsub_logger
+ elif logging_integration == "generic_api":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, _GENERIC_API_LOGGER_CLS):
+ return callback
+ generic_api_logger: Final = GenericAPILogger()
+ _in_memory_loggers.append(generic_api_logger)
+ return generic_api_logger
+ elif logging_integration == "resend_email":
+ resend_email_loggers: Final = _enterprise_alerting_loggers()
+ for callback in _in_memory_loggers:
+ if isinstance(callback, resend_email_loggers.resend_email):
+ return callback
+ resend_email_logger: Final = resend_email_loggers.resend_email()
+ _in_memory_loggers.append(resend_email_logger)
+ return resend_email_logger
+ elif logging_integration == "sendgrid_email":
+ sendgrid_email_loggers: Final = _enterprise_alerting_loggers()
+ for callback in _in_memory_loggers:
+ if isinstance(callback, sendgrid_email_loggers.sendgrid_email):
+ return callback
+ sendgrid_email_logger: Final = sendgrid_email_loggers.sendgrid_email()
+ _in_memory_loggers.append(sendgrid_email_logger)
+ return sendgrid_email_logger
+ elif logging_integration == "smtp_email":
+ smtp_email_loggers: Final = _enterprise_alerting_loggers()
+ for callback in _in_memory_loggers:
+ if isinstance(callback, smtp_email_loggers.smtp_email):
+ return callback
+ smtp_email_logger: Final = smtp_email_loggers.smtp_email()
+ _in_memory_loggers.append(smtp_email_logger)
+ return smtp_email_logger
+ elif logging_integration == "humanloop":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, HumanloopLogger):
+ return callback
+
+ humanloop_logger: Final = HumanloopLogger()
+ _in_memory_loggers.append(humanloop_logger)
+ return humanloop_logger
+ elif logging_integration == "dotprompt":
+ for callback in _in_memory_loggers:
+ if isinstance(callback, DotpromptManager):
+ return callback
+
+ dotprompt_logger: Final = DotpromptManager()
+ _in_memory_loggers.append(dotprompt_logger)
+ return dotprompt_logger
+ elif logging_integration == "bitbucket":
+ from litellm.integrations.bitbucket.bitbucket_prompt_manager import (
+ BitBucketPromptManager,
+ )
+
+ for callback in _in_memory_loggers:
+ if isinstance(callback, BitBucketPromptManager):
+ return callback
+
+ # Get global BitBucket config
+ bitbucket_config: Final = getattr(litellm, "global_bitbucket_config", None)
+ if bitbucket_config is None:
+ raise ValueError("BitBucket configuration not found. Please set litellm.global_bitbucket_config first.")
+
+ bitbucket_logger: Final = BitBucketPromptManager(bitbucket_config=bitbucket_config)
+ _in_memory_loggers.append(bitbucket_logger)
+ return bitbucket_logger
+ elif logging_integration == "gitlab":
+ from litellm.integrations.gitlab.gitlab_prompt_manager import (
+ GitLabPromptManager,
+ )
+
+ for callback in _in_memory_loggers:
+ if isinstance(callback, GitLabPromptManager):
+ return callback
+
+ # Get global BitBucket config
+ gitlab_config: Final = getattr(litellm, "global_gitlab_config", None)
+ if gitlab_config is None:
+ raise ValueError("Gitlab configuration not found. Please set litellm.global_gitlab_config first.")
+
+ gitlab_logger: Final = GitLabPromptManager(gitlab_config=gitlab_config)
+ _in_memory_loggers.append(gitlab_logger)
+ return gitlab_logger
+ elif logging_integration == "newrelic":
+ if custom_logger_init_args_value.get("newrelic_api_key"):
+ # Team-scoped credentials: per-team METRICS logger, isolated per
+ # credential set via DynamicLoggingCache. The trace logger for
+ # this name stays on the global path below.
+ from litellm.integrations.newrelic.newrelic_team_handler import (
+ NewRelicHandler,
+ )
+
+ return NewRelicHandler.get_newrelic_logger_for_request(
+ standard_callback_dynamic_params=cast( # cast-ok: callback options come from dynamic config
+ StandardCallbackDynamicParams, custom_logger_init_args_value
+ ),
+ in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
+ )
+
+ _v2 = _maybe_construct_otel_v2("newrelic", _in_memory_loggers)
+ if _v2 is not None:
+ return _v2
+ for callback in _in_memory_loggers:
+ if isinstance(callback, NewRelicLogger):
+ return callback
+ newrelic_logger: Final = NewRelicLogger()
+ _in_memory_loggers.append(newrelic_logger)
+ return newrelic_logger
+ return None
+
+
def _init_custom_logger_compatible_class(
- logging_integration: _custom_logger_compatible_callbacks_literal,
+ logging_integration: str,
internal_usage_cache: DualCache | None,
llm_router: object,
custom_logger_init_args: dict[str, object] | None = {},
+ raise_import_errors: bool = False,
) -> CustomLogger | None:
"""
Initialize a custom logger compatible class
"""
try:
- custom_logger_init_args_value: Final[dict[str, object]] = custom_logger_init_args or {}
- if logging_integration == "agentops": # Add AgentOps initialization
- _v2 = _maybe_construct_otel_v2("agentops", _in_memory_loggers)
- if _v2 is not None:
- return _v2
- for callback in _in_memory_loggers:
- if isinstance(callback, AgentOps):
- return callback
-
- agentops_logger: Final = AgentOps()
- _in_memory_loggers.append(agentops_logger)
- return agentops_logger
- elif logging_integration == "lago":
- for callback in _in_memory_loggers:
- if isinstance(callback, LagoLogger):
- return callback
-
- lago_logger: Final = LagoLogger()
- _in_memory_loggers.append(lago_logger)
- return lago_logger
- elif logging_integration == "openmeter":
- for callback in _in_memory_loggers:
- if isinstance(callback, OpenMeterLogger):
- return callback
-
- _openmeter_logger: Final = OpenMeterLogger()
- _in_memory_loggers.append(_openmeter_logger)
- return _openmeter_logger
- elif logging_integration == "posthog":
- for callback in _in_memory_loggers:
- if isinstance(callback, PostHogLogger):
- return callback
-
- _posthog_logger: Final = PostHogLogger()
- _in_memory_loggers.append(_posthog_logger)
- return _posthog_logger
- elif logging_integration == "braintrust":
- from litellm.integrations.braintrust_logging import BraintrustLogger
-
- for callback in _in_memory_loggers:
- if isinstance(callback, BraintrustLogger):
- return callback
-
- braintrust_logger: Final = BraintrustLogger()
- _in_memory_loggers.append(braintrust_logger)
- return braintrust_logger
- elif logging_integration == "langsmith":
- for callback in _in_memory_loggers:
- if isinstance(callback, LangsmithLogger):
- return callback
-
- _langsmith_logger: Final = LangsmithLogger()
- _in_memory_loggers.append(_langsmith_logger)
- return _langsmith_logger
- elif logging_integration == "argilla":
- for callback in _in_memory_loggers:
- if isinstance(callback, ArgillaLogger):
- return callback
-
- _argilla_logger: Final = ArgillaLogger()
- _in_memory_loggers.append(_argilla_logger)
- return _argilla_logger
- elif logging_integration == "literalai":
- for callback in _in_memory_loggers:
- if isinstance(callback, LiteralAILogger):
- return callback
-
- _literalai_logger: Final = LiteralAILogger()
- _in_memory_loggers.append(_literalai_logger)
- return _literalai_logger
- elif logging_integration == "litellm_agent":
- for callback in _in_memory_loggers:
- if isinstance(callback, LiteLLMAgentModelResolver):
- return callback
-
- _litellm_agent_resolver: Final = LiteLLMAgentModelResolver()
- _in_memory_loggers.append(_litellm_agent_resolver)
- return _litellm_agent_resolver
- elif logging_integration == "prometheus":
- PrometheusLogger: Final = _get_cached_prometheus_logger()
-
- for callback in _in_memory_loggers:
- if isinstance(callback, PrometheusLogger):
- return callback
-
- _prometheus_logger: Final = PrometheusLogger()
- _in_memory_loggers.append(_prometheus_logger)
- return _prometheus_logger
- elif logging_integration == "datadog":
- # Check if team-scoped credentials are provided
- _dd_api_key: Final = custom_logger_init_args_value.get("dd_api_key")
- _dd_site: Final = custom_logger_init_args_value.get("dd_site")
- _dd_agent_host: Final = custom_logger_init_args_value.get("dd_agent_host")
- _dd_agent_port: Final = custom_logger_init_args_value.get("dd_agent_port")
-
- if _dd_api_key or _dd_site or _dd_agent_host:
- # Team-scoped credentials: use DynamicLoggingCache for per-credential isolation
- from litellm.integrations.datadog.datadog_team_handler import (
- DataDogHandler,
- )
-
- return DataDogHandler.get_datadog_logger_for_request(
- standard_callback_dynamic_params=cast( # cast-ok: callback options come from dynamic config
- StandardCallbackDynamicParams, custom_logger_init_args_value
- ),
- in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
- )
-
- # Global (env-var based): reuse cached instance
- for callback in _in_memory_loggers:
- if isinstance(callback, DataDogLogger):
- return callback
-
- _datadog_logger: Final = DataDogLogger()
- _in_memory_loggers.append(_datadog_logger)
- return _datadog_logger
- elif logging_integration == "datadog_metrics":
- for callback in _in_memory_loggers:
- if isinstance(callback, DatadogMetricsLogger):
- return callback
-
- _datadog_metrics_logger: Final = DatadogMetricsLogger()
- _in_memory_loggers.append(_datadog_metrics_logger)
- return _datadog_metrics_logger
- elif logging_integration == "datadog_llm_observability":
- _datadog_llm_obs_logger: Final = DataDogLLMObsLogger()
- _in_memory_loggers.append(_datadog_llm_obs_logger)
- return _datadog_llm_obs_logger
- elif logging_integration == "azure_sentinel":
- for callback in _in_memory_loggers:
- if isinstance(callback, AzureSentinelLogger):
- return callback
-
- _azure_sentinel_logger: Final = AzureSentinelLogger()
- _in_memory_loggers.append(_azure_sentinel_logger)
- return _azure_sentinel_logger
- elif logging_integration == "gcs_bucket":
- for callback in _in_memory_loggers:
- if isinstance(callback, GCSBucketLogger):
- return callback
-
- _gcs_bucket_logger: Final = GCSBucketLogger()
- _in_memory_loggers.append(_gcs_bucket_logger)
- return _gcs_bucket_logger
- elif logging_integration == "s3_v2":
- for callback in _in_memory_loggers:
- if isinstance(callback, S3V2Logger):
- return callback
-
- _s3_v2_logger: Final = S3V2Logger()
- _in_memory_loggers.append(_s3_v2_logger)
- return _s3_v2_logger
- elif logging_integration == "clickhouse":
- for callback in _in_memory_loggers:
- if isinstance(callback, ClickHouseSpendLogger):
- return callback
-
- _clickhouse_spend_logger: Final = ClickHouseSpendLogger()
- _in_memory_loggers.append(_clickhouse_spend_logger)
- return _clickhouse_spend_logger
- elif logging_integration == "pointfive":
- for callback in _in_memory_loggers:
- if isinstance(callback, PointFiveLogger):
- return callback
-
- _pointfive_logger: Final = PointFiveLogger()
- _in_memory_loggers.append(_pointfive_logger)
- return _pointfive_logger
- elif logging_integration == "zerobus":
- for callback in _in_memory_loggers:
- if isinstance(callback, ZerobusLogger):
- return callback
-
- _zerobus_logger: Final = ZerobusLogger()
- _in_memory_loggers.append(_zerobus_logger)
- return _zerobus_logger
- elif logging_integration == "aws_sqs":
- for callback in _in_memory_loggers:
- if isinstance(callback, SQSLogger):
- return callback
-
- _aws_sqs_logger: Final = SQSLogger()
- _in_memory_loggers.append(_aws_sqs_logger)
- return _aws_sqs_logger
- elif logging_integration == "azure_storage":
- for callback in _in_memory_loggers:
- if isinstance(callback, AzureBlobStorageLogger):
- return callback
-
- _azure_storage_logger: Final = AzureBlobStorageLogger()
- _in_memory_loggers.append(_azure_storage_logger)
- return _azure_storage_logger
- elif logging_integration == "opik":
- for callback in _in_memory_loggers:
- if isinstance(callback, OpikLogger):
- return callback
-
- _opik_logger: Final = OpikLogger()
- _in_memory_loggers.append(_opik_logger)
- return _opik_logger
- elif logging_integration == "arize":
- _v2 = _maybe_construct_otel_v2("arize", _in_memory_loggers)
- if _v2 is not None:
- return _v2
- from litellm.integrations.opentelemetry import (
- OpenTelemetry,
- OpenTelemetryConfig,
- )
-
- arize_config: Final = ArizeLogger.get_arize_config()
- if arize_config.endpoint is None:
- raise ValueError(
- "No valid endpoint found for Arize, please set 'ARIZE_ENDPOINT' to your GRPC endpoint or 'ARIZE_HTTP_ENDPOINT' to your HTTP endpoint"
- )
- otel_config = OpenTelemetryConfig(
- exporter=arize_config.protocol,
- endpoint=arize_config.endpoint,
- service_name=arize_config.project_name,
- )
-
- os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = (
- f"space_id={arize_config.space_key or arize_config.space_id},api_key={arize_config.api_key}"
- )
- for callback in _in_memory_loggers:
- if isinstance(callback, ArizeLogger) and callback.callback_name == "arize":
- return callback
- _arize_otel_logger: Final = ArizeLogger(config=otel_config, callback_name="arize")
- _in_memory_loggers.append(_arize_otel_logger)
- return _arize_otel_logger
- elif logging_integration == "arize_phoenix":
- _v2 = _maybe_construct_otel_v2("arize_phoenix", _in_memory_loggers)
- if _v2 is not None:
- return _v2
- from litellm.integrations.opentelemetry import (
- OpenTelemetry,
- OpenTelemetryConfig,
- )
-
- arize_phoenix_config: Final = ArizePhoenixLogger.get_arize_phoenix_config()
- otel_config = OpenTelemetryConfig(
- exporter=arize_phoenix_config.protocol,
- endpoint=arize_phoenix_config.endpoint,
- headers=arize_phoenix_config.otlp_auth_headers,
- )
-
- # auth can be disabled on local deployments of arize phoenix
- if arize_phoenix_config.otlp_auth_headers is not None:
- os.environ["OTEL_EXPORTER_OTLP_TRACES_HEADERS"] = arize_phoenix_config.otlp_auth_headers
-
- for callback in _in_memory_loggers:
- if isinstance(callback, ArizePhoenixLogger) and callback.callback_name == "arize_phoenix":
- return callback
- _arize_phoenix_otel_logger: Final = ArizePhoenixLogger(config=otel_config, callback_name="arize_phoenix")
- _in_memory_loggers.append(_arize_phoenix_otel_logger)
- return _arize_phoenix_otel_logger
- elif logging_integration == "levo":
- _v2 = _maybe_construct_otel_v2("levo", _in_memory_loggers)
- if _v2 is not None:
- return _v2
- from litellm.integrations.levo.levo import LevoLogger
- from litellm.integrations.opentelemetry import (
- OpenTelemetry,
- OpenTelemetryConfig,
- )
-
- levo_config: Final = LevoLogger.get_levo_config()
- otel_config = OpenTelemetryConfig(
- exporter=levo_config.protocol,
- endpoint=levo_config.endpoint,
- headers=levo_config.otlp_auth_headers,
- )
-
- # Check if LevoLogger instance already exists
- for callback in _in_memory_loggers:
- if isinstance(callback, LevoLogger) and callback.callback_name == "levo":
- return callback
-
- _levo_otel_logger: Final = LevoLogger(config=otel_config, callback_name="levo")
- _in_memory_loggers.append(_levo_otel_logger)
- return _levo_otel_logger
- elif logging_integration == "otel":
- # Gate the new typed V2 adapter behind LITELLM_OTEL_V2. When off,
- # the legacy 3,227-line god-class is used unchanged. The two are
- # never registered simultaneously — the dedup loop below treats
- # any module under ``litellm.integrations.otel`` or
- # ``litellm.integrations.opentelemetry`` as "the OTel callback".
- from litellm.integrations.otel.model.config import is_otel_v2_enabled
-
- if is_otel_v2_enabled():
- from litellm.integrations.otel.logger import OpenTelemetryV2, build_otel_v2_logger
- from litellm.integrations.otel.model.config import OpenTelemetryV2Config
-
- for callback in _in_memory_loggers:
- if isinstance(callback, OpenTelemetryV2):
- return callback
- otel_settings: Final = _get_custom_logger_settings_from_proxy_server(callback_name=logging_integration)
- otel_logger_v2: Final = build_otel_v2_logger(
- config=OpenTelemetryV2Config(**otel_settings), settings=otel_settings
- )
- _in_memory_loggers.append(otel_logger_v2)
- _maybe_auto_initialize_arize_phoenix(_in_memory_loggers)
- return otel_logger_v2
-
- from litellm.integrations.opentelemetry import OpenTelemetry
-
- for callback in _in_memory_loggers:
- if type(callback) is OpenTelemetry:
- return callback
- otel_logger: Final = OpenTelemetry(
- **_get_custom_logger_settings_from_proxy_server(callback_name=logging_integration)
- )
- _in_memory_loggers.append(otel_logger)
-
- # Auto-initialize Arize Phoenix if Phoenix env vars are configured
- # This allows users to get nested traces in both OTEL and Phoenix
- # by only specifying "otel" in callbacks
- _maybe_auto_initialize_arize_phoenix(_in_memory_loggers)
-
- return otel_logger
-
- elif logging_integration == "galileo":
- for callback in _in_memory_loggers:
- if isinstance(callback, GalileoObserve):
- return callback
-
- galileo_logger: Final = GalileoObserve()
- _in_memory_loggers.append(galileo_logger)
- return galileo_logger
- elif logging_integration == "cloudzero":
- from litellm.integrations.cloudzero.cloudzero import CloudZeroLogger
-
- for callback in _in_memory_loggers:
- if isinstance(callback, CloudZeroLogger):
- return callback
- cloudzero_logger: Final = CloudZeroLogger()
- _in_memory_loggers.append(cloudzero_logger)
- return cloudzero_logger
- elif logging_integration == "focus":
- from litellm.integrations.focus.focus_logger import FocusLogger
-
- for callback in _in_memory_loggers:
- if type(callback) is FocusLogger: # exact match; exclude subclasses like VantageLogger
- return callback
- focus_logger: Final = FocusLogger()
- _in_memory_loggers.append(focus_logger)
- return focus_logger
- elif logging_integration == "mavvrik":
- from litellm.integrations.mavvrik_focus.mavvrik_focus_logger import (
- MavvrikFocusLogger,
- )
-
- for callback in _in_memory_loggers:
- if type(callback) is MavvrikFocusLogger:
- return callback
- mavvrik_focus_logger: Final = MavvrikFocusLogger()
- _in_memory_loggers.append(mavvrik_focus_logger)
- return mavvrik_focus_logger
- elif logging_integration == "vantage":
- from litellm.integrations.vantage.vantage_logger import VantageLogger
-
- for callback in _in_memory_loggers:
- if isinstance(callback, VantageLogger):
- return callback
- vantage_logger: Final = VantageLogger()
- _in_memory_loggers.append(vantage_logger)
- return vantage_logger
- elif logging_integration == "ternary":
- from litellm.integrations.ternary.ternary_logger import TernaryLogger
-
- for callback in _in_memory_loggers:
- if isinstance(callback, TernaryLogger):
- return callback
- ternary_logger: Final = TernaryLogger()
- _in_memory_loggers.append(ternary_logger)
- return ternary_logger
- elif logging_integration == "deepeval":
- for callback in _in_memory_loggers:
- if isinstance(callback, DeepEvalLogger):
- return callback
- deepeval_logger: Final = DeepEvalLogger()
- _in_memory_loggers.append(deepeval_logger)
- return deepeval_logger
-
- elif logging_integration == "logfire":
- if "LOGFIRE_TOKEN" not in os.environ:
- raise ValueError("LOGFIRE_TOKEN not found in environment variables")
- from litellm.integrations.opentelemetry import (
- OpenTelemetry,
- OpenTelemetryConfig,
- )
-
- logfire_base_url: Final = os.getenv("LOGFIRE_BASE_URL", "https://logfire-api.pydantic.dev")
- otel_config = OpenTelemetryConfig(
- exporter="otlp_http",
- endpoint=f"{logfire_base_url.rstrip('/')}/v1/traces",
- headers=f"Authorization={os.getenv('LOGFIRE_TOKEN')}",
- )
- for callback in _in_memory_loggers:
- # Use exact type check to avoid matching ArizePhoenixLogger (subclass)
- if type(callback) is OpenTelemetry:
- return callback
- _otel_logger = OpenTelemetry(config=otel_config)
- _in_memory_loggers.append(_otel_logger)
- return _otel_logger
- elif logging_integration == "dynamic_rate_limiter":
- from litellm.proxy.hooks.dynamic_rate_limiter import (
- PROXY_DynamicRateLimitHandler,
- )
-
- for callback in _in_memory_loggers:
- if isinstance(callback, PROXY_DynamicRateLimitHandler):
- return callback
-
- if internal_usage_cache is None:
- raise Exception(f"Internal Error: Cache cannot be empty - internal_usage_cache={internal_usage_cache}")
-
- dynamic_rate_limiter_obj: Final = PROXY_DynamicRateLimitHandler(internal_usage_cache=internal_usage_cache)
-
- if llm_router is not None and isinstance(llm_router, litellm.Router):
- dynamic_rate_limiter_obj.update_variables(llm_router=llm_router)
- _in_memory_loggers.append(dynamic_rate_limiter_obj)
- return dynamic_rate_limiter_obj
- elif logging_integration == "dynamic_rate_limiter_v3":
- from litellm.proxy.hooks.dynamic_rate_limiter_v3 import (
- PROXY_DynamicRateLimitHandlerV3,
- )
-
- for callback in _in_memory_loggers:
- if isinstance(callback, PROXY_DynamicRateLimitHandlerV3):
- return callback
-
- if internal_usage_cache is None:
- raise Exception(f"Internal Error: Cache cannot be empty - internal_usage_cache={internal_usage_cache}")
-
- dynamic_rate_limiter_obj_v3: Final = PROXY_DynamicRateLimitHandlerV3(
- internal_usage_cache=internal_usage_cache
- )
-
- if llm_router is not None and isinstance(llm_router, litellm.Router):
- dynamic_rate_limiter_obj_v3.update_variables(llm_router=llm_router)
- _in_memory_loggers.append(dynamic_rate_limiter_obj_v3)
- return dynamic_rate_limiter_obj_v3
- elif logging_integration == "langtrace":
- if "LANGTRACE_API_KEY" not in os.environ:
- raise ValueError("LANGTRACE_API_KEY not found in environment variables")
- _v2 = _maybe_construct_otel_v2("langtrace", _in_memory_loggers)
- if _v2 is not None:
- return _v2
-
- from litellm.integrations.opentelemetry import (
- OpenTelemetry,
- OpenTelemetryConfig,
- )
-
- otel_config = OpenTelemetryConfig(
- exporter="otlp_http",
- endpoint=langtrace_trace_endpoint(os.getenv("LANGTRACE_API_HOST")),
- headers=f"x-api-key={os.environ['LANGTRACE_API_KEY']}",
- )
- for callback in _in_memory_loggers:
- if isinstance(callback, OpenTelemetry) and callback.callback_name == "langtrace":
- return callback
- _otel_logger = OpenTelemetry(config=otel_config, callback_name="langtrace")
- _in_memory_loggers.append(_otel_logger)
- return _otel_logger
-
- elif logging_integration == "signoz":
- from litellm.integrations.otel.presets.signoz import (
- SIGNOZ_INGESTION_ENDPOINT_ENV,
- )
-
- _signoz_endpoint: Final = os.getenv(SIGNOZ_INGESTION_ENDPOINT_ENV)
- if not _signoz_endpoint:
- raise ValueError(f"{SIGNOZ_INGESTION_ENDPOINT_ENV} not found in environment variables")
-
- _signoz_v2: Final = _maybe_construct_otel_v2("signoz", _in_memory_loggers)
- if _signoz_v2 is not None:
- return _signoz_v2
-
- from litellm.integrations.opentelemetry import (
- OpenTelemetry,
- OpenTelemetryConfig,
- )
-
- _signoz_base: Final = _signoz_endpoint.rstrip("/")
- _signoz_key: Final = os.getenv("SIGNOZ_INGESTION_KEY")
- _signoz_config: Final = OpenTelemetryConfig(
- exporter="otlp_http",
- endpoint=(_signoz_base if _signoz_base.endswith("/v1/traces") else f"{_signoz_base}/v1/traces"),
- headers=(f"signoz-ingestion-key={_signoz_key}" if _signoz_key else None),
- )
- for callback in _in_memory_loggers:
- if isinstance(callback, OpenTelemetry) and callback.callback_name == "signoz":
- return callback
- _signoz_logger: Final = OpenTelemetry(config=_signoz_config, callback_name="signoz")
- _in_memory_loggers.append(_signoz_logger)
- return _signoz_logger
-
- elif logging_integration == "mlflow":
- for callback in _in_memory_loggers:
- if isinstance(callback, MlflowLogger):
- return callback
-
- _mlflow_logger: Final = MlflowLogger()
- _in_memory_loggers.append(_mlflow_logger)
- return _mlflow_logger
- elif logging_integration == "langfuse":
- for callback in _in_memory_loggers:
- if isinstance(callback, LangfusePromptManagement):
- return callback
-
- langfuse_logger: Final = LangfusePromptManagement()
- _in_memory_loggers.append(langfuse_logger)
- return langfuse_logger
- elif logging_integration == "langfuse_otel":
- _v2 = _maybe_construct_otel_v2("langfuse_otel", _in_memory_loggers)
- if _v2 is not None:
- return _v2
- from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger
-
- for callback in _in_memory_loggers:
- if isinstance(callback, LangfuseOtelLogger) and callback.callback_name == "langfuse_otel":
- return callback
- # Allow LangfuseOtelLogger to initialize its own config safely
- # This prevents startup crashes if LANGFUSE keys are not in env (e.g. for dynamic usage)
- _otel_logger = LangfuseOtelLogger(config=None, callback_name="langfuse_otel")
- _in_memory_loggers.append(_otel_logger)
- return _otel_logger
- elif logging_integration == "weave_otel":
- _v2 = _maybe_construct_otel_v2("weave_otel", _in_memory_loggers)
- if _v2 is not None:
- return _v2
- from litellm.integrations.opentelemetry import OpenTelemetryConfig
- from litellm.integrations.weave.weave_otel import (
- WeaveOtelLogger,
- get_weave_otel_config,
- )
-
- weave_otel_config: Final = get_weave_otel_config()
-
- otel_config = OpenTelemetryConfig(
- exporter=weave_otel_config.protocol,
- endpoint=weave_otel_config.endpoint,
- headers=weave_otel_config.otlp_auth_headers,
- )
-
- for callback in _in_memory_loggers:
- if isinstance(callback, WeaveOtelLogger) and callback.callback_name == "weave_otel":
- return callback
- _otel_logger = WeaveOtelLogger(config=otel_config, callback_name="weave_otel")
- _in_memory_loggers.append(_otel_logger)
- return _otel_logger
- elif logging_integration == "pagerduty":
- pagerduty_loggers: Final = _enterprise_alerting_loggers()
- for callback in _in_memory_loggers:
- if isinstance(callback, pagerduty_loggers.pagerduty):
- return callback
- pagerduty_logger: Final = pagerduty_loggers.pagerduty(**custom_logger_init_args_value)
- _in_memory_loggers.append(pagerduty_logger)
- return pagerduty_logger
- elif logging_integration == "anthropic_cache_control_hook":
- for callback in _in_memory_loggers:
- if isinstance(callback, AnthropicCacheControlHook):
- return callback
- anthropic_cache_control_hook: Final = AnthropicCacheControlHook()
- _in_memory_loggers.append(anthropic_cache_control_hook)
- return anthropic_cache_control_hook
- elif logging_integration == "vector_store_pre_call_hook":
- from litellm.integrations.vector_store_integrations.vector_store_pre_call_hook import (
- VectorStorePreCallHook,
- )
-
- for callback in _in_memory_loggers:
- if isinstance(callback, VectorStorePreCallHook):
- return callback
- vector_store_pre_call_hook: Final = VectorStorePreCallHook()
- _in_memory_loggers.append(vector_store_pre_call_hook)
- return vector_store_pre_call_hook
- elif logging_integration == "gcs_pubsub":
- for callback in _in_memory_loggers:
- if isinstance(callback, GcsPubSubLogger):
- return callback
- _gcs_pubsub_logger: Final = GcsPubSubLogger()
- _in_memory_loggers.append(_gcs_pubsub_logger)
- return _gcs_pubsub_logger
- elif logging_integration == "generic_api":
- for callback in _in_memory_loggers:
- if isinstance(callback, _GENERIC_API_LOGGER_CLS):
- return callback
- generic_api_logger: Final = GenericAPILogger()
- _in_memory_loggers.append(generic_api_logger)
- return generic_api_logger
- elif logging_integration == "resend_email":
- resend_email_loggers: Final = _enterprise_alerting_loggers()
- for callback in _in_memory_loggers:
- if isinstance(callback, resend_email_loggers.resend_email):
- return callback
- resend_email_logger: Final = resend_email_loggers.resend_email()
- _in_memory_loggers.append(resend_email_logger)
- return resend_email_logger
- elif logging_integration == "sendgrid_email":
- sendgrid_email_loggers: Final = _enterprise_alerting_loggers()
- for callback in _in_memory_loggers:
- if isinstance(callback, sendgrid_email_loggers.sendgrid_email):
- return callback
- sendgrid_email_logger: Final = sendgrid_email_loggers.sendgrid_email()
- _in_memory_loggers.append(sendgrid_email_logger)
- return sendgrid_email_logger
- elif logging_integration == "smtp_email":
- smtp_email_loggers: Final = _enterprise_alerting_loggers()
- for callback in _in_memory_loggers:
- if isinstance(callback, smtp_email_loggers.smtp_email):
- return callback
- smtp_email_logger: Final = smtp_email_loggers.smtp_email()
- _in_memory_loggers.append(smtp_email_logger)
- return smtp_email_logger
- elif logging_integration == "humanloop":
- for callback in _in_memory_loggers:
- if isinstance(callback, HumanloopLogger):
- return callback
-
- humanloop_logger: Final = HumanloopLogger()
- _in_memory_loggers.append(humanloop_logger)
- return humanloop_logger
- elif logging_integration == "dotprompt":
- for callback in _in_memory_loggers:
- if isinstance(callback, DotpromptManager):
- return callback
-
- dotprompt_logger: Final = DotpromptManager()
- _in_memory_loggers.append(dotprompt_logger)
- return dotprompt_logger
- elif logging_integration == "bitbucket":
- from litellm.integrations.bitbucket.bitbucket_prompt_manager import (
- BitBucketPromptManager,
- )
-
- for callback in _in_memory_loggers:
- if isinstance(callback, BitBucketPromptManager):
- return callback
-
- # Get global BitBucket config
- bitbucket_config: Final = getattr(litellm, "global_bitbucket_config", None)
- if bitbucket_config is None:
- raise ValueError("BitBucket configuration not found. Please set litellm.global_bitbucket_config first.")
-
- bitbucket_logger: Final = BitBucketPromptManager(bitbucket_config=bitbucket_config)
- _in_memory_loggers.append(bitbucket_logger)
- return bitbucket_logger
- elif logging_integration == "gitlab":
- from litellm.integrations.gitlab.gitlab_prompt_manager import (
- GitLabPromptManager,
- )
-
- for callback in _in_memory_loggers:
- if isinstance(callback, GitLabPromptManager):
- return callback
-
- # Get global BitBucket config
- gitlab_config: Final = getattr(litellm, "global_gitlab_config", None)
- if gitlab_config is None:
- raise ValueError("Gitlab configuration not found. Please set litellm.global_gitlab_config first.")
-
- gitlab_logger: Final = GitLabPromptManager(gitlab_config=gitlab_config)
- _in_memory_loggers.append(gitlab_logger)
- return gitlab_logger
- elif logging_integration == "newrelic":
- if custom_logger_init_args_value.get("newrelic_api_key"):
- # Team-scoped credentials: per-team METRICS logger, isolated per
- # credential set via DynamicLoggingCache. The trace logger for
- # this name stays on the global path below.
- from litellm.integrations.newrelic.newrelic_team_handler import (
- NewRelicHandler,
- )
-
- return NewRelicHandler.get_newrelic_logger_for_request(
- standard_callback_dynamic_params=cast( # cast-ok: callback options come from dynamic config
- StandardCallbackDynamicParams, custom_logger_init_args_value
- ),
- in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
- )
-
- _v2 = _maybe_construct_otel_v2("newrelic", _in_memory_loggers)
- if _v2 is not None:
- return _v2
- for callback in _in_memory_loggers:
- if isinstance(callback, NewRelicLogger):
- return callback
- newrelic_logger: Final = NewRelicLogger()
- _in_memory_loggers.append(newrelic_logger)
- return newrelic_logger
- return None
+ return build_custom_logger_compatible_class(
+ logging_integration,
+ internal_usage_cache=internal_usage_cache,
+ llm_router=llm_router,
+ custom_logger_init_args=custom_logger_init_args,
+ )
except Exception as e:
+ if raise_import_errors and isinstance(e, ImportError):
+ raise ImportError(
+ f"Could not build the '{logging_integration}' callback configured in litellm_settings: {e}"
+ ) from e
verbose_logger.exception("[Non-Blocking Error] Error initializing custom logger: %s", e)
return None
- return None
+
+
+def init_custom_logger_compatible_class_at_startup(
+ callback: str,
+ internal_usage_cache: DualCache | None,
+ llm_router: object,
+) -> CustomLogger | None:
+ return _init_custom_logger_compatible_class(
+ callback, internal_usage_cache=internal_usage_cache, llm_router=llm_router, raise_import_errors=True
+ )
def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[CustomLogger]) -> "OpenTelemetryV2 | None":
diff --git a/litellm/litellm_core_utils/prompt_templates/image_handling.py b/litellm/litellm_core_utils/prompt_templates/image_handling.py
index 7924bb9bf4c..4cf76e5bd5e 100644
--- a/litellm/litellm_core_utils/prompt_templates/image_handling.py
+++ b/litellm/litellm_core_utils/prompt_templates/image_handling.py
@@ -4,12 +4,14 @@ Helper functions to handle images passed in messages
import asyncio
import base64
+import time
from collections.abc import Callable, Iterable, Mapping
from dataclasses import dataclass
from itertools import chain
from types import MappingProxyType
from typing import Final
+import httpx
from httpx import Response
import litellm
@@ -22,6 +24,7 @@ from litellm.types.llms.openai import AllMessageValues
MAX_IMGS_IN_MEMORY: Final = 10
MAX_CONCURRENT_REMOTE_MEDIA_FETCHES: Final = 20
+_FETCH_ATTEMPTS: Final = 3
in_memory_cache: Final = InMemoryCache(max_size_in_memory=MAX_IMGS_IN_MEMORY)
@@ -73,6 +76,13 @@ def _process_image_response(response: Response, url: str) -> str:
return result
+def _attempt_timeouts(timeout: float | httpx.Timeout | None) -> Iterable[float | httpx.Timeout | None]:
+ if timeout is None or isinstance(timeout, httpx.Timeout):
+ return (timeout,) * _FETCH_ATTEMPTS
+ deadline: Final = time.monotonic() + timeout
+ return (remaining for _ in range(_FETCH_ATTEMPTS) if (remaining := deadline - time.monotonic()) > 0)
+
+
def _rejected_image_fetch(url: str, verdict: SSRFError) -> "litellm.ImageFetchError":
verbose_logger.warning("Image fetch of %s rejected before any request went out: %s", url, verdict)
return litellm.ImageFetchError(
@@ -81,7 +91,7 @@ def _rejected_image_fetch(url: str, verdict: SSRFError) -> "litellm.ImageFetchEr
)
-async def async_convert_url_to_base64(url: str) -> str:
+async def async_convert_url_to_base64(url: str, timeout: float | httpx.Timeout | None = None) -> str:
if url.startswith("data:") and ";base64," in url:
return url
@@ -96,9 +106,9 @@ async def async_convert_url_to_base64(url: str) -> str:
return cached_result
client: Final = litellm.module_level_aclient
- for _ in range(3):
+ for attempt_timeout in _attempt_timeouts(timeout):
try:
- response = await async_safe_get(client, url)
+ response = await async_safe_get(client, url, timeout=attempt_timeout)
return _process_image_response(response, url)
except litellm.ImageFetchError:
raise
@@ -109,7 +119,7 @@ async def async_convert_url_to_base64(url: str) -> str:
raise litellm.ImageFetchError(f"Error: Unable to fetch image from URL after 3 attempts. url={url}")
-def convert_url_to_base64(url: str) -> str:
+def convert_url_to_base64(url: str, timeout: float | httpx.Timeout | None = None) -> str:
if url.startswith("data:") and ";base64," in url:
return url
@@ -124,9 +134,9 @@ def convert_url_to_base64(url: str) -> str:
return cached_result
client: Final = litellm.module_level_client
- for _ in range(3):
+ for attempt_timeout in _attempt_timeouts(timeout):
try:
- response = safe_get(client, url)
+ response = safe_get(client, url, timeout=attempt_timeout)
return _process_image_response(response, url)
except litellm.ImageFetchError:
raise
@@ -286,14 +296,16 @@ def _inline_message(
return inlined_message # pyright: ignore[reportReturnType] # the same message with its remote parts inlined
-async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore) -> str:
+async def _fetch_data_url(url: str, in_flight: asyncio.Semaphore, timeout: float | httpx.Timeout | None) -> str:
async with in_flight:
- return await async_convert_url_to_base64(url)
+ return await async_convert_url_to_base64(url, timeout=timeout)
-async def _fetch_data_urls(remote_urls: tuple[str, ...]) -> tuple[str, ...]:
+async def async_convert_urls_to_base64(
+ remote_urls: tuple[str, ...], timeout: float | httpx.Timeout | None = None
+) -> tuple[str, ...]:
in_flight: Final = asyncio.Semaphore(MAX_CONCURRENT_REMOTE_MEDIA_FETCHES)
- fetches: Final = tuple(asyncio.create_task(_fetch_data_url(url, in_flight)) for url in remote_urls)
+ fetches: Final = tuple(asyncio.create_task(_fetch_data_url(url, in_flight, timeout)) for url in remote_urls)
try:
return tuple(await asyncio.gather(*fetches))
except BaseException:
@@ -329,6 +341,6 @@ async def async_inline_remote_media(
remote_urls: Final = _remote_urls_to_inline(messages, should_inline)
if not remote_urls:
return messages
- data_urls: Final = await _fetch_data_urls(remote_urls)
+ data_urls: Final = await async_convert_urls_to_base64(remote_urls)
inlined: Final = MappingProxyType(dict(zip(remote_urls, data_urls, strict=True)))
return [_inline_message(message, inlined, should_inline) for message in messages]
diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py
index a78e6f17cc9..236286a493b 100644
--- a/litellm/litellm_core_utils/streaming_handler.py
+++ b/litellm/litellm_core_utils/streaming_handler.py
@@ -1,6 +1,7 @@
import asyncio
import collections.abc
import datetime
+import inspect
import json
import logging
import threading
@@ -234,6 +235,7 @@ class CustomStreamWrapper:
self.completion_stream = completion_stream
self.sent_first_chunk = False
self.sent_last_chunk = False
+ self._closed = False
self._stream_created_time: float = time.time()
litellm_params: Final[GenericLiteLLMParams] = GenericLiteLLMParams.model_validate(
@@ -328,6 +330,18 @@ class CustomStreamWrapper:
def __aiter__(self) -> AsyncIterator["ModelResponseStream"]:
return self
+ def __enter__(self) -> "CustomStreamWrapper":
+ return self
+
+ def __exit__(self, *exc_info: object) -> None:
+ self.close()
+
+ async def __aenter__(self) -> "CustomStreamWrapper":
+ return self
+
+ async def __aexit__(self, *exc_info: object) -> None:
+ await self.aclose()
+
def _restore_consumer_correlation_context(self, *, guarded: bool = False) -> None:
"""Restore trace_id/session_id in the *consuming* thread/task/context.
@@ -392,11 +406,26 @@ class CustomStreamWrapper:
"""
self._restore_consumer_correlation_context(guarded=True)
+ def close(self) -> None:
+ stream_to_close: Final = cast(object, self.completion_stream) # cast-ok: completion_stream is untyped
+ self.completion_stream = None
+ self._closed = True
+ close_stream: Final = getattr(stream_to_close, "close", None)
+ if callable(close_stream):
+ try:
+ result: Final = close_stream()
+ if inspect.iscoroutine(result):
+ result.close()
+ except Exception as e: # noqa: BLE001 # best-effort cleanup; must not raise into the caller
+ verbose_logger.debug("CustomStreamWrapper.close: error closing completion_stream: %s", e)
+ self._restore_consumer_correlation_context()
+
async def aclose(self):
# Restore the consumer's outer context only after the underlying
# provider stream's own close (and its diagnostic logging below, if
# closing fails) completes - not before - so those log lines still
# carry this closing stream's own trace_id/session_id.
+ self._closed = True
if self.completion_stream is not None:
stream_to_close: Final = self.completion_stream
self.completion_stream = None
@@ -1780,6 +1809,8 @@ class CustomStreamWrapper:
response.hidden_params["additional_headers"]["llm_provider-x-litellm-response-cost"] = _cost
def __next__(self) -> "ModelResponseStream":
+ if self._closed:
+ raise StopIteration
cache_hit = False
if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response":
cache_hit = True
@@ -1950,7 +1981,7 @@ class CustomStreamWrapper:
if self.sent_stream_usage is False and self.send_stream_usage is True:
self.sent_stream_usage = True
return response
- self._restore_consumer_correlation_context()
+ self.close()
raise # Re-raise StopIteration
else:
self.sent_last_chunk = True
@@ -1986,6 +2017,7 @@ class CustomStreamWrapper:
traceback_exception: Final = traceback.format_exc()
# LOG FAILURE - handle streaming failure logging in the _next_ object, remove `handle_failure` once it's deprecated
threading.Thread(target=self.logging_obj.failure_handler, args=(e, traceback_exception)).start()
+ self.close()
self._handle_stream_fallback_error(e)
def fetch_sync_stream(self):
@@ -2005,6 +2037,8 @@ class CustomStreamWrapper:
return self.completion_stream
async def __anext__(self) -> "ModelResponseStream":
+ if self._closed:
+ raise StopAsyncIteration
cache_hit = False
if self.custom_llm_provider is not None and self.custom_llm_provider == "cached_response":
cache_hit = True
@@ -2125,12 +2159,15 @@ class CustomStreamWrapper:
asyncio.create_task(
self.logging_obj.dispatch_failure_handlers(e, traceback_exception, prefer_async_handlers=True)
)
+ await self.aclose()
self._handle_stream_fallback_error(e)
except (httpx.ReadError, httpx.RemoteProtocolError) as e:
if self.received_finish_reason is None:
+ await self.aclose()
self._log_stream_failure_and_raise(e)
return await self._finalize_completed_stream(cache_hit=cache_hit)
except Exception as e:
+ await self.aclose()
self._log_stream_failure_and_raise(e)
async def _finalize_completed_stream(self, cache_hit: bool) -> "ModelResponseStream":
@@ -2224,7 +2261,7 @@ class CustomStreamWrapper:
)
)
- self._restore_consumer_correlation_context()
+ await self.aclose()
raise StopAsyncIteration # Re-raise StopIteration
else:
self.sent_last_chunk = True
diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py
index eb0de9d9b60..dff10473ae7 100644
--- a/litellm/llms/anthropic/chat/handler.py
+++ b/litellm/llms/anthropic/chat/handler.py
@@ -25,7 +25,10 @@ from litellm.litellm_core_utils.json_fragment_accumulator import JSONFragmentAcc
from litellm.llms.base_llm.base_utils import (
RESPONSE_FORMAT_ENVELOPE_NAME_KEYS,
ResponseFormatSchema,
+ names_a_response_format_property,
+ partial_response_format_answer,
response_format_envelope_contents,
+ split_response_format_answer,
)
from litellm.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
@@ -85,6 +88,9 @@ _NAME_VALUE_PREFIX: Final = re.compile(
r'^\s*\{\s*"(?:' + _ENVELOPE_NAME_KEY_PATTERN + r')"\s*:\s*(?:"(?:[^"\\]|\\.)*\\?)?$'
)
_STRING_VALUE_THEN_KEY: Final = re.compile(r'^"(?:[^"\\]|\\.)*"\s*,\s*"((?:[^"\\]|\\.)*)"\s*:')
+_LABEL_THEN_KEY: Final = re.compile(r'"((?:[^"\\]|\\.)*)"\s*,\s*"((?:[^"\\]|\\.)*)"\s*:\s*')
+_LEADING_LABELS: Final = re.compile(r'(?:"(?:[^"\\]|\\.)*"\s*,\s*"(?:[^"\\]|\\.)*"\s*:\s*)*')
+_CLOSED_STRING: Final = re.compile(r'"((?:[^"\\]|\\.)*)"')
_UNDECIDED_PREFIX_LIMIT: Final = 4096
@@ -95,14 +101,45 @@ def _key_after_a_string_verdict(string_value_start: str, schema: ResponseFormatS
return "pass_through" if next_key.group(1) in schema.property_names else "hold"
-def _first_value_verdict(after_first_key: str, schema: ResponseFormatSchema) -> _EnvelopeVerdict:
+def _is_a_label(text: str, schema: ResponseFormatSchema, tool_name: str) -> bool:
+ return text in schema.property_names or text == tool_name
+
+
+def _label_verdict(label: str, next_key: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
+ if _is_a_label(label, schema, tool_name):
+ return "hold"
+ return "pass_through" if next_key in schema.property_names else "undecided"
+
+
+def _leading_labels_verdict(string_value_start: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
+ last_value_start: Final = _LEADING_LABELS.sub("", string_value_start, count=1)
+ labelled: Final = string_value_start[: len(string_value_start) - len(last_value_start)]
+ verdicts: Final[tuple[_EnvelopeVerdict, ...]] = tuple(
+ _label_verdict(pair.group(1), pair.group(2), schema, tool_name) for pair in _LABEL_THEN_KEY.finditer(labelled)
+ )
+ decided: Final[_EnvelopeVerdict | None] = next((verdict for verdict in verdicts if verdict != "undecided"), None)
+ if decided is not None:
+ return decided
+ last_label: Final = _CLOSED_STRING.match(last_value_start)
+ if last_label is not None and _is_a_label(last_label.group(1), schema, tool_name):
+ return "hold"
+ return "undecided" if not last_value_start or last_value_start.startswith('"') else "pass_through"
+
+
+def _string_value_verdict(string_value_start: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
+ if schema.required_names:
+ return _key_after_a_string_verdict(string_value_start, schema)
+ return _leading_labels_verdict(string_value_start, schema, tool_name)
+
+
+def _first_value_verdict(after_first_key: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
value_start: Final = after_first_key.lstrip()
if not value_start:
return "undecided"
- if value_start.startswith("{"):
+ if value_start.startswith('"'):
+ return _string_value_verdict(value_start, schema, tool_name)
+ if value_start.startswith("{") and not schema.accepts_a_lone_undeclared_key:
return "hold"
- if value_start.startswith('"') and schema.required_names:
- return _key_after_a_string_verdict(value_start, schema)
return "pass_through"
@@ -116,9 +153,11 @@ def _prefix_verdict(buffered_json: str, schema: ResponseFormatSchema, tool_name:
return "hold"
if first_value is None and _NAME_VALUE_PREFIX.match(buffered_json):
return "undecided"
- if first_key.group(1) in schema.property_names or not schema.property_names or schema.accepts_a_lone_undeclared_key:
+ if first_key.group(1) == "values" and "values" not in schema.property_names:
+ return "hold"
+ if first_key.group(1) in schema.property_names or not schema.property_names:
return "pass_through"
- return _first_value_verdict(buffered_json[first_key.end() :], schema)
+ return _first_value_verdict(buffered_json[first_key.end() :], schema, tool_name)
def _held_answer(held: str, schema: ResponseFormatSchema) -> str:
@@ -130,6 +169,13 @@ def _held_answer(held: str, schema: ResponseFormatSchema) -> str:
return held if contents is None else json.dumps(contents)
+def _split_answer_text(blocks: tuple[str, ...], schema: ResponseFormatSchema) -> str:
+ answer: Final = split_response_format_answer(blocks, schema)
+ if answer is not None:
+ return json.dumps(answer)
+ return "".join(_held_answer(block, schema) for block in blocks)
+
+
def _envelope_verdict(buffered_json: str, schema: ResponseFormatSchema, tool_name: str) -> _EnvelopeVerdict:
verdict: Final = _prefix_verdict(buffered_json, schema, tool_name)
if verdict == "undecided" and len(buffered_json) > _UNDECIDED_PREFIX_LIMIT:
@@ -186,6 +232,7 @@ async def make_call(
speed=speed,
tool_name_reverse_map=tool_name_reverse_map,
response_format_schema=response_format_schema,
+ response=response,
)
# LOGGING
@@ -256,6 +303,7 @@ def make_sync_call(
speed=speed,
tool_name_reverse_map=tool_name_reverse_map,
response_format_schema=response_format_schema,
+ response=response,
)
# LOGGING
@@ -629,9 +677,12 @@ class ModelResponseIterator:
speed: str | None = None,
tool_name_reverse_map: dict[str, str] | None = None,
response_format_schema: ResponseFormatSchema | None = None,
+ response: httpx.Response | None = None,
):
self.streaming_response = streaming_response
self.response_iterator = self.streaming_response
+ self.sync_stream = sync_stream
+ self._response = response
self.content_blocks: list[ContentBlockDelta] = []
self.tool_index = -1
self.json_mode = json_mode
@@ -655,6 +706,9 @@ class ModelResponseIterator:
self.response_format_schema = response_format_schema
self._response_format_buffer: io.StringIO = io.StringIO()
self._response_format_verdict: _EnvelopeVerdict = "undecided"
+ self._split_answer_blocks: tuple[str, ...] = ()
+ self._withheld_answer_close: str = ""
+ self._streamed_answer_keys: frozenset[str] | None = None
# For handling partial JSON chunks from fragmentation
# See: https://github.com/BerriAI/litellm/issues/17473
@@ -917,7 +971,7 @@ class ModelResponseIterator:
# Track current content block type for filtering deltas
self.current_content_block_type = content_block_start["content_block"]["type"]
if content_block_start["content_block"]["type"] == "text":
- text = content_block_start["content_block"]["text"]
+ text = self._release_answer_close() + content_block_start["content_block"]["text"]
elif (
content_block_start["content_block"]["type"] == "tool_use"
or content_block_start["content_block"]["type"] == "server_tool_use"
@@ -1049,6 +1103,7 @@ class ModelResponseIterator:
pass
elif type_chunk == "message_delta":
finish_reason, usage, container = self._handle_message_delta(chunk)
+ text = self._release_answer_close() + self._drain_split_answer_blocks()
if container:
provider_specific_fields["container"] = container
# Store container_id and re-emit code_interpreter_results
@@ -1155,6 +1210,8 @@ class ModelResponseIterator:
# (per Anthropic's fine-grained streaming pattern)
tool_name: Final = tool_use.get("function", {}).get("name", "")
self.is_response_format_tool = tool_name == RESPONSE_FORMAT_TOOL_NAME
+ if self.is_response_format_tool:
+ self.tool_index -= 1
if self.is_response_format_tool and self.response_format_schema is not None:
self.converted_response_format_tool = True
@@ -1176,25 +1233,69 @@ class ModelResponseIterator:
) -> str:
if tool_use.get("id") is not None:
self._response_format_buffer = io.StringIO()
- self._response_format_verdict = "undecided"
+ held_back: Final = self._split_answer_blocks or self._streamed_answer_keys is not None
+ self._response_format_verdict = "hold" if held_back else "undecided"
return ""
fragment: Final = tool_use.get("function", {}).get("arguments") or ""
- if self._response_format_verdict == "pass_through":
- return fragment
self._response_format_buffer.write(fragment)
+ if self._response_format_verdict == "pass_through":
+ return self._withhold_answer_close(fragment)
if self._response_format_verdict == "hold":
return ""
buffered: Final = self._response_format_buffer.getvalue()
self._response_format_verdict = _envelope_verdict(buffered, schema, RESPONSE_FORMAT_TOOL_NAME)
- return buffered if self._response_format_verdict == "pass_through" else ""
+ return self._withhold_answer_close(buffered) if self._response_format_verdict == "pass_through" else ""
+
+ def _withhold_answer_close(self, text: str) -> str:
+ pending: Final = self._withheld_answer_close + text
+ body: Final = pending.rstrip()
+ cut: Final = len(body) - 1 if body.endswith("}") else len(pending)
+ self._withheld_answer_close = pending[cut:]
+ return pending[:cut]
+
+ def _release_answer_close(self) -> str:
+ close: Final = self._withheld_answer_close
+ self._withheld_answer_close = ""
+ self._streamed_answer_keys = None
+ return close
+
+ def _join_streamed_answer(self, held: str, schema: ResponseFormatSchema) -> str | None:
+ streamed: Final = self._streamed_answer_keys
+ fields: Final = partial_response_format_answer(held, schema)
+ if streamed is None or fields is None or not streamed.isdisjoint(fields):
+ return None
+ self._streamed_answer_keys = streamed | frozenset(fields)
+ separator: Final = ", " if streamed else ""
+ return f"{separator}{json.dumps(fields)[1:-1]}" if fields else ""
def _flush_response_format_buffer(self) -> str:
- if self.response_format_schema is None or self._response_format_verdict == "pass_through":
+ schema: Final = self.response_format_schema
+ if schema is None:
return ""
held: Final = self._response_format_buffer.getvalue()
+ verdict: Final = self._response_format_verdict
self._response_format_buffer = io.StringIO()
self._response_format_verdict = "undecided"
- return _held_answer(held, self.response_format_schema) if held else ""
+ if verdict == "pass_through":
+ answer: Final = partial_response_format_answer(held, schema)
+ close: Final = "" if answer is not None else self._release_answer_close()
+ self._streamed_answer_keys = None if answer is None else frozenset(answer)
+ return close
+ joined: Final = self._join_streamed_answer(held, schema)
+ if joined is not None:
+ return joined
+ released: Final = self._release_answer_close()
+ if self._split_answer_blocks or names_a_response_format_property(held, schema):
+ self._split_answer_blocks = (*self._split_answer_blocks, held)
+ return released
+ return released + (_held_answer(held, schema) if held else "")
+
+ def _drain_split_answer_blocks(self) -> str:
+ blocks: Final = self._split_answer_blocks
+ self._split_answer_blocks = ()
+ if self.response_format_schema is None or not blocks:
+ return ""
+ return _split_answer_text(blocks, self.response_format_schema)
def _handle_message_delta(self, chunk: dict) -> tuple[str, Usage | None, dict[str, Any] | None]:
"""
@@ -1207,11 +1308,13 @@ class ModelResponseIterator:
Tuple of (finish_reason, usage, container)
"""
message_delta: Final = MessageBlockDelta(**chunk)
- finish_reason = map_finish_reason(finish_reason=message_delta["delta"].get("stop_reason", "stop") or "stop")
- # Override finish_reason to "stop" if we converted response_format tools
- # (matches OpenAI behavior and non-streaming Anthropic implementation)
- if self.converted_response_format_tool:
- finish_reason = "stop"
+ provider_finish_reason: Final = map_finish_reason(
+ finish_reason=message_delta["delta"].get("stop_reason", "stop") or "stop"
+ )
+ called_only_the_response_format_tool: Final = (
+ provider_finish_reason == "tool_calls" and self.converted_response_format_tool and self.tool_index < 0
+ )
+ finish_reason: Final = "stop" if called_only_the_response_format_tool else provider_finish_reason
usage: Final = (
self._handle_usage(anthropic_usage_chunk=message_delta["usage"]) if "usage" in message_delta else None
)
@@ -1277,6 +1380,22 @@ class ModelResponseIterator:
return self._handle_accumulated_json_chunk(data_str)
# Sync iterator
+ def close(self) -> None:
+ response: Final = self._response
+ self._response = None
+ if response is not None:
+ response.close()
+
+ async def aclose(self) -> None:
+ response: Final = self._response
+ self._response = None
+ if response is None:
+ return
+ if self.sync_stream:
+ response.close()
+ return
+ await response.aclose()
+
def __iter__(self):
return self
diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py
index c43e5f9b05b..7abc13d5508 100644
--- a/litellm/llms/anthropic/chat/transformation.py
+++ b/litellm/llms/anthropic/chat/transformation.py
@@ -39,6 +39,7 @@ from litellm.llms.base_llm.base_utils import (
ResponseFormatSchema,
as_json_object,
response_format_schema,
+ split_response_format_answer,
type_to_response_format_param,
unwrap_response_format_envelope,
)
@@ -294,6 +295,17 @@ DROP_UNSUPPORTED_SPEED_WARNING: Final = (
"Dropping unsupported `speed` for model=%s (drop_params=True). Fast mode is only supported on select Opus models."
)
+SERVICE_TIER_PROVIDERS: Final = frozenset({"anthropic", "azure_ai"})
+
+OPENAI_TO_ANTHROPIC_SERVICE_TIER: Final = MappingProxyType(
+ {"auto": "auto", "default": "standard_only", "standard_only": "standard_only"}
+)
+
+DROP_UNSUPPORTED_SERVICE_TIER_WARNING: Final = (
+ "Dropping unsupported `service_tier` for model=%s (drop_params=True). "
+ "Anthropic only accepts 'auto' and 'standard_only'."
+)
+
class AnthropicConfig(AnthropicModelInfo, BaseConfig):
"""
@@ -335,6 +347,11 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
def _resolved_provider(self) -> str:
return self.custom_llm_provider or "anthropic"
+ def uses_native_structured_output(self, model: str) -> bool:
+ return AnthropicConfig._supports_model_capability(
+ model, "supports_native_structured_output", self._resolved_provider
+ )
+
@classmethod
def get_config(cls, *, model: str | None = None):
config: Final = super().get_config()
@@ -495,6 +512,23 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
)
optional_params.pop("speed", None)
+ @staticmethod
+ def _map_service_tier(value: str, model: str, drop_params: bool) -> str | None:
+ anthropic_tier: Final = OPENAI_TO_ANTHROPIC_SERVICE_TIER.get(value)
+ if anthropic_tier is not None:
+ return anthropic_tier
+ if not (litellm.drop_params or drop_params):
+ raise litellm.utils.UnsupportedParamsError(
+ message=(
+ f"{model} does not support service_tier={value!r}. Anthropic accepts 'auto', "
+ "'standard_only', or OpenAI's 'default' (sent as 'standard_only'). "
+ "To drop unsupported params, set `litellm.drop_params = True`."
+ ),
+ status_code=400,
+ )
+ litellm.verbose_logger.warning(DROP_UNSUPPORTED_SERVICE_TIER_WARNING, model)
+ return None
+
@classmethod
def maybe_drop_speed_param(
cls,
@@ -568,6 +602,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
params.append("thinking")
params.append("reasoning_effort")
+ if self._resolved_provider in SERVICE_TIER_PROVIDERS:
+ params.append("service_tier")
+
return params
@staticmethod
@@ -1600,11 +1637,7 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
output_key=param,
)
elif param == "response_format" and isinstance(value, dict):
- if AnthropicConfig._supports_model_capability(
- model,
- "supports_native_structured_output",
- self._resolved_provider,
- ):
+ if self.uses_native_structured_output(model):
_output_format = self.map_response_format_to_anthropic_output_format(value)
if _output_format is not None:
optional_params["output_format"] = _output_format
@@ -1719,6 +1752,13 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
elif param == "cache_control" and isinstance(value, dict):
# Pass through top-level cache_control for automatic prompt caching
optional_params["cache_control"] = value
+ elif (
+ param == "service_tier"
+ and isinstance(value, str)
+ and self._resolved_provider in SERVICE_TIER_PROVIDERS
+ and (anthropic_tier := AnthropicConfig._map_service_tier(value, model, drop_params)) is not None
+ ):
+ optional_params["service_tier"] = anthropic_tier # rebind-ok: in-place fill is this method's contract
## handle thinking tokens
self.update_optional_params_with_thinking_tokens(
@@ -2229,13 +2269,12 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
if json_tool.get("function", {}).get("arguments") is None:
return None, tool_calls, None
_message: Final = AnthropicConfig.convert_tool_response_to_message(
- tool_calls=[json_tool], response_format_schema=response_format_schema
+ tool_calls=tool_calls, response_format_schema=response_format_schema
)
return _message, [], None
- first_json: Final = tool_calls[json_indices[0]]
json_msg: Final = AnthropicConfig.convert_tool_response_to_message(
- [first_json], response_format_schema=response_format_schema
+ [tool_calls[i] for i in json_indices], response_format_schema=response_format_schema
)
extra_content: Final[str | None] = json_msg.content if json_msg is not None else None
filtered_tools: Final = [t for i, t in enumerate(tool_calls) if i not in json_indices]
@@ -2856,6 +2895,16 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig):
json_mode_content_str: Final[str | None] = tool_calls[0]["function"].get("arguments")
if json_mode_content_str is None:
return None
+ split_answer: Final = (
+ None
+ if response_format_schema is None
+ else split_response_format_answer(
+ tuple(tool_call["function"].get("arguments") or "" for tool_call in tool_calls),
+ response_format_schema,
+ )
+ )
+ if split_answer is not None:
+ return litellm.Message(content=json.dumps(split_answer))
try:
payload: Final[object] = json.loads(json_mode_content_str)
except (json.JSONDecodeError, RecursionError):
@@ -2906,11 +2955,13 @@ class JsonModeRequest(BaseModel):
return None if json_tool is None else response_format_schema(json_tool.get("input_schema"))
-def _valid_user_id(user_id: str) -> bool:
+def _valid_user_id(user_id: object) -> bool:
"""
- Validate that user_id is not an email or phone number.
- Returns: bool: True if valid (not email or phone), False otherwise
+ Validate that user_id is a string and not an email or phone number.
+ Returns: bool: True if valid (string, not email or phone), False otherwise
"""
+ if not isinstance(user_id, str):
+ return False
email_pattern: Final = r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$"
phone_pattern: Final = r"^\+?[\d\s\(\)-]{7,}$"
diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py
index 1bdb2b2db2c..9c8ab02d552 100644
--- a/litellm/llms/anthropic/pass_through/adapters/handler.py
+++ b/litellm/llms/anthropic/pass_through/adapters/handler.py
@@ -1,4 +1,4 @@
-from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping
+from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence
from typing import (
TYPE_CHECKING,
Final,
@@ -6,6 +6,7 @@ from typing import (
cast,
)
+from pydantic import TypeAdapter
from typing_extensions import TypedDict
import litellm
@@ -19,10 +20,20 @@ from litellm.llms.anthropic.pass_through.context_management import (
PolyfillResult,
apply_context_management,
)
+from litellm.llms.anthropic.pass_through.safeguards import (
+ SafeguardsEvaluator,
+ build_safeguards_evaluator,
+ read_classifier_model_setting,
+ requested_dangerous_tool_use,
+ with_safeguard_results,
+)
from litellm.llms.anthropic.pass_through.utils import (
+ ProxySubCallLoggingKwargs,
is_reasoning_auto_summary_enabled,
litellm_logging_obj_from_kwargs,
local_model_name,
+ proxy_sub_call_logging_kwargs,
+ proxy_sub_call_trace_id,
)
from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS
from litellm.types.llms.anthropic_messages.anthropic_response import (
@@ -38,6 +49,7 @@ if TYPE_CHECKING:
# Anthropic-only keys already mapped by the translator; strip on extra_kwargs re-merge.
ANTHROPIC_ONLY_REQUEST_KEYS: Final[frozenset[str]] = frozenset({"output_config", "safeguards"})
+_REQUEST_KWARGS: Final = TypeAdapter(Mapping[str, object])
_AnthropicMessages: TypeAlias = "list[dict[str, object]]"
_AnthropicSystem: TypeAlias = "str | list[dict[str, object]] | None"
@@ -69,6 +81,10 @@ def _proxy_router_fallback() -> "Router | None":
return _proxy_router
+def _router_or_proxy_fallback(requested: "Router | None") -> "Router | None":
+ return requested if requested is not None else _proxy_router_fallback()
+
+
def _extract_proxy_litellm_metadata(
kwargs: Mapping[str, object],
) -> "tuple[dict[str, object], UserAPIKeyAuth | None] | tuple[None, None]":
@@ -90,6 +106,31 @@ def _extract_proxy_litellm_metadata(
return litellm_metadata, user_api_key_auth
+def _requested_router(kwargs: Mapping[str, object]) -> "Router | None":
+ from litellm.router import Router
+
+ requested: Final = kwargs.get("litellm_router")
+ return requested if isinstance(requested, Router) else None
+
+
+def request_safeguards_evaluator(
+ kwargs: Mapping[str, object], messages: Sequence[Mapping[str, object]]
+) -> SafeguardsEvaluator | None:
+ safeguards: Final = kwargs.get("safeguards")
+ if requested_dangerous_tool_use(safeguards) is None or read_classifier_model_setting() is None:
+ return None
+ proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs)
+ return build_safeguards_evaluator(
+ safeguards=safeguards,
+ messages=messages,
+ litellm_metadata=proxy_litellm_metadata,
+ user_api_key_auth=user_api_key_auth,
+ llm_router=_router_or_proxy_fallback(_requested_router(kwargs)),
+ logging_kwargs=proxy_sub_call_logging_kwargs(kwargs),
+ parent_trace_id=proxy_sub_call_trace_id(kwargs),
+ )
+
+
async def _prepare_context_managed_request(
*,
model: str,
@@ -101,6 +142,8 @@ async def _prepare_context_managed_request(
additional_drop_params: list[str] | None,
llm_router: "Router | None",
user_api_key_auth: "UserAPIKeyAuth | None" = None,
+ sub_call_logging_kwargs: ProxySubCallLoggingKwargs | None = None,
+ parent_trace_id: str | None = None,
) -> PolyfillResult | None:
"""Apply client compaction history, then optional context_management polyfill."""
from litellm.llms.anthropic.pass_through.context_management.editors.compact import (
@@ -140,6 +183,8 @@ async def _prepare_context_managed_request(
additional_drop_params=additional_drop_params,
llm_router=llm_router,
user_api_key_auth=user_api_key_auth,
+ sub_call_logging_kwargs=sub_call_logging_kwargs,
+ parent_trace_id=parent_trace_id,
)
if polyfill_result is not None:
@@ -262,6 +307,8 @@ async def _run_polyfill_if_enabled(
additional_drop_params: list[str] | None,
llm_router: "Router | None",
user_api_key_auth: "UserAPIKeyAuth | None" = None,
+ sub_call_logging_kwargs: ProxySubCallLoggingKwargs | None = None,
+ parent_trace_id: str | None = None,
) -> PolyfillResult | None:
"""Run the async context_management polyfill if a spec is present.
@@ -288,6 +335,8 @@ async def _run_polyfill_if_enabled(
litellm_metadata=litellm_metadata,
llm_router=llm_router,
user_api_key_auth=user_api_key_auth,
+ sub_call_logging_kwargs=sub_call_logging_kwargs,
+ parent_trace_id=parent_trace_id,
)
except AnthropicContextManagementError:
# Surface validation errors so the endpoint can emit an Anthropic-format
@@ -622,12 +671,12 @@ class LiteLLMMessagesToCompletionTransformationHandler:
"""Handle non-Anthropic models asynchronously using the adapter"""
context_management: Final = kwargs.pop("context_management", None)
additional_drop_params: Final[list[str] | None] = kwargs.get("additional_drop_params", None)
- requested_router: Final[Router | None] = kwargs.pop("litellm_router", None)
- litellm_router: Final[Router | None] = (
- requested_router if requested_router is not None else _proxy_router_fallback()
- )
+ litellm_router: Final[Router | None] = _router_or_proxy_fallback(kwargs.pop("litellm_router", None))
proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs)
+ request_kwargs: Final = _REQUEST_KWARGS.validate_python(kwargs)
+ sub_call_logging_kwargs: Final = proxy_sub_call_logging_kwargs(request_kwargs)
+ sub_call_trace_id: Final = proxy_sub_call_trace_id(request_kwargs)
polyfill_result: Final = await _prepare_context_managed_request(
model=model,
@@ -639,10 +688,21 @@ class LiteLLMMessagesToCompletionTransformationHandler:
additional_drop_params=additional_drop_params,
llm_router=litellm_router,
user_api_key_auth=user_api_key_auth,
+ sub_call_logging_kwargs=sub_call_logging_kwargs,
+ parent_trace_id=sub_call_trace_id,
)
effective_messages: Final = polyfill_result.messages if polyfill_result is not None else messages
effective_system: Final = polyfill_result.system if polyfill_result is not None else system
+ safeguards_evaluator: Final = build_safeguards_evaluator(
+ safeguards=kwargs.get("safeguards"),
+ messages=messages,
+ litellm_metadata=proxy_litellm_metadata,
+ user_api_key_auth=user_api_key_auth,
+ llm_router=litellm_router,
+ logging_kwargs=sub_call_logging_kwargs,
+ parent_trace_id=sub_call_trace_id,
+ )
(
completion_kwargs,
@@ -675,6 +735,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
polyfill_result=polyfill_result,
is_async=True,
litellm_logging_obj=litellm_logging_obj_from_kwargs(kwargs),
+ safeguards_evaluator=safeguards_evaluator,
)
if transformed_stream is not None:
return transformed_stream
@@ -686,7 +747,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
polyfill_result=polyfill_result,
)
if anthropic_response is not None:
- return anthropic_response
+ return await with_safeguard_results(anthropic_response, safeguards_evaluator)
raise ValueError("Failed to transform response to Anthropic format")
@staticmethod
@@ -763,6 +824,7 @@ class LiteLLMMessagesToCompletionTransformationHandler:
polyfill_result: PolyfillResult | None = None
else:
proxy_litellm_metadata, user_api_key_auth = _extract_proxy_litellm_metadata(kwargs)
+ request_kwargs: Final = _REQUEST_KWARGS.validate_python(kwargs)
polyfill_result = run_async_function(
_prepare_context_managed_request,
model=model,
@@ -774,6 +836,8 @@ class LiteLLMMessagesToCompletionTransformationHandler:
additional_drop_params=additional_drop_params,
llm_router=litellm_router,
user_api_key_auth=user_api_key_auth,
+ sub_call_logging_kwargs=proxy_sub_call_logging_kwargs(request_kwargs),
+ parent_trace_id=proxy_sub_call_trace_id(request_kwargs),
)
effective_messages: Final = polyfill_result.messages if polyfill_result is not None else messages
diff --git a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py
index 636051d34ce..8d26b1e6387 100644
--- a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py
+++ b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py
@@ -15,12 +15,15 @@ from typing import (
get_args,
)
+from pydantic import TypeAdapter
from typing_extensions import assert_never
from litellm._logging import verbose_logger
from litellm._uuid import uuid
from litellm.exceptions import ContextWindowExceededError, MidStreamFallbackError
from litellm.litellm_core_utils.hidden_params import set_hidden_params
+from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults
+from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame
from litellm.llms.base_llm.chat.transformation import BaseLLMException
from litellm.types.llms.anthropic import (
AppliedEdit,
@@ -37,11 +40,13 @@ from litellm.types.utils import AdapterCompletionStreamWrapper, Delta
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
+ from litellm.llms.anthropic.pass_through.safeguards import SafeguardsEvaluator
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import ModelResponseStream
_STREAMING_DELTA_TYPES: Final = frozenset(get_args(StreamingContentBlockDeltaType))
+_SSE_EVENT: Final = TypeAdapter(Mapping[str, object])
class _UsageDeltaWithIterations(UsageDelta, total=False):
@@ -342,6 +347,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
compaction_block: CompactionBlock | None = None,
iterations_usage: list[UsageIteration] | None = None,
litellm_logging_obj: "LiteLLMLoggingObject | None" = None,
+ safeguards_evaluator: "SafeguardsEvaluator | None" = None,
):
# Wrap the upstream stream so chunks that carry both content and a
# finish_reason (fake-streamed providers) are split into two — see
@@ -352,6 +358,7 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
if litellm_logging_obj is not None:
litellm_logging_obj.record_streamed_anthropic_message_id(self._message_id)
self.litellm_logging_obj = litellm_logging_obj
+ self.safeguards_evaluator: Final = safeguards_evaluator
# Mapping of truncated tool names to original names (for OpenAI's 64-char limit)
self.tool_name_mapping = tool_name_mapping or {}
# Polyfill applied_edits on final message_delta.
@@ -1048,14 +1055,17 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper):
Async version of anthropic_sse_wrapper.
Convert AnthropicStreamWrapper dict chunks to Server-Sent Events format.
"""
+ safeguard_results: Final = (
+ StreamedSafeguardResults(self.safeguards_evaluator) if self.safeguards_evaluator is not None else None
+ )
try:
async for chunk in self:
- if isinstance(chunk, dict):
- event_type: str = str(chunk.get("type", "message"))
- payload = f"event: {event_type}\ndata: {json.dumps(chunk)}\n\n"
- yield payload.encode()
- else:
+ if not isinstance(chunk, dict):
yield chunk
+ elif safeguard_results is None:
+ yield anthropic_sse_frame(chunk)
+ else:
+ yield anthropic_sse_frame(await safeguard_results.observe(_SSE_EVENT.validate_python(chunk)))
except Exception as e: # noqa: BLE001 # boundary before the socket: any upstream failure becomes an Anthropic error event
verbose_logger.exception("Anthropic Adapter - mid-stream error: %s", e)
logging_obj: Final = self.litellm_logging_obj
diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py
index d2a0849ffef..1e276b7c9c7 100644
--- a/litellm/llms/anthropic/pass_through/adapters/transformation.py
+++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py
@@ -206,6 +206,7 @@ from .streaming_iterator import AnthropicSSEStream, AnthropicStreamWrapper
if TYPE_CHECKING:
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
+ from litellm.llms.anthropic.pass_through.safeguards import SafeguardsEvaluator
from litellm.types.llms.anthropic import ContentBlockContentBlockDict
ToolResultContent: TypeAlias = str | list[ToolMessageContentPart]
@@ -312,6 +313,7 @@ class AnthropicAdapter:
polyfill_result: PolyfillResult | None = None,
is_async: bool = True,
litellm_logging_obj: "LiteLLMLoggingObject | None" = None,
+ safeguards_evaluator: "SafeguardsEvaluator | None" = None,
) -> AsyncIterator[bytes] | Iterator[bytes] | None:
"""
Translate OpenAI streaming response to Anthropic format.
@@ -339,6 +341,7 @@ class AnthropicAdapter:
compaction_block=compaction_block,
iterations_usage=iterations_usage,
litellm_logging_obj=litellm_logging_obj,
+ safeguards_evaluator=safeguards_evaluator,
)
# Return the SSE-wrapped version for proper event formatting.
if is_async:
diff --git a/litellm/llms/anthropic/pass_through/caller_model_checks.py b/litellm/llms/anthropic/pass_through/caller_model_checks.py
new file mode 100644
index 00000000000..7eb6a8f9266
--- /dev/null
+++ b/litellm/llms/anthropic/pass_through/caller_model_checks.py
@@ -0,0 +1,498 @@
+"""Checks a proxy-made sub-call runs against the caller of the parent request.
+
+The compaction summary call and the safeguards classifier call go to a model the
+admin configured, not the one the client asked for, and they never pass back
+through the proxy's auth, budget, or rate-limit hooks. These checks apply the
+caller's model allowlists, per-model budgets, and RPM/TPM limits to that model
+before the sub-call runs.
+"""
+
+from collections.abc import Awaitable, Mapping, Sequence
+from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeAlias
+
+from pydantic import TypeAdapter, ValidationError
+
+import litellm
+from litellm._logging import verbose_logger
+
+if TYPE_CHECKING:
+ from starlette.requests import Request
+
+ from litellm.proxy.hooks.parallel_request_limiter_v3 import (
+ RateLimitDescriptor,
+ RateLimitDescriptorRateLimitObject,
+ RateLimitResponse,
+ )
+ from litellm.router import Router
+ from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
+
+CallerModelVerdict: TypeAlias = Literal["allowed", "denied", "denied_for_now"]
+_DEFAULT_SUB_CALL_ROUTE: Final = "/v1/messages"
+_MODEL_SCOPED_DESCRIPTOR_PREFIX: Final = "model_per_"
+_HEADERS: Final = TypeAdapter(Mapping[str, str])
+_LASTING_REFUSAL_STATUSES: Final = frozenset({401, 403})
+_LASTING_REFUSAL_CODES: Final = frozenset(str(status) for status in _LASTING_REFUSAL_STATUSES)
+
+
+class _TeamMemberModelBudgetLimiter(Protocol):
+ async def is_team_member_within_model_budget(
+ self,
+ user_id: str,
+ team_id: str,
+ team_member_model_max_budget: Mapping[str, object],
+ model: str,
+ ) -> bool: ...
+
+
+class _CreateRateLimitDescriptors(Protocol):
+ def __call__(
+ self,
+ *,
+ user_api_key_dict: "UserAPIKeyAuth",
+ data: Mapping[str, str],
+ rpm_limit_type: object,
+ tpm_limit_type: object,
+ model_has_failures: bool,
+ ) -> "Sequence[RateLimitDescriptor]": ...
+
+
+class _AddModelRateLimitDescriptor(Protocol):
+ def __call__(
+ self,
+ *,
+ user_api_key_dict: "UserAPIKeyAuth",
+ requested_model: str,
+ descriptors: "Sequence[RateLimitDescriptor]",
+ ) -> None: ...
+
+
+class _CreateOrgRateLimitDescriptors(Protocol):
+ def __call__(
+ self, user_api_key_dict: "UserAPIKeyAuth", requested_model: str | None = None
+ ) -> "Sequence[RateLimitDescriptor]": ...
+
+
+class _GetProxyHook(Protocol):
+ def __call__(self, hook: str) -> object: ...
+
+
+class _ShouldRateLimit(Protocol):
+ def __call__(
+ self,
+ *,
+ descriptors: "Sequence[RateLimitDescriptor]",
+ parent_otel_span: object,
+ read_only: bool,
+ ) -> "Awaitable[RateLimitResponse]": ...
+
+
+async def caller_model_verdict(
+ user_api_key_auth: Optional["UserAPIKeyAuth"],
+ model: str,
+ llm_router: Optional["Router"],
+ parent_litellm_metadata: Mapping[str, object] | None = None,
+ parent_trace_id: str | None = None,
+) -> CallerModelVerdict:
+ """``denied`` when an allowlist or a per-model budget refuses ``model`` for the caller, which only an
+ admin or the next budget window changes; ``denied_for_now`` when a rate-limit window is full or a check
+ could not run, which the caller's next request may clear."""
+ access: Final = await caller_can_call_model(
+ user_api_key_auth=user_api_key_auth,
+ model=model,
+ llm_router=llm_router,
+ parent_litellm_metadata=parent_litellm_metadata,
+ parent_trace_id=parent_trace_id,
+ )
+ if access != "allowed":
+ return access
+ budget: Final = await caller_within_model_budget(user_api_key_auth=user_api_key_auth, model=model)
+ if budget != "allowed":
+ return budget
+ return await caller_within_model_rate_limit(user_api_key_auth=user_api_key_auth, model=model)
+
+
+async def caller_can_call_model(
+ user_api_key_auth: Optional["UserAPIKeyAuth"],
+ model: str,
+ llm_router: Optional["Router"],
+ parent_litellm_metadata: Mapping[str, object] | None = None,
+ parent_trace_id: str | None = None,
+) -> CallerModelVerdict:
+ """Run the proxy's own model-access checks for ``model`` against the caller of the parent request.
+
+ These are the checks the parent request passed for its own model, so every scope they cover (key, team,
+ team member, user, project, end user, managed agent grants) applies to the sub-call's model too, and a
+ deployment that skips them for the parent (no-auth dev mode, custom auth without common checks) skips
+ them here. Budgets are checked too, so a spent budget the parent skipped for a free model still refuses
+ a paid one. ``denied`` answers a budget refusal or a 401 or 403, which only an admin or the next budget
+ window changes. ``denied_for_now`` answers any other failure, such as a team read that 404s, since the
+ parent passed the same checks this turn and the next turn may clear it. ``allowed`` answers a caller
+ outside the proxy.
+ """
+ if user_api_key_auth is None:
+ return "allowed"
+ try:
+ from fastapi import HTTPException
+
+ from litellm.proxy._types import ProxyException
+ from litellm.proxy.auth.user_api_key_auth import (
+ enforce_key_and_fallback_model_access, # pyright: ignore[reportUnknownVariableType] # takes the bare dict and list the proxy passes
+ run_centralized_common_checks,
+ )
+ from litellm.proxy.proxy_server import (
+ llm_model_list, # pyright: ignore[reportUnknownVariableType] # proxy global has untyped deployment entries
+ )
+ except ImportError:
+ return "allowed"
+
+ route: Final = user_api_key_auth.request_route or _DEFAULT_SUB_CALL_ROUTE
+ try:
+ request: Final = _sub_call_request(route, parent_litellm_metadata)
+ request_data: Final = _sub_call_request_data(model, parent_litellm_metadata, parent_trace_id)
+ caller: Final = user_api_key_auth.model_copy(update={"budget_reservation": None})
+ await enforce_key_and_fallback_model_access(
+ valid_token=caller,
+ request_data=request_data,
+ route=route,
+ request=request,
+ llm_model_list=llm_model_list, # pyright: ignore[reportUnknownArgumentType] # proxy global has untyped deployment entries
+ llm_router=llm_router,
+ )
+ await run_centralized_common_checks(caller, request, request_data, route, sub_call=True)
+ except litellm.BudgetExceededError:
+ return "denied"
+ except HTTPException as e:
+ return "denied" if e.status_code in _LASTING_REFUSAL_STATUSES else _check_failed_this_turn(model, e)
+ except ProxyException as e:
+ return "denied" if e.code in _LASTING_REFUSAL_CODES else _check_failed_this_turn(model, e)
+ except Exception as e: # noqa: BLE001 # a check that cannot run must not invoke the model with the caller's limits dropped
+ return _check_failed_this_turn(model, e)
+ return "allowed"
+
+
+def _check_failed_this_turn(model: str, error: Exception) -> CallerModelVerdict:
+ verbose_logger.warning(
+ "proxy sub-call: model access check for model=%s failed this turn; denying for now: %s", model, error
+ )
+ return "denied_for_now"
+
+
+def _sub_call_request(route: str, parent_litellm_metadata: Mapping[str, object] | None) -> "Request":
+ from starlette.requests import Request
+
+ headers: Final = _parent_headers(parent_litellm_metadata)
+ return Request(
+ scope={
+ "type": "http",
+ "method": "POST",
+ "path": route,
+ "scheme": "http",
+ "query_string": b"",
+ "root_path": "",
+ "headers": tuple(
+ (name.lower().encode("latin-1"), value.encode("latin-1")) for name, value in headers.items()
+ ),
+ }
+ )
+
+
+def _parent_headers(parent_litellm_metadata: Mapping[str, object] | None) -> Mapping[str, str]:
+ if not parent_litellm_metadata:
+ return {}
+ try:
+ return _HEADERS.validate_python(parent_litellm_metadata.get("headers") or {})
+ except ValidationError:
+ return {}
+
+
+def _sub_call_request_data(
+ model: str, parent_litellm_metadata: Mapping[str, object] | None, parent_trace_id: str | None
+) -> dict[str, object]: # mutable-ok: the proxy's checks merge metadata and tags into the body they are handed
+ parent: Final[Mapping[str, object]] = parent_litellm_metadata or {}
+ end_user_id: Final = parent.get("user_api_key_end_user_id")
+ metadata_trace_id: Final = parent.get("trace_id")
+ trace_id: Final = metadata_trace_id if isinstance(metadata_trace_id, str) and metadata_trace_id else parent_trace_id
+ return {
+ "model": model,
+ **({"user": end_user_id} if isinstance(end_user_id, str) and end_user_id else {}),
+ **({"litellm_trace_id": trace_id} if isinstance(trace_id, str) and trace_id else {}),
+ }
+
+
+async def caller_within_model_budget(
+ user_api_key_auth: Optional["UserAPIKeyAuth"],
+ model: str,
+) -> CallerModelVerdict:
+ """Return ``allowed`` when the caller is within their per-model budget for ``model``.
+
+ Mirrors the per-model budget enforcement that ``user_api_key_auth`` runs for
+ the client-requested model. Returns ``allowed`` outside the proxy or when no
+ per-model budget is configured, ``denied`` over a budget, and ``denied_for_now``
+ when a budget read fails.
+
+ Every scope is checked because the sub-call's spend is charged to every
+ scope: the key, team, user, team-member and end-user budgets ride into the
+ sub-call's metadata, so skipping one of them would let the sub-call
+ increment a counter it can never be refused by.
+ """
+ if user_api_key_auth is None:
+ return "allowed"
+ try:
+ from litellm.proxy.proxy_server import model_max_budget_limiter
+ except Exception:
+ return "allowed"
+
+ model_max_budget: Final = getattr(user_api_key_auth, "model_max_budget", None)
+ token: Final = getattr(user_api_key_auth, "token", None)
+ if isinstance(model_max_budget, dict) and model_max_budget and token is not None:
+ try:
+ await model_max_budget_limiter.is_key_within_model_budget(
+ user_api_key_dict=user_api_key_auth,
+ model=model,
+ )
+ except litellm.BudgetExceededError:
+ return "denied"
+ except Exception as e:
+ verbose_logger.warning(
+ "proxy sub-call: unexpected error during key model-budget check for model=%s; denying: %s",
+ model,
+ e,
+ )
+ return "denied_for_now"
+
+ user_model_max_budget: Final = user_api_key_auth.user_model_max_budget
+ user_id: Final = user_api_key_auth.user_id
+ if isinstance(user_model_max_budget, dict) and user_model_max_budget and user_id is not None:
+ try:
+ await model_max_budget_limiter.is_user_within_model_budget(
+ user_id=user_id,
+ user_model_max_budget=user_model_max_budget,
+ model=model,
+ )
+ except litellm.BudgetExceededError:
+ return "denied"
+ except Exception as e: # noqa: BLE001 # a budget gate denies on any failure, as the key and end-user scopes do
+ verbose_logger.warning(
+ "proxy sub-call: unexpected error during user model-budget check for model=%s; denying: %s",
+ model,
+ e,
+ )
+ return "denied_for_now"
+
+ team_model_max_budget: Final = user_api_key_auth.team_model_max_budget
+ team_id: Final = user_api_key_auth.team_id
+ if isinstance(team_model_max_budget, dict) and team_model_max_budget and team_id is not None:
+ try:
+ await model_max_budget_limiter.is_team_within_model_budget(
+ team_id=team_id,
+ team_model_max_budget=team_model_max_budget,
+ key_model_max_budget=model_max_budget if isinstance(model_max_budget, dict) else None,
+ model=model,
+ )
+ except litellm.BudgetExceededError:
+ return "denied"
+ except Exception as e: # noqa: BLE001 # a budget gate denies on any failure, as the other scopes do
+ verbose_logger.warning(
+ "proxy sub-call: unexpected error during team model-budget check for model=%s; denying: %s",
+ model,
+ e,
+ )
+ return "denied_for_now"
+
+ end_user_model_max_budget: Final[dict[str, object] | None] = getattr(
+ user_api_key_auth, "end_user_model_max_budget", None
+ )
+ end_user_id: Final[str | None] = getattr(user_api_key_auth, "end_user_id", None)
+ if isinstance(end_user_model_max_budget, dict) and end_user_model_max_budget and end_user_id is not None:
+ try:
+ await model_max_budget_limiter.is_end_user_within_model_budget(
+ end_user_id=end_user_id,
+ end_user_model_max_budget=end_user_model_max_budget,
+ model=model,
+ )
+ except litellm.BudgetExceededError:
+ return "denied"
+ except Exception as e:
+ verbose_logger.warning(
+ "proxy sub-call: unexpected error during end-user model-budget check for model=%s; denying: %s",
+ model,
+ e,
+ )
+ return "denied_for_now"
+
+ return await _caller_within_team_member_model_budget(
+ user_api_key_auth=user_api_key_auth,
+ model_max_budget_limiter=model_max_budget_limiter,
+ model=model,
+ )
+
+
+async def _caller_within_team_member_model_budget(
+ user_api_key_auth: "UserAPIKeyAuth",
+ model_max_budget_limiter: _TeamMemberModelBudgetLimiter,
+ model: str,
+) -> CallerModelVerdict:
+ team_member_model_max_budget: Final = user_api_key_auth.team_member_model_max_budget
+ user_id: Final = user_api_key_auth.user_id
+ team_id: Final = user_api_key_auth.team_id
+ if (
+ not isinstance(team_member_model_max_budget, Mapping)
+ or not team_member_model_max_budget
+ or not isinstance(user_id, str)
+ or not isinstance(team_id, str)
+ ):
+ return "allowed"
+ try:
+ await model_max_budget_limiter.is_team_member_within_model_budget(
+ user_id=user_id,
+ team_id=team_id,
+ team_member_model_max_budget=team_member_model_max_budget,
+ model=model,
+ )
+ except litellm.BudgetExceededError:
+ return "denied"
+ except Exception as e: # noqa: BLE001 # a budget gate denies on any failure, as the other scopes do
+ verbose_logger.warning(
+ "proxy sub-call: unexpected error during team member model-budget check for model=%s; denying: %s",
+ model,
+ e,
+ )
+ return "denied_for_now"
+ return "allowed"
+
+
+def _sub_call_descriptors(
+ descriptors: "Sequence[RateLimitDescriptor]",
+) -> "tuple[RateLimitDescriptor, ...]":
+ return tuple(_sub_call_descriptor(descriptor) for descriptor in descriptors)
+
+
+def _sub_call_descriptor(descriptor: "RateLimitDescriptor") -> "RateLimitDescriptor":
+ rate_limit: Final = descriptor.get("rate_limit")
+ if rate_limit is None:
+ return descriptor
+ caller_wide: Final = not descriptor["key"].startswith(_MODEL_SCOPED_DESCRIPTOR_PREFIX)
+ windowed_limits: Final[RateLimitDescriptorRateLimitObject] = {
+ "requests_per_unit": None if caller_wide else rate_limit.get("requests_per_unit"),
+ "tokens_per_unit": rate_limit.get("tokens_per_unit"),
+ "window_size": rate_limit.get("window_size"),
+ }
+ return {**descriptor, "rate_limit": windowed_limits}
+
+
+async def caller_within_model_rate_limit(
+ user_api_key_auth: Optional["UserAPIKeyAuth"],
+ model: str,
+) -> CallerModelVerdict:
+ """Return ``allowed`` when the caller is within their configured RPM/TPM limits for ``model``.
+
+ This mirrors the read side of
+ ``_PROXY_MaxParallelRequestsHandler_v3.async_pre_call_hook`` for ``model``:
+ it builds the same descriptor set and runs the check in ``read_only`` mode
+ so no counter is reserved or incremented. The sub-call's tokens are still
+ charged once by the limiter's post-call success hook (via the propagated
+ ``litellm_metadata``), but it takes no request slot, so it does not count
+ toward the caller's RPM. ``max_parallel_requests`` gauges and the RPM of
+ every caller-wide scope (key, user, team, org, ...) are left out of the
+ check: the sub-call runs inside the caller's already admitted request,
+ whose own slot would otherwise count against it. The RPM of a limit on
+ ``model`` itself still gates the sub-call, as do token limits.
+
+ Returns ``allowed`` outside the proxy, when the active limiter does not
+ expose the read-only descriptor check (legacy limiter), or when the
+ descriptor set cannot be built. ``denied_for_now`` comes from a definitive
+ ``OVER_LIMIT`` response, which the next window clears, and from the limiter's own fail-closed rejection
+ (``RateLimitUnverifiableError``, raised when ``fail_closed_rate_limit_enforcement``
+ is on and the counters could not be verified), so any other internal error
+ here lets the sub-call run rather than blocking every one.
+ """
+ if user_api_key_auth is None:
+ return "allowed"
+ try:
+ from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitUnverifiableError
+ from litellm.proxy.proxy_server import proxy_logging_obj
+ except Exception:
+ return "allowed"
+
+ get_proxy_hook: Final[_GetProxyHook | None] = getattr(proxy_logging_obj, "get_proxy_hook", None)
+ limiter: Final[object] = get_proxy_hook("parallel_request_limiter") if get_proxy_hook is not None else None
+ return await rate_limit_verdict(
+ user_api_key_auth=user_api_key_auth, model=model, limiter=limiter, unverifiable=RateLimitUnverifiableError
+ )
+
+
+async def rate_limit_verdict(
+ user_api_key_auth: "UserAPIKeyAuth",
+ model: str,
+ limiter: object,
+ unverifiable: type[Exception],
+) -> CallerModelVerdict:
+ should_rate_limit_check: Final[_ShouldRateLimit | None] = getattr(limiter, "should_rate_limit", None)
+ create_descriptors: Final[_CreateRateLimitDescriptors | None] = getattr(
+ limiter, "create_rate_limit_descriptors", None
+ )
+ add_team_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr(
+ limiter, "_add_team_model_rate_limit_descriptor_from_metadata", None
+ )
+ add_project_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr(
+ limiter, "_add_project_model_rate_limit_descriptor_from_metadata", None
+ )
+ create_org_descriptors: Final[_CreateOrgRateLimitDescriptors | None] = getattr(
+ limiter, "create_organization_rate_limit_descriptor", None
+ )
+ if (
+ limiter is None
+ or should_rate_limit_check is None
+ or create_descriptors is None
+ or add_team_descriptor is None
+ or add_project_descriptor is None
+ or create_org_descriptors is None
+ ):
+ return "allowed"
+
+ try:
+ metadata: Final[Mapping[str, object]] = getattr(user_api_key_auth, "metadata", None) or {}
+ data: Final = {"model": model}
+ base_descriptors: Final = create_descriptors(
+ user_api_key_dict=user_api_key_auth,
+ data=data,
+ rpm_limit_type=metadata.get("rpm_limit_type"),
+ tpm_limit_type=metadata.get("tpm_limit_type"),
+ model_has_failures=False,
+ )
+ add_team_descriptor(
+ user_api_key_dict=user_api_key_auth,
+ requested_model=model,
+ descriptors=base_descriptors,
+ )
+ add_project_descriptor(
+ user_api_key_dict=user_api_key_auth,
+ requested_model=model,
+ descriptors=base_descriptors,
+ )
+ descriptors: Final = _sub_call_descriptors(
+ (*base_descriptors, *create_org_descriptors(user_api_key_auth, model))
+ )
+ if not descriptors:
+ return "allowed"
+ parent_otel_span: Final[object] = getattr(user_api_key_auth, "parent_otel_span", None)
+ response: Final[RateLimitResponse] = await should_rate_limit_check(
+ descriptors=descriptors,
+ parent_otel_span=parent_otel_span,
+ read_only=True,
+ )
+ except unverifiable as e:
+ verbose_logger.warning(
+ "proxy sub-call: rate-limit counters for model=%s could not be verified; denying for now: %s",
+ model,
+ e,
+ )
+ return "denied_for_now"
+ except Exception as e:
+ verbose_logger.warning(
+ "proxy sub-call: unexpected error during rate-limit check for model=%s; allowing: %s",
+ model,
+ e,
+ )
+ return "allowed"
+ return "allowed" if response.get("overall_code") != "OVER_LIMIT" else "denied_for_now"
diff --git a/litellm/llms/anthropic/pass_through/context_management/dispatcher.py b/litellm/llms/anthropic/pass_through/context_management/dispatcher.py
index 9b4e12ef2f9..27b7412292e 100644
--- a/litellm/llms/anthropic/pass_through/context_management/dispatcher.py
+++ b/litellm/llms/anthropic/pass_through/context_management/dispatcher.py
@@ -13,6 +13,7 @@ from .editors import apply_clear_tool_uses_20250919, apply_compact_20260112
from .result import PolyfillResult
if TYPE_CHECKING:
+ from litellm.llms.anthropic.pass_through.utils import ProxySubCallLoggingKwargs
from litellm.router import Router
from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth
@@ -82,6 +83,8 @@ async def apply_context_management(
litellm_metadata: dict[str, object] | None = None,
llm_router: "Router | None" = None,
user_api_key_auth: "UserAPIKeyAuth | None" = None,
+ sub_call_logging_kwargs: "ProxySubCallLoggingKwargs | None" = None,
+ parent_trace_id: str | None = None,
) -> PolyfillResult:
"""Run edits in order; return a single ``PolyfillResult``.
@@ -123,6 +126,8 @@ async def apply_context_management(
litellm_metadata=litellm_metadata,
llm_router=llm_router,
user_api_key_auth=user_api_key_auth,
+ sub_call_logging_kwargs=sub_call_logging_kwargs,
+ parent_trace_id=parent_trace_id,
)
if editor_is_async
else await asyncify(editor)(
diff --git a/litellm/llms/anthropic/pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/pass_through/context_management/editors/compact.py
index deb3d849ae9..09f49e9b663 100644
--- a/litellm/llms/anthropic/pass_through/context_management/editors/compact.py
+++ b/litellm/llms/anthropic/pass_through/context_management/editors/compact.py
@@ -14,13 +14,23 @@ Mirrors Anthropic's native ``compact_20260112`` for non-Anthropic providers:
import re
from collections.abc import Awaitable, Mapping, Sequence
-from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypeVar, Union, cast
+from typing import TYPE_CHECKING, Final, Optional, Protocol, TypeVar, Union, cast
from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack
import litellm
from litellm._logging import verbose_logger
from litellm.litellm_core_utils.asyncify import asyncify
+from litellm.llms.anthropic.pass_through.caller_model_checks import (
+ caller_can_call_model,
+ caller_within_model_budget,
+ caller_within_model_rate_limit,
+)
+from litellm.llms.anthropic.pass_through.utils import (
+ ProxySubCallLoggingKwargs,
+ proxy_general_settings,
+ proxy_sub_call_metadata,
+)
from litellm.types.llms.anthropic import (
AppliedEdit,
CompactionBlock,
@@ -29,11 +39,6 @@ from litellm.types.llms.anthropic import (
if TYPE_CHECKING:
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
- from litellm.proxy.hooks.parallel_request_limiter_v3 import (
- RateLimitDescriptor,
- RateLimitDescriptorRateLimitObject,
- RateLimitResponse,
- )
from litellm.router import Router
from litellm.types.llms.anthropic import (
AllAnthropicPassThroughMessageValues,
@@ -58,38 +63,6 @@ from ..constants import (
from ..errors import AnthropicContextManagementError
from ..result import PolyfillResult
-# Auth metadata fields propagated from the parent request to the summary call
-# so the summary's spend is attributed to the same scopes. The list mirrors the
-# fields populated by
-# ``LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata``.
-# The three ``*_model_max_budget`` fields
-# are what ``_PROXY_VirtualKeyModelMaxBudgetLimiter`` reads post-call to update
-# the per-model spend caches, so without them the summary spend would never
-# count against the caller's model budget. ``user_api_key_end_user_id`` /
-# ``user_api_key_project_id`` are the scope identifiers the post-call spend hook
-# and rate limiter key their counters on, and ``user_api_end_user_max_budget``
-# is the end-user budget the cost callback enforces — without these the summary
-# tokens escape the caller's end-user/project budgets and counters.
-_PROPAGATED_METADATA_KEYS: Final = (
- "user_api_key",
- "user_api_key_alias",
- "user_api_key_team_id",
- "user_api_key_team_alias",
- "user_api_key_user_id",
- "user_api_key_user_email",
- "user_api_key_org_id",
- "user_api_key_project_id",
- "user_api_key_end_user_id",
- "user_api_end_user_max_budget",
- "user_api_key_model_max_budget",
- "user_api_key_team_model_max_budget",
- "user_api_key_user_model_max_budget",
- "user_api_key_team_member_model_max_budget",
- "user_api_key_end_user_model_max_budget",
- "litellm_call_id",
- "litellm_parent_otel_span",
-)
-
_SUMMARY_TAG_RE: Final = re.compile(r"Error: invalid_issuer
" + "This authorization response came from a different identity provider than the one this " + "MCP server is configured to use.
" + "You can close this window and try again.
" + ), wrong.text + assert "location" not in wrong.headers, wrong.headers + assert wrong_code not in wrong.text, wrong.text + _assert_cleared_oauth_state_cookie(wrong, wrong_cookie) + + correct_relay, correct_cookie, correct_cookie_value = _start_callback_flow(gateway, alias, key, client_id) + correct: Final = gateway.client.get( + "/callback", + params={"code": "c", "state": correct_relay, "iss": auth.issuer}, + cookies={correct_cookie: correct_cookie_value}, + ) + assert correct.status_code == 302, correct.text + correct_location: Final = urlsplit(correct.headers["location"]) + assert correct_location.scheme + "://" + correct_location.netloc + correct_location.path == CLIENT_REDIRECT, ( + correct.headers["location"] + ) + correct_query: Final = parse_qs(correct_location.query) + assert set(correct_query) == {"code", "state"}, correct.headers["location"] + assert correct_query["state"] == ["client-state"], correct.headers["location"] + assert len(correct_query["code"]) == 1 and correct_query["code"][0].startswith("llm_scope_"), correct.headers[ + "location" + ] + _assert_cleared_oauth_state_cookie(correct, correct_cookie) + + error_relay, error_cookie, error_cookie_value = _start_callback_flow(gateway, alias, key, client_id) + error: Final = gateway.client.get( + "/callback", + params={"error": "access_denied", "error_description": "x", "state": error_relay}, + cookies={error_cookie: error_cookie_value}, + ) + assert error.status_code == 302, error.text + error_location: Final = urlsplit(error.headers["location"]) + assert parse_qs(error_location.query) == { + "error": ["access_denied"], + "error_description": ["x"], + "state": ["client-state"], + }, error.headers["location"] + _assert_cleared_oauth_state_cookie(error, error_cookie) + + error_issuer_relay, error_issuer_cookie, error_issuer_cookie_value = _start_callback_flow( + gateway, alias, key, client_id + ) + error_issuer: Final = gateway.client.get( + "/callback", + params={ + "error": "access_denied", + "error_description": "x", + "state": error_issuer_relay, + "iss": "http://127.0.0.1:1/other", + }, + cookies={error_issuer_cookie: error_issuer_cookie_value}, + ) + assert error_issuer.status_code == 400, error_issuer.text + assert error_issuer.headers["content-type"] == "text/html; charset=utf-8", error_issuer.headers + assert error_issuer.text == ( + "Error: invalid_issuer
" + "Unexpected authorization issuer
" + "You can close this window and try again.
" + ), error_issuer.text + assert "location" not in error_issuer.headers, error_issuer.headers + _assert_cleared_oauth_state_cookie(error_issuer, error_issuer_cookie) + + missing: Final = gateway.client.get("/callback") + assert missing.status_code == 400, missing.text + assert missing.headers["content-type"] == "text/html; charset=utf-8", missing.headers + assert missing.text == ( + "Error: invalid_request
" + "Missing authorization 'code' and 'state' parameter(s).
" + "You can close this window and try again.
" + ), missing.text + + +def test_callback_refuses_a_response_without_iss_when_the_authorization_server_advertises_iss_support( + gateway: Gateway, +) -> None: + wire_holder: Final[list[Wire]] = [] + + def respond(request: Request) -> Reply: + path: Final = urlsplit(request.target).path + if request.method == "GET" and path in ( + "/.well-known/oauth-authorization-server", + "/.well-known/openid-configuration", + ): + issuer: Final = wire_holder[0].url + return Reply( + body=json.dumps( + { + "issuer": issuer, + "authorization_endpoint": issuer + "/authorize", + "token_endpoint": issuer + "/token", + "registration_endpoint": issuer + "/register", + "response_types_supported": ["code"], + "grant_types_supported": ["authorization_code", "refresh_token"], + "code_challenge_methods_supported": ["S256"], + "token_endpoint_auth_methods_supported": ["client_secret_post"], + "authorization_response_iss_parameter_supported": True, + } + ).encode() + ) + return Reply(status=404, body=json.dumps({"error": "not_found"}).encode()) + + with wire_server(respond) as wire, mcp_peer() as peer, gateway.scenario() as scenario: + wire_holder.append(wire) + alias: Final = "rt5" + uuid.uuid4().hex[:10] + identity: Final = register_mcp( + scenario, + peer, + alias, + issuer=wire.url, + auth_type="oauth2", + oauth2_flow="authorization_code", + credentials={"client_id": "ac-client", "client_secret": "ac-secret", "scopes": ["tools.call"]}, + ) + key: Final = scenario.key(object_permission={"mcp_servers": [identity]}) + + relay, cookie, cookie_value = _start_callback_flow( + gateway, alias, key, "ac-client", expected_authorization_url=wire.url + "/authorize" + ) + code: Final = "callback-code-" + uuid.uuid4().hex + missing_issuer: Final = gateway.client.get( + "/callback", + params={"code": code, "state": relay}, + cookies={cookie: cookie_value}, + ) + assert missing_issuer.status_code == 400, missing_issuer.text + assert missing_issuer.headers["content-type"] == "text/html; charset=utf-8", missing_issuer.headers + assert missing_issuer.text == ( + "Error: invalid_issuer
" + "This authorization response came from a different identity provider than the one this " + "MCP server is configured to use.
" + "You can close this window and try again.
" + ), missing_issuer.text + assert "location" not in missing_issuer.headers, missing_issuer.headers + assert code not in missing_issuer.text, missing_issuer.text + _assert_cleared_oauth_state_cookie(missing_issuer, cookie) + + trusted_relay, trusted_cookie, trusted_cookie_value = _start_callback_flow( + gateway, alias, key, "ac-client", expected_authorization_url=wire.url + "/authorize" + ) + trusted_code: Final = "callback-code-" + uuid.uuid4().hex + trusted: Final = gateway.client.get( + "/callback", + params={"code": trusted_code, "state": trusted_relay, "iss": wire.url}, + cookies={trusted_cookie: trusted_cookie_value}, + ) + assert trusted.status_code == 302, trusted.text + trusted_location: Final = urlsplit(trusted.headers["location"]) + assert trusted_location.scheme + "://" + trusted_location.netloc + trusted_location.path == CLIENT_REDIRECT, ( + trusted.headers["location"] + ) + trusted_query: Final = parse_qs(trusted_location.query) + assert set(trusted_query) == {"code", "state"}, trusted.headers["location"] + assert trusted_query["state"] == ["client-state"], trusted.headers["location"] + assert len(trusted_query["code"]) == 1 and trusted_query["code"][0].startswith("llm_scope_"), trusted.headers[ + "location" + ] + assert trusted_query["code"][0] != trusted_code, trusted.headers["location"] + _assert_cleared_oauth_state_cookie(trusted, trusted_cookie) + + wire_requests: Final = wire.drain() + wire_method_paths: Final = tuple((request.method, urlsplit(request.target).path) for request in wire_requests) + assert wire_method_paths == (("GET", "/.well-known/oauth-authorization-server"),), wire_method_paths + assert not any(method == "POST" and path == "/token" for method, path in wire_method_paths), wire_method_paths @pytest.mark.covers("other.mcp.oauth.discovery_cannot_erase_configured_authorization_endpoint") diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_count_tokens_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_count_tokens_wire.py new file mode 100644 index 00000000000..92512a304bc --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_count_tokens_wire.py @@ -0,0 +1,240 @@ +import json +from collections.abc import Callable +from typing import Final + +import anthropic +import pytest +from anthropic.types import MessageTokensCount +from integration._support.client import Gateway +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "claude-sonnet-4-6" +_API_KEY: Final = "synthetic-anthropic-key" +_COUNT_BETA: Final = "token-counting-2024-11-01" +_CLIENT_BETA: Final = "client-future-beta-2099-01-01" +_MERGED_BETA: Final = f"{_CLIENT_BETA},{_COUNT_BETA}" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this synthetic request."}] +_SYSTEM: Final[list[JsonValue]] = [{"type": "text", "text": "Count the supplied message and tool."}] +_TOOLS: Final[list[JsonValue]] = [ + { + "name": "lookup", + "description": "Look up a synthetic record.", + "input_schema": { + "type": "object", + "properties": {"record_id": {"type": "string"}}, + "required": ["record_id"], + }, + } +] +_COUNTED_INPUT_TOKENS: Final = 73 +_UPSTREAM_ERROR: Final = { + "type": "error", + "error": {"type": "invalid_request_error", "message": "The scripted count request was rejected."}, +} + + +def _client(gateway: Gateway, key: str) -> anthropic.Anthropic: + return anthropic.Anthropic( + api_key=key, + base_url=str(gateway.client.base_url), + max_retries=0, + http_client=gateway.client, + ) + + +def _count_peer( + expected_body: dict[str, JsonValue], result: Reply, expected_beta: str = _COUNT_BETA +) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages/count_tokens", request.target + assert request.headers["x-api-key"] == _API_KEY, request.headers + assert request.headers["anthropic-version"] == "2023-06-01", request.headers + assert request.headers["anthropic-beta"] == expected_beta, request.headers + assert _JSON_OBJECT.validate_json(request.body) == expected_body, request.body + return result + + return respond + + +def _sdk_count(client: anthropic.Anthropic, model: str, spelling: str, beta: str) -> MessageTokensCount: + if spelling == "beta": + beta_result: Final = client.beta.messages.count_tokens( + model=model, + messages=_MESSAGES, + system=_SYSTEM, + tools=_TOOLS, + betas=[beta], + ) + return MessageTokensCount.model_validate(beta_result.model_dump(exclude_none=True)) + return client.messages.count_tokens( + model=model, + messages=_MESSAGES, + system=_SYSTEM, + tools=_TOOLS, + extra_headers={"anthropic-beta": beta}, + ) + + +@pytest.mark.parametrize("spelling", ["messages", "beta"], ids=["messages_count_tokens", "beta_messages_count_tokens"]) +def test_anthropic_count_tokens_forwards_body_and_refuses_a_key_without_model_access( + gateway: Gateway, spelling: str +) -> None: + with gateway.scenario() as scenario: + with wire_server( + _count_peer( + {"model": _BACKEND, "messages": _MESSAGES, "system": _SYSTEM, "tools": _TOOLS}, + Reply(body=json.dumps({"input_tokens": _COUNTED_INPUT_TOKENS}).encode()), + ) + ) as wire: + model: Final = scenario.model( + model=f"anthropic/{_BACKEND}", + api_base=wire.url, + api_key=_API_KEY, + num_retries=0, + ) + restricted_model: Final = scenario.model( + model="anthropic/claude-haiku-4-5-20251001", + api_base=wire.url, + api_key=_API_KEY, + num_retries=0, + ) + restricted_key: Final = scenario.key(models=[restricted_model]) + + client: Final = _client(gateway, gateway.key) + result: Final = _sdk_count(client, model, spelling, _COUNT_BETA) + assert result.model_dump() == {"input_tokens": _COUNTED_INPUT_TOKENS} + requests: Final = wire.drain() + assert len(requests) == 1, requests + + restricted_client: Final = _client(gateway, restricted_key) + with pytest.raises(anthropic.PermissionDeniedError) as raised: + _sdk_count(restricted_client, model, spelling, _COUNT_BETA) + assert raised.value.status_code == 403, raised.value.response.text + assert model in raised.value.message, raised.value.message + assert wire.drain() == () + + +def test_anthropic_count_tokens_relays_an_upstream_server_error(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + with wire_server( + _count_peer( + {"model": _BACKEND, "messages": _MESSAGES, "system": _SYSTEM, "tools": _TOOLS}, + Reply(status=500, body=json.dumps(_UPSTREAM_ERROR).encode()), + ) + ) as wire: + model: Final = scenario.model( + model=f"anthropic/{_BACKEND}", + api_base=wire.url, + api_key=_API_KEY, + num_retries=0, + ) + client: Final = _client(gateway, gateway.key) + with pytest.raises(anthropic.InternalServerError) as raised: + client.messages.count_tokens( + model=model, + messages=_MESSAGES, + system=_SYSTEM, + tools=_TOOLS, + extra_headers={"anthropic-beta": _COUNT_BETA}, + ) + assert raised.value.status_code == 500 + assert "The scripted count request was rejected." in raised.value.message + assert len(wire.drain()) == 1 + + +def test_anthropic_count_tokens_relays_an_upstream_client_error(gateway: Gateway) -> None: + with gateway.scenario() as scenario: + with wire_server( + _count_peer( + {"model": _BACKEND, "messages": _MESSAGES, "system": _SYSTEM, "tools": _TOOLS}, + Reply(status=400, body=json.dumps(_UPSTREAM_ERROR).encode()), + ) + ) as wire: + model: Final = scenario.model( + model=f"anthropic/{_BACKEND}", + api_base=wire.url, + api_key=_API_KEY, + num_retries=0, + ) + client: Final = _client(gateway, gateway.key) + with pytest.raises(anthropic.BadRequestError) as raised: + client.messages.count_tokens( + model=model, + messages=_MESSAGES, + system=_SYSTEM, + tools=_TOOLS, + extra_headers={"anthropic-beta": _COUNT_BETA}, + ) + assert raised.value.status_code == 400 + assert "The scripted count request was rejected." in raised.value.message + assert len(wire.drain()) == 1 + + +def test_anthropic_count_tokens_preserves_thinking_tool_choice_and_output_config(gateway: Gateway) -> None: + pytest.skip( + "BUG: /v1/messages/count_tokens drops thinking, tool_choice and output_config before the provider count call" + ) + + thinking: Final = {"type": "enabled", "budget_tokens": 1024} + tool_choice: Final = {"type": "auto", "disable_parallel_tool_use": True} + output_config: Final = {"effort": "medium"} + expected_body: Final = { + "model": _BACKEND, + "messages": _MESSAGES, + "system": _SYSTEM, + "tools": _TOOLS, + "thinking": thinking, + "tool_choice": tool_choice, + "output_config": output_config, + } + + with gateway.scenario() as scenario: + with wire_server( + _count_peer(expected_body, Reply(body=json.dumps({"input_tokens": _COUNTED_INPUT_TOKENS}).encode())) + ) as wire: + model: Final = scenario.model( + model=f"anthropic/{_BACKEND}", + api_base=wire.url, + api_key=_API_KEY, + num_retries=0, + ) + client: Final = _client(gateway, gateway.key) + result: Final = client.messages.count_tokens( + model=model, + messages=_MESSAGES, + system=_SYSTEM, + tools=_TOOLS, + thinking=thinking, + tool_choice=tool_choice, + output_config=output_config, + extra_headers={"anthropic-beta": _COUNT_BETA}, + ) + assert result.input_tokens == _COUNTED_INPUT_TOKENS + assert len(wire.drain()) == 1 + + +@pytest.mark.parametrize("spelling", ["messages", "beta"], ids=["messages_count_tokens", "beta_messages_count_tokens"]) +def test_anthropic_count_tokens_merges_the_client_beta_with_the_token_counting_beta( + gateway: Gateway, spelling: str +) -> None: + pytest.skip("BUG: /v1/messages/count_tokens drops the client's anthropic-beta before the provider count call") + + with gateway.scenario() as scenario: + with wire_server( + _count_peer( + {"model": _BACKEND, "messages": _MESSAGES, "system": _SYSTEM, "tools": _TOOLS}, + Reply(body=json.dumps({"input_tokens": _COUNTED_INPUT_TOKENS}).encode()), + _MERGED_BETA, + ) + ) as wire: + model: Final = scenario.model( + model=f"anthropic/{_BACKEND}", + api_base=wire.url, + api_key=_API_KEY, + num_retries=0, + ) + result: Final = _sdk_count(_client(gateway, gateway.key), model, spelling, _CLIENT_BETA) + assert result.model_dump() == {"input_tokens": _COUNTED_INPUT_TOKENS} + assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py new file mode 100644 index 00000000000..8a79d5c6a7a --- /dev/null +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_gateway_wire.py @@ -0,0 +1,359 @@ +import hashlib +import uuid +from collections.abc import Iterator +from pathlib import Path +from typing import Final + +import pytest +import yaml +from anthropic.types import Message +from integration._support import claude_code as cc +from integration._support.client import Gateway, eventually, gateway_from_environment +from integration._support.database import read_rows +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_BACKEND: Final = "claude-sonnet-4-6" +_API_KEY: Final = "synthetic-anthropic-key" +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MESSAGE: Final = TypeAdapter(Message) + + +def _expected_stream_events( + stream_reply: tuple[bytes, ...], model: str +) -> tuple[tuple[str, dict[str, JsonValue]], ...]: + return tuple( + ( + event, + { + **data, + **( + { + "message": { + **_JSON_OBJECT.validate_python(data["message"]), + "model": model, + } + } + if event == "message_start" + else {} + ), + }, + ) + for event, data in cc.sse_events(b"".join(stream_reply).decode()) + ) + + +def _owned_config(path: Path) -> Path: + config: Final = _JSON_OBJECT.validate_python( + yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + ) + general_settings: Final = _JSON_OBJECT.validate_python(config.get("general_settings", {})) + litellm_settings: Final = _JSON_OBJECT.validate_python(config.get("litellm_settings", {})) + path.write_text( + yaml.safe_dump( + { + **config, + "general_settings": {**general_settings, "enable_claude_code_gateway": True}, + "litellm_settings": {**litellm_settings, "cache": False}, + } + ) + ) + return path + + +def _cli_bearer_headers(key: str) -> dict[str, str]: + return { + **{name: value for name, value in cc.cli_headers(key).items() if name != "x-api-key"}, + "authorization": f"Bearer {key}", + } + + +@pytest.fixture(scope="module") +def enabled_gateway(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Gateway]: + with gateway_from_environment() as base_gateway: + directory: Final = tmp_path_factory.mktemp("anthropic-gateway") + config: Final = _owned_config(directory / "gateway.yaml") + with owned_proxy_process(base_gateway, directory, {}, config=config, workers=2) as owned: + yield owned.gateway + + +@pytest.mark.parametrize( + "path", + ["/claude_code_gateway/v1/messages?beta=true", "/claude_code_gateway/v1/messages/count_tokens?beta=true"], + ids=["messages", "count_tokens"], +) +def test_shared_gateway_route_is_disabled_without_reaching_provider(gateway: Gateway, path: str) -> None: + def respond(request: Request) -> Reply: + raise AssertionError(f"disabled gateway must not reach the provider: {request.target}") + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=f"anthropic/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + response: Final = gateway.client.post( + path, + json={"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": "disabled control"}]}, + headers=_cli_bearer_headers(gateway.key), + ) + assert response.status_code == 404, response.text + assert _JSON_OBJECT.validate_json(response.content) == {"detail": "Claude Code gateway is not enabled"}, ( + response.text + ) + assert wire.drain() == () + + +def _restricted_model_responses(enabled_gateway: Gateway) -> tuple[int, str, int, str, str, tuple[Request, ...]]: + def respond(request: Request) -> Reply: + raise AssertionError(f"restricted model must not reach the provider: {request.target}") + + with enabled_gateway.scenario() as scenario, wire_server(respond) as wire: + model_a: Final = scenario.model(model=f"anthropic/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + model_b: Final = scenario.model(model=f"anthropic/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + key: Final = scenario.key(models=[model_a]) + request_body: Final = {**cc.claude_code_request("restricted-model"), "model": model_b} + gateway_response: Final = enabled_gateway.client.post( + "/claude_code_gateway/v1/messages?beta=true", + json=request_body, + headers=_cli_bearer_headers(key), + ) + messages_response: Final = enabled_gateway.client.post( + "/v1/messages?beta=true", + json=request_body, + headers=cc.cli_headers(key), + ) + return ( + gateway_response.status_code, + gateway_response.text, + messages_response.status_code, + messages_response.text, + model_b, + wire.drain(), + ) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +@pytest.mark.timeout(300) +def test_enabled_gateway_matches_native_messages_route(enabled_gateway: Gateway, stream: bool) -> None: + identity: Final = "msg_gateway_" + ("stream" if stream else "json") + usage: Final = {"input_tokens": 6, "output_tokens": 1} + content: Final = ({"type": "text", "text": "PONG"},) + request_body: Final = { + **cc.claude_code_request(f"cache-bust-{identity}"), + "model": "", + "stream": stream, + } + expected_stream: Final = cc.message_stream(identity, _BACKEND, content, usage) + expected_provider_body: Final = {**request_body, "model": _BACKEND} + client_betas: Final = sorted(cc.CLI_BETA.split(",")) + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + assert request.headers["x-api-key"] == _API_KEY, request.headers + assert request.headers["anthropic-version"] == "2023-06-01", request.headers + assert sorted(request.headers["anthropic-beta"].split(",")) == client_betas, request.headers + assert "authorization" not in request.headers, request.headers + assert _JSON_OBJECT.validate_json(request.body) == expected_provider_body, request.body + if stream: + return Reply(chunks=expected_stream, content_type="text/event-stream") + return Reply(body=cc.message_reply(identity, _BACKEND, content, usage)) + + with enabled_gateway.scenario() as scenario, wire_server(respond) as wire: + model_a: Final = scenario.model(model=f"anthropic/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + key: Final = scenario.key(models=[model_a]) + expected_message: Final = _MESSAGE.validate_python( + {**_JSON_OBJECT.validate_json(cc.message_reply(identity, _BACKEND, content, usage)), "model": model_a} + ) + expected_stream_events: Final = _expected_stream_events(expected_stream, model_a) + body: Final = {**request_body, "model": model_a} + headers: Final = _cli_bearer_headers(key) + gateway_response: Final = enabled_gateway.client.post( + "/claude_code_gateway/v1/messages?beta=true", + json=body, + headers=headers, + ) + assert gateway_response.status_code == 200, gateway_response.text + gateway_requests: Final = wire.drain() + assert len(gateway_requests) == 1, gateway_requests + + messages_response: Final = enabled_gateway.client.post( + "/v1/messages?beta=true", + json=body, + headers=headers, + ) + assert messages_response.status_code == 200, messages_response.text + messages_requests: Final = wire.drain() + assert len(messages_requests) == 1, messages_requests + if stream: + gateway_events: Final = cc.sse_events(gateway_response.text) + messages_events: Final = cc.sse_events(messages_response.text) + assert gateway_events == messages_events, (gateway_response.text, messages_response.text) + assert gateway_events == expected_stream_events, (gateway_response.text, messages_response.text) + else: + gateway_message: Final = _MESSAGE.validate_json(gateway_response.content) + messages_message: Final = _MESSAGE.validate_json(messages_response.content) + assert gateway_message.model_dump(mode="json") == expected_message.model_dump(mode="json"), ( + gateway_response.text, + ) + assert messages_message.model_dump(mode="json") == expected_message.model_dump(mode="json"), ( + messages_response.text, + ) + + (gateway_request,) = gateway_requests + (messages_request,) = messages_requests + assert gateway_request.method == messages_request.method == "POST" + assert gateway_request.target == messages_request.target == "/v1/messages" + assert gateway_request.headers["x-api-key"] == messages_request.headers["x-api-key"] == _API_KEY + assert ( + gateway_request.headers["anthropic-version"] + == messages_request.headers["anthropic-version"] + == "2023-06-01" + ) + assert gateway_request.headers["anthropic-beta"] == messages_request.headers["anthropic-beta"] + assert sorted(gateway_request.headers["anthropic-beta"].split(",")) == client_betas + assert _JSON_OBJECT.validate_json(gateway_request.body) == expected_provider_body, gateway_request.body + assert _JSON_OBJECT.validate_json(messages_request.body) == expected_provider_body, messages_request.body + + +@pytest.mark.timeout(300) +def test_enabled_gateway_count_tokens_alias_matches_native_count_tokens_route(enabled_gateway: Gateway) -> None: + messages: Final = [{"role": "user", "content": "Count this gateway request."}] + system: Final = [{"type": "text", "text": "Count the gateway message and tool."}] + tools: Final = [ + { + "name": "lookup", + "description": "Look up a synthetic record.", + "input_schema": {"type": "object", "properties": {"record_id": {"type": "string"}}}, + } + ] + expected_provider_body: Final = {"model": _BACKEND, "messages": messages, "system": system, "tools": tools} + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages/count_tokens", request.target + assert request.headers["x-api-key"] == _API_KEY, request.headers + assert request.headers["anthropic-version"] == "2023-06-01", request.headers + assert request.headers["anthropic-beta"] == "token-counting-2024-11-01", request.headers + assert "authorization" not in request.headers, request.headers + assert _JSON_OBJECT.validate_json(request.body) == expected_provider_body, request.body + return Reply(body=b'{"input_tokens": 41}') + + with enabled_gateway.scenario() as scenario, wire_server(respond) as wire: + model: Final = scenario.model(model=f"anthropic/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + restricted_model: Final = scenario.model(model=f"anthropic/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + key: Final = scenario.key(models=[model]) + body: Final = {"model": model, "messages": messages, "system": system, "tools": tools} + headers: Final = {**_cli_bearer_headers(key), "anthropic-beta": "token-counting-2024-11-01"} + gateway_response: Final = enabled_gateway.client.post( + "/claude_code_gateway/v1/messages/count_tokens?beta=true", json=body, headers=headers + ) + assert gateway_response.status_code == 200, gateway_response.text + assert _JSON_OBJECT.validate_json(gateway_response.content) == {"input_tokens": 41}, gateway_response.text + native_response: Final = enabled_gateway.client.post( + "/v1/messages/count_tokens?beta=true", json=body, headers=headers + ) + assert native_response.status_code == 200, native_response.text + assert native_response.content == gateway_response.content, (native_response.text, gateway_response.text) + assert tuple((request.method, request.target) for request in wire.drain()) == ( + ("POST", "/v1/messages/count_tokens"), + ("POST", "/v1/messages/count_tokens"), + ) + + restricted_response: Final = enabled_gateway.client.post( + "/claude_code_gateway/v1/messages/count_tokens?beta=true", + json={**body, "model": restricted_model}, + headers=headers, + ) + assert restricted_response.status_code == 403, restricted_response.text + assert _JSON_OBJECT.validate_json(restricted_response.content) == { + "error": { + "message": ( + f"The requested model '{restricted_model}' is not available for this API key, or the model name " + "is invalid. Check the models available to you and try again." + ), + "type": "key_model_access_denied", + "param": "model", + "code": "403", + } + }, restricted_response.text + assert wire.drain() == () + + +@pytest.mark.timeout(300) +def test_enabled_gateway_refuses_model_outside_virtual_key(enabled_gateway: Gateway) -> None: + result: Final = _restricted_model_responses(enabled_gateway) + expected_body: Final = { + "error": { + "message": ( + f"The requested model '{result[4]}' is not available for this API key, or the model name is invalid. " + "Check the models available to you and try again." + ), + "type": "key_model_access_denied", + "param": "model", + "code": "403", + } + } + assert result[0] == result[2] == 403, (result[1], result[3]) + assert _JSON_OBJECT.validate_json(result[1]) == expected_body, result[1] + assert _JSON_OBJECT.validate_json(result[3]) == expected_body, result[3] + assert result[1] == result[3], (result[1], result[3]) + assert result[5] == () + + +@pytest.mark.timeout(300) +def test_enabled_gateway_attributes_spend_to_virtual_key(enabled_gateway: Gateway) -> None: + gateway_identity: Final = f"msg_gateway_spend_gateway_{uuid.uuid4().hex}" + native_identity: Final = f"msg_gateway_spend_native_{uuid.uuid4().hex}" + response_ids: Final = iter((gateway_identity, native_identity)) + content: Final = ({"type": "text", "text": "PONG"},) + usage: Final = {"input_tokens": 6, "output_tokens": 1} + request_body: Final = {**cc.claude_code_request("cache-bust-msg-gateway-spend"), "model": "", "stream": False} + expected_provider_body: Final = {**request_body, "model": _BACKEND} + client_betas: Final = sorted(cc.CLI_BETA.split(",")) + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + assert request.headers["x-api-key"] == _API_KEY, request.headers + assert request.headers["anthropic-version"] == "2023-06-01", request.headers + assert sorted(request.headers["anthropic-beta"].split(",")) == client_betas, request.headers + assert "authorization" not in request.headers, request.headers + assert _JSON_OBJECT.validate_json(request.body) == expected_provider_body, request.body + return Reply(body=cc.message_reply(next(response_ids), _BACKEND, content, usage)) + + with enabled_gateway.scenario() as scenario, wire_server(respond) as wire: + model: Final = scenario.model(model=f"anthropic/{_BACKEND}", api_base=wire.url, api_key=_API_KEY) + key: Final = scenario.key(models=[model]) + native_key: Final = scenario.key(models=[model]) + body: Final = {**request_body, "model": model} + response: Final = enabled_gateway.client.post( + "/claude_code_gateway/v1/messages?beta=true", + json=body, + headers=_cli_bearer_headers(key), + ) + assert response.status_code == 200, response.text + native_response: Final = enabled_gateway.client.post( + "/v1/messages?beta=true", + json=body, + headers=cc.cli_headers(native_key), + ) + assert native_response.status_code == 200, native_response.text + requests: Final = wire.drain() + assert tuple(request.target for request in requests) == ("/v1/messages", "/v1/messages"), requests + gateway_key_hash: Final = hashlib.sha256(key.encode()).hexdigest() + native_key_hash: Final = hashlib.sha256(native_key.encode()).hexdigest() + spend_rows: Final = eventually( + lambda: ( + read_rows( + 'SELECT api_key, model_group FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (gateway_key_hash,), + ), + read_rows( + 'SELECT api_key, model_group FROM "LiteLLM_SpendLogs" WHERE api_key=%s', + (native_key_hash,), + ), + ), + lambda rows: len(rows[0]) == 1 and len(rows[1]) == 1, + seconds=70, + return_last_on_timeout=True, + ) + assert spend_rows == ( + [{"api_key": gateway_key_hash, "model_group": model}], + [{"api_key": native_key_hash, "model_group": model}], + ), spend_rows diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py index 7e9c5be227a..dcfd1d19ac2 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_anthropic_wire.py @@ -1,12 +1,20 @@ +import base64 import json import time import uuid from typing import Final +import anthropic import pytest +from anthropic.types import Message +from integration._support import claude_code as cc from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MESSAGE: Final = TypeAdapter(Message) @pytest.mark.covers( @@ -204,3 +212,563 @@ def test_anthropic_messages_slow_upstream_is_cut_off_at_the_deployment_request_t assert response.status_code == 408, response.text assert "Timeout" in response.json()["error"]["message"], response.text assert eventually(wire.drain, lambda requests: len(requests) == 1, seconds=5, return_last_on_timeout=True) + + +@pytest.mark.parametrize("stream", [False, True], ids=["non_stream", "stream"]) +def test_anthropic_messages_preserves_cli_beta_safeguards_and_result(gateway: Gateway, stream: bool) -> None: + identity: Final = "anthropic-safeguards-" + uuid.uuid4().hex + beta: Final = cc.CLI_BETA + ",made-up-future-beta-2099-01-01" + safeguards: Final = {"mode": "strict", "policies": ["sensitive-data"]} + safeguard_results: Final = {"status": "allowed", "policy": "sensitive-data"} + system: Final = [{"type": "text", "text": "Synthetic policy.", "cache_control": {"type": "ephemeral"}}] + tools: Final = [ + { + "name": "lookup", + "description": "Look up a synthetic record.", + "input_schema": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]}, + }, + { + "name": "update", + "description": "Update a synthetic record.", + "input_schema": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]}, + }, + ] + response_body: Final = { + "id": identity, + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "Record checked."}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 6, "output_tokens": 3}, + "safeguard_results": safeguard_results, + } + stream_reply: Final = ( + cc.sse_frame( + "message_start", + { + "type": "message_start", + "message": { + **response_body, + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 6, "output_tokens": 0}, + }, + }, + ), + cc.sse_frame( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + cc.sse_frame( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Record checked."}}, + ), + cc.sse_frame("content_block_stop", {"type": "content_block_stop", "index": 0}), + cc.sse_frame( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 3}, + }, + ), + cc.sse_frame("message_stop", {"type": "message_stop"}), + ) + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + assert request.headers["x-api-key"] == "synthetic-anthropic-key", request.headers + assert sorted(request.headers["anthropic-beta"].split(",")) == sorted(beta.split(",")), request.headers + expected_body: Final = { + "model": "claude-sonnet-4-6", + "max_tokens": 32, + "messages": [{"role": "user", "content": "Check the synthetic record."}], + "system": system, + "tools": tools, + "metadata": {"user_id": cc.METADATA_USER_ID}, + "safeguards": safeguards, + "stream": stream, + } + assert _JSON_OBJECT.validate_json(request.body) == expected_body, request.body + if not stream: + return Reply(body=json.dumps(response_body).encode()) + return Reply(chunks=stream_reply, content_type="text/event-stream") + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-6", + api_base=wire.url, + api_key="synthetic-anthropic-key", + ) + client_body: Final = { + "model": model, + "max_tokens": 32, + "messages": [{"role": "user", "content": "Check the synthetic record."}], + "system": system, + "tools": tools, + "metadata": {"user_id": cc.METADATA_USER_ID}, + "safeguards": safeguards, + "stream": stream, + } + response: Final = gateway.client.post( + "/v1/messages?beta=true", + json=client_body, + headers=cc.cli_headers(gateway.key, beta=beta), + ) + assert response.status_code == 200, response.text + (sent,) = wire.drain() + assert sent.target == "/v1/messages", sent.target + + if stream: + expected_stream_events: Final = tuple( + ( + event, + { + **data, + **( + { + "message": { + **_JSON_OBJECT.validate_python(data["message"]), + "model": model, + } + } + if event == "message_start" + else {} + ), + }, + ) + for event, data in cc.sse_events(b"".join(stream_reply).decode()) + ) + assert cc.sse_events(response.text) == expected_stream_events, response.text + + typed_message: Final = ( + _MESSAGE.validate_python( + next(data["message"] for event, data in cc.sse_events(response.text) if event == "message_start") + ) + if stream + else _MESSAGE.validate_json(response.content) + ) + assert typed_message.model_extra == {"safeguard_results": safeguard_results}, response.text + + +@pytest.mark.parametrize( + ("status", "error_type", "error_class", "litellm_error"), + [ + pytest.param(429, "rate_limit_error", anthropic.RateLimitError, "RateLimitError", id="rate_limit"), + pytest.param(400, "invalid_request_error", anthropic.BadRequestError, "BadRequestError", id="invalid_request"), + ], +) +def test_anthropic_messages_preserves_upstream_error_status_and_body( + gateway: Gateway, + status: int, + error_type: str, + error_class: type[anthropic.APIStatusError], + litellm_error: str, +) -> None: + message: Final = f"synthetic {error_type}" + expected_error: Final = {"type": "error", "error": {"type": error_type, "message": message}} + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + assert request.headers["x-api-key"] == "synthetic-anthropic-key", request.headers + assert _JSON_OBJECT.validate_json(request.body) == { + "model": "claude-sonnet-4-6", + "max_tokens": 16, + "messages": [{"role": "user", "content": "Trigger the scripted error."}], + "stream": False, + }, request.body + return Reply( + status=status, + body=json.dumps(expected_error).encode(), + headers={"retry-after": "7"} if status == 429 else {}, + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-6", + api_base=wire.url, + api_key="synthetic-anthropic-key", + ) + with ( + anthropic.Anthropic( + base_url=str(gateway.client.base_url), + api_key=gateway.key, + max_retries=0, + ) as client, + pytest.raises(error_class) as caught, + ): + client.messages.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": "Trigger the scripted error."}], + ) + assert caught.value.status_code == status, caught.value.response.text + assert caught.value.response.json() == { + "type": "error", + "error": { + "type": error_type, + "message": ( + f"litellm.{litellm_error}: AnthropicException - {json.dumps(expected_error)}\n\n" + f"LiteLLM: model group '{model}' failed with the error above. No fallback was attempted." + ), + }, + }, caught.value.response.text + assert len(wire.drain()) == 1 + + +def test_anthropic_messages_preserves_upstream_rate_limit_retry_after(gateway: Gateway) -> None: + pytest.skip("BUG: /v1/messages drops the upstream retry-after header on 429") + + message: Final = "synthetic rate_limit_error" + upstream_error: Final = { + "type": "error", + "error": {"type": "rate_limit_error", "message": message}, + } + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + assert request.headers["x-api-key"] == "synthetic-anthropic-key", request.headers + assert _JSON_OBJECT.validate_json(request.body) == { + "model": "claude-sonnet-4-6", + "max_tokens": 16, + "messages": [{"role": "user", "content": "Trigger the scripted error."}], + "stream": False, + }, request.body + return Reply(status=429, body=json.dumps(upstream_error).encode(), headers={"retry-after": "7"}) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-6", + api_base=wire.url, + api_key="synthetic-anthropic-key", + ) + with ( + anthropic.Anthropic( + base_url=str(gateway.client.base_url), + api_key=gateway.key, + max_retries=0, + ) as client, + pytest.raises(anthropic.RateLimitError) as caught, + ): + client.messages.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": "Trigger the scripted error."}], + ) + assert caught.value.status_code == 429, caught.value.response.text + assert caught.value.response.headers.get("retry-after") == "7", caught.value.response.text + assert caught.value.response.json() == { + "type": "error", + "error": { + "type": "rate_limit_error", + "message": ( + f"litellm.RateLimitError: AnthropicException - {json.dumps(upstream_error)}\n\n" + f"LiteLLM: model group '{model}' failed with the error above. No fallback was attempted." + ), + }, + }, caught.value.response.text + assert len(wire.drain()) == 1 + + +def test_anthropic_messages_preserves_upstream_overloaded_error(gateway: Gateway) -> None: + pytest.skip("BUG: /v1/messages maps an upstream 529 overloaded_error to HTTP 500") + + upstream_error: Final = { + "type": "error", + "error": {"type": "overloaded_error", "message": "synthetic overloaded error"}, + } + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + assert request.headers["x-api-key"] == "synthetic-anthropic-key", request.headers + assert _JSON_OBJECT.validate_json(request.body) == { + "model": "claude-sonnet-4-6", + "max_tokens": 16, + "messages": [{"role": "user", "content": "Trigger the scripted overload."}], + "stream": False, + }, request.body + return Reply(status=529, body=json.dumps(upstream_error).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-6", + api_base=wire.url, + api_key="synthetic-anthropic-key", + ) + with ( + anthropic.Anthropic( + base_url=str(gateway.client.base_url), + api_key=gateway.key, + max_retries=0, + ) as client, + pytest.raises(anthropic.InternalServerError) as caught, + ): + client.messages.create( + model=model, + max_tokens=16, + messages=[{"role": "user", "content": "Trigger the scripted overload."}], + ) + assert caught.value.status_code == 529, caught.value.response.text + response_body: Final = _JSON_OBJECT.validate_json(caught.value.response.content) + assert response_body == upstream_error, caught.value.response.text + assert len(wire.drain()) == 1 + + +def test_anthropic_messages_stream_surfaces_upstream_error_event(gateway: Gateway) -> None: + error: Final = { + "type": "error", + "error": {"type": "overloaded_error", "message": "synthetic stream overload"}, + } + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + assert request.headers["x-api-key"] == "synthetic-anthropic-key", request.headers + assert _JSON_OBJECT.validate_json(request.body) == { + "model": "claude-sonnet-4-6", + "max_tokens": 16, + "messages": [{"role": "user", "content": "Trigger the scripted stream error."}], + "stream": True, + }, request.body + return Reply( + chunks=( + cc.sse_frame( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_stream_error", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 0}, + }, + }, + ), + cc.sse_frame("error", error), + ), + content_type="text/event-stream", + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-6", + api_base=wire.url, + api_key="synthetic-anthropic-key", + ) + response: Final = gateway.client.post( + "/v1/messages", + json={ + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "Trigger the scripted stream error."}], + "stream": True, + }, + headers=cc.cli_headers(gateway.key), + ) + assert response.status_code == 200, response.text + assert cc.sse_events(response.text) == ( + ( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_stream_error", + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 0}, + }, + }, + ), + ("error", error), + ), response.text + assert len(wire.drain()) == 1 + + +def test_anthropic_messages_stream_ending_without_message_stop_surfaces_an_error_event(gateway: Gateway) -> None: + pytest.skip("BUG: /v1/messages relays an upstream stream that closes before message_stop as a clean 200") + message_start: Final = { + "type": "message_start", + "message": { + "id": "msg_stream_truncated", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 0}, + }, + } + block_start: Final = {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} + block_delta: Final = {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Partial"}} + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + assert request.headers["x-api-key"] == "synthetic-anthropic-key", request.headers + assert _JSON_OBJECT.validate_json(request.body) == { + "model": "claude-sonnet-4-6", + "max_tokens": 16, + "messages": [{"role": "user", "content": "Trigger the scripted truncated stream."}], + "stream": True, + }, request.body + return Reply( + chunks=( + cc.sse_frame("message_start", message_start), + cc.sse_frame("content_block_start", block_start), + cc.sse_frame("content_block_delta", block_delta), + ), + content_type="text/event-stream", + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-6", + api_base=wire.url, + api_key="synthetic-anthropic-key", + ) + response: Final = gateway.client.post( + "/v1/messages", + json={ + "model": model, + "max_tokens": 16, + "messages": [{"role": "user", "content": "Trigger the scripted truncated stream."}], + "stream": True, + }, + headers=cc.cli_headers(gateway.key), + ) + assert response.status_code == 200, response.text + assert cc.sse_events(response.text) == ( + ("message_start", {**message_start, "message": {**message_start["message"], "model": model}}), + ("content_block_start", block_start), + ("content_block_delta", block_delta), + ( + "error", + { + "type": "error", + "error": { + "type": "api_error", + "message": "Provider stream ended before emitting a message_stop event; the response is " + "incomplete and any partial content (e.g. tool_use input JSON) may be truncated.", + }, + }, + ), + ), response.text + assert len(wire.drain()) == 1 + + +def test_anthropic_messages_preserves_image_document_and_tool_result_blocks(gateway: Gateway) -> None: + image_data: Final = base64.b64encode(b"synthetic png bytes").decode() + pdf_data: Final = base64.b64encode(b"%PDF-1.7 synthetic document").decode() + messages: Final = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Inspect these attachments."}, + { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": image_data}, + }, + {"type": "image", "source": {"type": "url", "url": "https://example.invalid/synthetic.png"}}, + { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": pdf_data}, + }, + ], + }, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "lookup-one", "name": "lookup", "input": {"id": "one"}}], + }, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "lookup-one", + "is_error": True, + "content": [ + {"type": "text", "text": "Synthetic lookup failed."}, + { + "type": "image", + "source": {"type": "base64", "media_type": "image/png", "data": image_data}, + }, + ], + } + ], + }, + { + "role": "assistant", + "content": [{"type": "tool_use", "id": "lookup-two", "name": "lookup", "input": {"id": "two"}}], + }, + { + "role": "user", + "content": [ + {"type": "tool_result", "tool_use_id": "lookup-two", "content": "Synthetic lookup succeeded."}, + {"type": "text", "text": "Continue with the first result."}, + ], + }, + ] + content: Final = [ + {"type": "text", "text": "The requested blocks were preserved."}, + {"type": "tool_use", "id": "lookup-three", "name": "lookup", "input": {"id": "three"}}, + {"type": "thinking", "thinking": "The history contains two tool calls.", "signature": "synthetic-signature"}, + ] + expected_response: Final = { + "id": "msg_anthropic_blocks", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": content, + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 20, "output_tokens": 12}, + } + tools: Final = [ + { + "name": "lookup", + "description": "Look up a synthetic record.", + "input_schema": {"type": "object", "properties": {"id": {"type": "string"}}, "required": ["id"]}, + } + ] + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/v1/messages", request.target + assert request.headers["x-api-key"] == "synthetic-anthropic-key", request.headers + assert _JSON_OBJECT.validate_json(request.body) == { + "model": "claude-sonnet-4-6", + "max_tokens": 128, + "messages": messages, + "stream": False, + "tools": tools, + }, request.body + return Reply(body=json.dumps(expected_response).encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-6", + api_base=wire.url, + api_key="synthetic-anthropic-key", + ) + with anthropic.Anthropic( + base_url=str(gateway.client.base_url), + api_key=gateway.key, + max_retries=0, + ) as client: + response: Final = client.messages.create( + model=model, + max_tokens=128, + messages=messages, + tools=tools, + ) + expected_message: Final = _MESSAGE.validate_python({**expected_response, "model": model}) + assert response.model_dump(mode="json") == expected_message.model_dump(mode="json"), response + assert len(wire.drain()) == 1 diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py index 80ace4c4535..63325e7a005 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_chaos.py @@ -15,8 +15,6 @@ import pytest from integration._support.client import Gateway, eventually from integration._support.database import read_rows from integration._support.pdf_document import ( - COUNT_REFUSED, - COUNT_TOKENS_TARGET, LETTER, messages_body, pdf_document, @@ -42,7 +40,7 @@ class _Counted: pages: int status: int text: str - input_tokens: int | None + total_tokens: int | None @dataclass(frozen=True, slots=True) @@ -54,8 +52,6 @@ class _Sent: def _peer(request: Request) -> Reply: - if request.target == COUNT_TOKENS_TARGET: - return COUNT_REFUSED return anthropic_peer(request) @@ -79,31 +75,31 @@ def _count_body(pages: int) -> dict[str, JsonValue]: return {"model": _MODEL, "messages": messages_body(_MODEL, blocks, _ASK)["messages"]} -def _input_tokens(response: httpx.Response) -> int | None: +def _total_tokens(response: httpx.Response) -> int | None: if response.status_code != 200: return None - counted: Final = _JSON_OBJECT.validate_json(response.content).get("input_tokens") + counted: Final = _JSON_OBJECT.validate_json(response.content).get("total_tokens") return counted if isinstance(counted, int) else None -async def _fire_counts(url: str, key: str, *, tolerate_transport_errors: bool = False) -> tuple[_Counted, ...]: +async def _fire_counts(url: str, key: str) -> tuple[_Counted, ...]: async def one(client: httpx.AsyncClient, index: int) -> _Counted: pages: Final = _PAGE_COUNTS[index % len(_PAGE_COUNTS)] response: Final = await client.post( - "/v1/messages/count_tokens", json=_count_body(pages), headers={"Authorization": f"Bearer {key}"} + "/utils/token_counter", + params={"call_endpoint": "false"}, + json=_count_body(pages), + headers={"Authorization": f"Bearer {key}"}, ) - return _Counted(pages, response.status_code, response.text, _input_tokens(response)) + return _Counted(pages, response.status_code, response.text, _total_tokens(response)) async with httpx.AsyncClient(base_url=url, timeout=90, trust_env=False) as client: - results: Final = await asyncio.gather( - *(one(client, index) for index in range(_COUNT_BURST)), return_exceptions=tolerate_transport_errors - ) - for result in results: - assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) - return tuple(result for result in results if isinstance(result, _Counted)) + return tuple(await asyncio.gather(*(one(client, index) for index in range(_COUNT_BURST)))) -async def _fire_messages(url: str, key: str) -> tuple[_Sent, ...]: +async def _fire_messages( + url: str, key: str, *, tolerate_transport_errors: bool = False +) -> tuple[_Sent, ...]: async def one(client: httpx.AsyncClient) -> _Sent: marker: Final = uuid.uuid4().hex body: Final = messages_body(_MODEL, [pdf_document((LETTER,))], f"{_ASK} marker-{marker}") @@ -111,13 +107,20 @@ async def _fire_messages(url: str, key: str) -> tuple[_Sent, ...]: return _Sent(marker, response.status_code, response.text, response.headers.get("x-litellm-call-id", "")) async with httpx.AsyncClient(base_url=url, timeout=90, trust_env=False) as client: - return tuple(await asyncio.gather(*(one(client) for _ in range(_MESSAGE_BURST)))) + results: Final = await asyncio.gather( + *(one(client) for _ in range(_MESSAGE_BURST)), return_exceptions=tolerate_transport_errors + ) + for result in results: + assert not isinstance(result, BaseException) or isinstance(result, httpx.TransportError), repr(result) + return tuple(result for result in results if isinstance(result, _Sent)) def _baseline(gateway: Gateway) -> int: - response: Final = gateway.request("POST", "/v1/messages/count_tokens", _count_body(0)) + response: Final = gateway.request( + "POST", "/utils/token_counter", _count_body(0), params={"call_endpoint": "false"} + ) assert response.status_code == 200, response.text - counted: Final = _input_tokens(response) + counted: Final = _total_tokens(response) assert counted is not None, response.text return counted @@ -125,11 +128,11 @@ def _baseline(gateway: Gateway) -> int: def _assert_exact(counts: tuple[_Counted, ...], baseline: int) -> None: for item in counts: assert item.status == 200, (item.pages, item.status, item.text) - assert item.input_tokens is not None and item.input_tokens - baseline == rendered_tokens( + assert item.total_tokens is not None and item.total_tokens - baseline == rendered_tokens( (LETTER,) * item.pages ), ( item.pages, - item.input_tokens, + item.total_tokens, baseline, ) @@ -155,7 +158,7 @@ def _held_upstream_connections(pid: int, upstream: str) -> int: @pytest.mark.timeout(2 * graceful_stop_seconds() + 180) -async def test_pdf_count_burst_across_two_workers_prices_every_page_and_logs_each_message_once( +async def test_pdf_local_count_burst_across_two_workers_prices_every_page_and_logs_each_message_once( gateway: Gateway, tmp_path: Path ) -> None: with wire_server(_peer) as wire: @@ -174,8 +177,8 @@ async def test_pdf_count_burst_across_two_workers_prices_every_page_and_logs_eac ) received: Final = wire.drain() targets: Final = [request.target for request in received] - assert targets.count(COUNT_TOKENS_TARGET) == _COUNT_BURST, targets assert targets.count("/v1/messages") == _MESSAGE_BURST, targets + assert len(targets) == _MESSAGE_BURST, targets _assert_exact(counts, baseline) assert {item.marker for item in messages} == { _JSON_OBJECT.validate_json(item.text)["id"][len("msg_") :] for item in messages if item.status == 200 @@ -185,7 +188,7 @@ async def test_pdf_count_burst_across_two_workers_prices_every_page_and_logs_eac @pytest.mark.timeout(2 * graceful_stop_seconds() + 180) -async def test_worker_sigkill_mid_pdf_count_burst_leaves_the_sibling_pricing_pages( +async def test_worker_sigkill_mid_pdf_message_burst_preserves_local_page_pricing_and_single_logs( gateway: Gateway, tmp_path: Path ) -> None: release: Final = threading.Event() @@ -198,29 +201,38 @@ async def test_worker_sigkill_mid_pdf_count_burst_leaves_the_sibling_pricing_pag seconds=graceful_stop_seconds(), ) owned_url: Final = str(owned.gateway.client.base_url) - burst: Final = asyncio.create_task( - _fire_counts(owned_url, owned.gateway.key, tolerate_transport_errors=True) - ) - try: - await asyncio.to_thread( - eventually, lambda: wire.received.qsize(), lambda size: size >= _COUNT_BURST, 90 + async with asyncio.TaskGroup() as tasks: + burst = tasks.create_task( + _fire_messages(owned_url, owned.gateway.key, tolerate_transport_errors=True) ) - held_by: Final = MappingProxyType({pid: _held_upstream_connections(pid, wire.url) for pid in workers}) - victim: Final = max(workers, key=held_by.__getitem__) - with owned.replacing_worker(victim) as kill: - await asyncio.to_thread(kill) + try: + await asyncio.to_thread( + eventually, lambda: wire.received.qsize(), lambda size: size >= _MESSAGE_BURST, 90 + ) + held_by: Final = MappingProxyType( + {pid: _held_upstream_connections(pid, wire.url) for pid in workers} + ) + victim: Final = max(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + message_requests: Final = wire.drain() + finally: release.set() - served: Final = await burst - wire.drain() - finally: - release.set() baseline: Final = _baseline(owned.gateway) after: Final = await _fire_counts(owned_url, owned.gateway.key) after_received: Final = wire.drain() - assert sum(held_by.values()) == _COUNT_BURST, held_by + assert sum(held_by.values()) == _MESSAGE_BURST, held_by assert held_by[victim] > 0, held_by - assert len(served) == _COUNT_BURST - held_by[victim], (len(served), held_by) + assert len(message_requests) == _MESSAGE_BURST, message_requests + assert all(request.target == "/v1/messages" for request in message_requests), message_requests + assert len(served) == _MESSAGE_BURST - held_by[victim], (len(served), held_by) assert len(after) == _COUNT_BURST, len(after) - assert [request.target for request in after_received] == [COUNT_TOKENS_TARGET] * (_COUNT_BURST + 1) - _assert_exact(served, baseline) + assert after_received == (), after_received + assert {item.marker for item in served} == { + _JSON_OBJECT.validate_json(item.text)["id"][len("msg_") :] for item in served if item.status == 200 + } + for item in served: + _single_spend_row(item) _assert_exact(after, baseline) diff --git a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py index 41d5f3d8685..8eeaab66553 100644 --- a/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py +++ b/tests/integration/messages_endpoint/providers/anthropic/test_pdf_document_count_tokens_wire.py @@ -22,8 +22,8 @@ from integration._support.pdf_document import ( pdf_bytes, pdf_data_url, pdf_document, - rendered_tokens, responses_input_file, + rendered_tokens, text_document, ) from integration._support.provider import SharedProvider @@ -70,31 +70,33 @@ def _turn(blocks: Sequence[JsonValue]) -> list[JsonValue]: return messages_body(_MODEL, blocks, _ASK)["messages"] -def _count(gateway: Gateway, provider: SharedProvider, blocks: Sequence[JsonValue]) -> int: - provider.expect(COUNT_REFUSED) - response: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": _MODEL, "messages": _turn(blocks)}) - assert response.status_code == 200, response.text - assert [(request.method, request.target) for request in provider.received()] == [("POST", COUNT_TOKENS_TARGET)] - return _int(_JSON_OBJECT.validate_json(response.content)["input_tokens"]) - - -def _local(gateway: Gateway, parts: Sequence[JsonValue]) -> int: +def _local(gateway: Gateway, provider: SharedProvider, parts: Sequence[JsonValue]) -> int: response: Final = gateway.request( "POST", "/utils/token_counter", {"model": _MODEL, "messages": _turn(parts)}, params={"call_endpoint": "false"} ) assert response.status_code == 200, response.text + assert provider.received() == () return _int(_JSON_OBJECT.validate_json(response.content)["total_tokens"]) -def _responses_count(gateway: Gateway, provider: SharedProvider, items: Sequence[JsonValue]) -> int: - provider.expect(COUNT_REFUSED) +def _responses_count(gateway: Gateway, model: str, pages: Sequence[Page]) -> int: response: Final = gateway.request( "POST", "/v1/responses/input_tokens", - {"model": _MODEL, "input": [{"role": "user", "content": [*items, {"type": "input_text", "text": _ASK}]}]}, + { + "model": model, + "input": [ + { + "role": "user", + "content": [ + responses_input_file(pdf_data_url(pages)), + {"type": "input_text", "text": _ASK}, + ], + } + ], + }, ) assert response.status_code == 200, response.text - assert [(request.method, request.target) for request in provider.received()] == [("POST", COUNT_TOKENS_TARGET)] payload: Final = _JSON_OBJECT.validate_json(response.content) assert payload["object"] == "response.input_tokens", response.text return _int(payload["input_tokens"]) @@ -144,108 +146,123 @@ def _prompt_tokens(call_id: str) -> int: @pytest.mark.parametrize("pages", [1, 3, 12]) -def test_count_tokens_fallback_prices_each_blank_page_as_a_rendered_image( +def test_utils_token_counter_prices_each_blank_page_as_a_rendered_image( gateway: Gateway, provider: SharedProvider, pages: int ) -> None: letter_pages: Final = (LETTER,) * pages - with_document: Final = _count(gateway, provider, [pdf_document(letter_pages)]) - without: Final = _count(gateway, provider, []) + with_document: Final = _local(gateway, provider, [pdf_document(letter_pages)]) + without: Final = _local(gateway, provider, []) assert with_document - without == rendered_tokens(letter_pages), (with_document, without) -def test_count_tokens_fallback_adds_a_page_text_to_its_rendered_image( +def test_utils_token_counter_adds_a_page_text_to_its_rendered_image( gateway: Gateway, provider: SharedProvider ) -> None: - text_page: Final = _count(gateway, provider, [pdf_document((Page(text=_TEXT),))]) - blank_page: Final = _count(gateway, provider, [pdf_document((LETTER,))]) - text_source: Final = _count(gateway, provider, [text_document(_TEXT)]) - without: Final = _count(gateway, provider, []) + text_page: Final = _local(gateway, provider, [pdf_document((Page(text=_TEXT),))]) + blank_page: Final = _local(gateway, provider, [pdf_document((LETTER,))]) + text_source: Final = _local(gateway, provider, [text_document(_TEXT)]) + without: Final = _local(gateway, provider, []) assert text_source - without > 0, (text_source, without) assert text_page - blank_page == text_source - without, (text_page, blank_page, text_source, without) -def test_count_tokens_fallback_prices_a_page_by_its_rendered_area(gateway: Gateway, provider: SharedProvider) -> None: +def test_utils_token_counter_prices_a_page_by_its_rendered_area( + gateway: Gateway, provider: SharedProvider +) -> None: mixed: Final = (LETTER, NARROW, LETTER) - without: Final = _count(gateway, provider, []) - narrow: Final = _count(gateway, provider, [pdf_document((NARROW,))]) + without: Final = _local(gateway, provider, []) + narrow: Final = _local(gateway, provider, [pdf_document((NARROW,))]) assert narrow - without == rendered_tokens((NARROW,)), (narrow, without) - assert _count(gateway, provider, [pdf_document(mixed)]) - without == rendered_tokens(mixed) + assert _local(gateway, provider, [pdf_document(mixed)]) - without == rendered_tokens(mixed) -def test_count_tokens_fallback_prices_every_document_in_the_turn(gateway: Gateway, provider: SharedProvider) -> None: - both: Final = _count(gateway, provider, [pdf_document((LETTER,) * 3), pdf_document((LETTER,))]) - without: Final = _count(gateway, provider, []) +def test_utils_token_counter_prices_every_document_in_the_turn(gateway: Gateway, provider: SharedProvider) -> None: + both: Final = _local(gateway, provider, [pdf_document((LETTER,) * 3), pdf_document((LETTER,))]) + without: Final = _local(gateway, provider, []) assert both - without == rendered_tokens((LETTER,) * 4), (both, without) -def test_count_tokens_fallback_prices_a_pdf_without_pages_as_nothing( - gateway: Gateway, provider: SharedProvider -) -> None: - assert _count(gateway, provider, [pdf_document(())]) == _count(gateway, provider, []) +def test_utils_token_counter_prices_a_pdf_without_pages_as_nothing(gateway: Gateway, provider: SharedProvider) -> None: + assert _local(gateway, provider, [pdf_document(())]) == _local(gateway, provider, []) -def test_count_tokens_fallback_adds_title_and_context_once_per_document( +def test_utils_token_counter_adds_title_and_context_once_per_document( gateway: Gateway, provider: SharedProvider ) -> None: - without: Final = _count(gateway, provider, []) - plain: Final = _count(gateway, provider, [pdf_document((LETTER,))]) - annotated: Final = _count(gateway, provider, [pdf_document((LETTER,), title=_TITLE, context=_CONTEXT)]) - title: Final = _count(gateway, provider, [text_document(_TITLE)]) - without - context: Final = _count(gateway, provider, [text_document(_CONTEXT)]) - without + without: Final = _local(gateway, provider, []) + plain: Final = _local(gateway, provider, [pdf_document((LETTER,))]) + annotated: Final = _local(gateway, provider, [pdf_document((LETTER,), title=_TITLE, context=_CONTEXT)]) + title: Final = _local(gateway, provider, [text_document(_TITLE)]) - without + context: Final = _local(gateway, provider, [text_document(_CONTEXT)]) - without assert title > 0 and context > 0, (title, context) assert annotated - plain == title + context, (annotated, plain, title, context) @pytest.mark.parametrize("label", list(_UNREADABLE)) -def test_count_tokens_fallback_prices_unreadable_pdf_bytes_like_one_image( +def test_utils_token_counter_prices_unreadable_pdf_bytes_like_one_image( gateway: Gateway, provider: SharedProvider, label: str ) -> None: data: Final = _UNREADABLE[label] - as_pdf: Final = _count(gateway, provider, [document(base64_source(data))]) - as_png: Final = _count(gateway, provider, [document(base64_source(data, media_type="image/png"))]) - without: Final = _count(gateway, provider, []) + as_pdf: Final = _local(gateway, provider, [document(base64_source(data))]) + as_png: Final = _local(gateway, provider, [document(base64_source(data, media_type="image/png"))]) + without: Final = _local(gateway, provider, []) assert as_pdf == as_png, (as_pdf, as_png) assert 0 <= as_pdf - without < rendered_tokens((LETTER,)), (as_pdf, without) -def test_count_tokens_fallback_reads_a_list_wrapped_base64_string_like_the_bare_string( +def test_utils_token_counter_reads_a_list_wrapped_base64_string_like_the_bare_string( gateway: Gateway, provider: SharedProvider ) -> None: raw: Final = encoded(pdf_bytes((LETTER,))) - wrapped: Final = _count(gateway, provider, [document(base64_source([raw]))]) - assert wrapped == _count(gateway, provider, [document(base64_source(raw))]), wrapped + wrapped: Final = _local(gateway, provider, [document(base64_source([raw]))]) + assert wrapped == _local(gateway, provider, [document(base64_source(raw))]), wrapped -def test_count_tokens_fallback_reads_an_owner_locked_pdf(gateway: Gateway, provider: SharedProvider) -> None: +def test_utils_token_counter_reads_an_owner_locked_pdf(gateway: Gateway, provider: SharedProvider) -> None: pages: Final = (LETTER, LETTER) locked: Final = document(base64_source(encoded(_locked(pdf_bytes(pages), "")))) - assert _count(gateway, provider, [locked]) - _count(gateway, provider, []) == rendered_tokens(pages) + assert _local(gateway, provider, [locked]) - _local(gateway, provider, []) == rendered_tokens(pages) -def test_count_tokens_fallback_parses_only_the_pdf_media_type(gateway: Gateway, provider: SharedProvider) -> None: +def test_utils_token_counter_parses_only_the_pdf_media_type(gateway: Gateway, provider: SharedProvider) -> None: raw: Final = encoded(pdf_bytes((LETTER,))) - without: Final = _count(gateway, provider, []) - labelled: Final = _count(gateway, provider, [document(base64_source(raw))]) + without: Final = _local(gateway, provider, []) + labelled: Final = _local(gateway, provider, [document(base64_source(raw))]) assert labelled - without == rendered_tokens((LETTER,)), (labelled, without) - mislabelled: Final = _count(gateway, provider, [document(base64_source(raw, media_type="application/x-pdf"))]) - assert mislabelled == _count(gateway, provider, [document(base64_source(raw, media_type="image/png"))]) + mislabelled: Final = _local(gateway, provider, [document(base64_source(raw, media_type="application/x-pdf"))]) + assert mislabelled == _local(gateway, provider, [document(base64_source(raw, media_type="image/png"))]) -def test_utils_token_counter_prices_a_chat_file_by_its_pages(gateway: Gateway) -> None: - one: Final = _local(gateway, [chat_file(pdf_data_url((LETTER,)))]) - three: Final = _local(gateway, [chat_file(pdf_data_url((LETTER,) * 3))]) +def test_messages_count_tokens_surfaces_provider_refusal_without_returning_tokens( + gateway: Gateway, provider: SharedProvider +) -> None: + provider.expect(COUNT_REFUSED) + response: Final = gateway.request( + "POST", + "/v1/messages/count_tokens", + {"model": _MODEL, "messages": _turn([pdf_document((LETTER,))])}, + ) + assert response.status_code == 400, response.text + assert "input_tokens" not in response.text, response.text + assert [(request.method, request.target) for request in provider.received()] == [("POST", COUNT_TOKENS_TARGET)] + + +def test_utils_token_counter_prices_a_chat_file_by_its_pages(gateway: Gateway, provider: SharedProvider) -> None: + one: Final = _local(gateway, provider, [chat_file(pdf_data_url((LETTER,)))]) + three: Final = _local(gateway, provider, [chat_file(pdf_data_url((LETTER,) * 3))]) assert three - one == rendered_tokens((LETTER,) * 2), (three, one) - text_page: Final = _local(gateway, [chat_file(pdf_data_url((Page(text=_TEXT),)))]) - text_part: Final = _local(gateway, [chat_file(pdf_data_url((LETTER,))), {"type": "text", "text": _TEXT}]) + text_page: Final = _local(gateway, provider, [chat_file(pdf_data_url((Page(text=_TEXT),)))]) + text_part: Final = _local(gateway, provider, [chat_file(pdf_data_url((LETTER,))), {"type": "text", "text": _TEXT}]) assert text_part - one > 0, (text_part, one) assert text_page - one == text_part - one, (text_page, text_part, one) -def test_responses_input_tokens_fallback_prices_an_input_file_by_its_pages( - gateway: Gateway, provider: SharedProvider -) -> None: - one: Final = _responses_count(gateway, provider, [responses_input_file(pdf_data_url((LETTER,)))]) - three: Final = _responses_count(gateway, provider, [responses_input_file(pdf_data_url((LETTER,) * 3))]) +def test_responses_input_tokens_prices_an_input_file_by_its_pages(gateway: Gateway, provider: SharedProvider) -> None: + with gateway.scenario() as scenario: + model: Final = scenario.model(model="xai/grok-4.20", api_base=f"{provider.wire.url}/v1") + one: Final = _responses_count(gateway, model, (LETTER,)) + three: Final = _responses_count(gateway, model, (LETTER,) * 3) + assert provider.received() == () assert three - one == rendered_tokens((LETTER,) * 2), (three, one) diff --git a/tests/integration/messages_endpoint/safeguards/conftest.py b/tests/integration/messages_endpoint/safeguards/conftest.py new file mode 100644 index 00000000000..22b3ca5e24d --- /dev/null +++ b/tests/integration/messages_endpoint/safeguards/conftest.py @@ -0,0 +1,12 @@ +from __future__ import annotations + +from collections.abc import Iterator + +import pytest +from integration._support.safeguards_rig import Rig, safeguards_rig + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + with safeguards_rig(tmp_path_factory.mktemp("safeguards")) as value: + yield value diff --git a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_bridges.py b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_bridges.py new file mode 100644 index 00000000000..9ecf65056b6 --- /dev/null +++ b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_bridges.py @@ -0,0 +1,244 @@ +"""Auto-mode safeguards on the chat and responses bridges behind /v1/messages: the proxy classifies with the +configured model, stamps ``safeguard_results`` on the reply, and keeps the request's safeguards away from the backend +exactly as the bridges did before.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Mapping +from typing import Final + +import anthropic +import pytest +from integration._support.process import graceful_stop_seconds +from integration._support.safeguards_rig import ( + BASH_TOOL, + CHAT_MODEL, + DANGEROUS_TOOL_USE_BETA, + RESPONSES_MODEL, + SAFEGUARDS, + Responder, + Rig, + auto_mode_body, + available, + chat_backend, + classifier_prompt, + classifier_rows, + flagged, + ids_under_review, + json_body, + key_hash, + messages_body, + nonce, + one_classifier_call, + request_json, + requested_max_tokens, + responses_backend, + spend_rows_for, + stream_results, + stream_tool_use_ids, + tool_id, + tool_use_ids, + unavailable, +) +from pydantic import JsonValue + +pytestmark = pytest.mark.timeout(graceful_stop_seconds() + 240) + +BRIDGES: Final = (CHAT_MODEL, RESPONSES_MODEL) + + +def backend(model: str, identity: str, *, stream: bool, text: str | None = None) -> Responder: + if model == CHAT_MODEL: + return chat_backend(identity, stream=stream, text=text, finish_reason="stop" if text else "tool_calls") + return responses_backend(identity, stream=stream, text=text) + + +def sdk_client(rig: Rig, key: str | None = None) -> anthropic.Anthropic: + return anthropic.Anthropic( + base_url=str(rig.gateway.client.base_url), + api_key=key or rig.gateway.key, + default_headers={"anthropic-beta": DANGEROUS_TOOL_USE_BETA}, + max_retries=0, + ) + + +def sdk_async_client(rig: Rig) -> anthropic.AsyncAnthropic: + return anthropic.AsyncAnthropic( + base_url=str(rig.gateway.client.base_url), + api_key=rig.gateway.key, + default_headers={"anthropic-beta": DANGEROUS_TOOL_USE_BETA}, + max_retries=0, + ) + + +def assert_classified_once(rig: Rig, request_nonce: str, ids: tuple[str, ...]) -> None: + classifier_call, backend_calls = one_classifier_call(rig, request_nonce) + assert len(backend_calls) == 1, [request.target for request in backend_calls] + assert ids_under_review(classifier_call) == ids, classifier_call.body + classifier_body: Final = request_json(classifier_call) + assert classifier_body["model"] == "gpt-5.4-mini", classifier_body + assert requested_max_tokens(classifier_body) == 4096, classifier_body + prompt: Final = classifier_prompt(classifier_call) + assert "never-forwarded" not in prompt and "remote_url" not in prompt, prompt + assert "/work/repo" in prompt and "permission_mode" in prompt, prompt + + +@pytest.mark.parametrize("stream", (False, True), ids=("non_stream", "stream")) +@pytest.mark.parametrize("model", BRIDGES) +def test_bridge_backend_body_carries_no_safeguards(rig: Rig, model: str, stream: bool) -> None: + request_nonce: Final = nonce() + rig.register(request_nonce, backend(model, tool_id(request_nonce), stream=stream)) + body: Final = auto_mode_body(model, request_nonce) + status: Final = rig.stream(body)[0] if stream else rig.post(body).status_code + assert status == 200, status + backend_calls: Final = rig.backend_calls(rig.received(request_nonce)) + assert len(backend_calls) == 1, [request.target for request in backend_calls] + assert "safeguards" not in request_json(backend_calls[0]), backend_calls[0].body + + +@pytest.mark.parametrize("model", BRIDGES) +def test_non_stream_reply_carries_the_classifier_verdict(rig: Rig, model: str) -> None: + request_nonce: Final = nonce() + rig.register(request_nonce, backend(model, tool_id(request_nonce), stream=False)) + body: Final = json_body(rig.post(auto_mode_body(model, request_nonce))) + ids: Final = tool_use_ids(body) + assert len(ids) == 1, body + assert body["safeguard_results"] == flagged(ids), body + assert_classified_once(rig, request_nonce, ids) + + +@pytest.mark.parametrize("model", BRIDGES) +def test_stream_final_delta_carries_the_classifier_verdict(rig: Rig, model: str) -> None: + request_nonce: Final = nonce() + rig.register(request_nonce, backend(model, tool_id(request_nonce), stream=True)) + status, raw = rig.stream(auto_mode_body(model, request_nonce)) + assert status == 200, raw + ids: Final = stream_tool_use_ids(raw) + assert len(ids) == 1, raw + assert stream_results(raw) == flagged(ids), raw + assert_classified_once(rig, request_nonce, ids) + + +@pytest.mark.parametrize("model", BRIDGES) +def test_anthropic_sdk_sync_client_reads_the_verdict(rig: Rig, model: str) -> None: + request_nonce: Final = nonce() + rig.register(request_nonce, backend(model, tool_id(request_nonce), stream=False)) + raw: Final = sdk_client(rig).messages.with_raw_response.create( + model=model, + max_tokens=256, + messages=[{"role": "user", "content": f"clean the build directory {request_nonce}"}], + tools=[dict(BASH_TOOL)], # pyright: ignore[reportArgumentType] # the SDK's ToolParam is a TypedDict of this shape + extra_body={"safeguards": list(SAFEGUARDS)}, + ) + parsed: Final = raw.parse().model_dump(mode="json") + ids: Final = tool_use_ids(parsed) + assert len(ids) == 1, parsed + assert parsed["safeguard_results"] == flagged(ids), parsed + assert raw.http_response.json()["safeguard_results"] == flagged(ids), raw.http_response.text + assert_classified_once(rig, request_nonce, ids) + + +@pytest.mark.parametrize("model", BRIDGES) +def test_anthropic_sdk_async_stream_reads_the_verdict(rig: Rig, model: str) -> None: + request_nonce: Final = nonce() + rig.register(request_nonce, backend(model, tool_id(request_nonce), stream=True)) + + async def consume() -> tuple[tuple[str, ...], list[JsonValue]]: + async with sdk_async_client(rig).messages.stream( + model=model, + max_tokens=256, + messages=[{"role": "user", "content": f"clean the build directory {request_nonce}"}], + tools=[dict(BASH_TOOL)], # pyright: ignore[reportArgumentType] # the SDK's ToolParam is a TypedDict of this shape + extra_body={"safeguards": list(SAFEGUARDS)}, + ) as stream: + events: Final = [event.model_dump(mode="json") async for event in stream] + starts: Final = [event["content_block"] for event in events if event["type"] == "content_block_start"] + ids: Final = tuple(block["id"] for block in starts if block["type"] == "tool_use") + deltas: Final = [ + event["delta"].get("safeguard_results") for event in events if event["type"] == "message_delta" + ] + return ids, deltas + + ids, deltas = asyncio.run(consume()) + assert len(ids) == 1, ids + assert deltas == [flagged(ids)], deltas + assert_classified_once(rig, request_nonce, ids) + + +def test_classifier_spend_lands_on_the_calling_key(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key() + request_nonce: Final = nonce() + rig.register(request_nonce, backend(CHAT_MODEL, tool_id(request_nonce), stream=False)) + body: Final = json_body(rig.post(auto_mode_body(CHAT_MODEL, request_nonce), key=key)) + assert body["safeguard_results"] == flagged(tool_use_ids(body)), body + rows: Final = spend_rows_for(key, 2) + assert len(rows) == 2, rows + assert {str(row["api_key"]) for row in rows} == {key_hash(key)}, rows + classifier: Final = classifier_rows(rows) + assert len(classifier) == 1, rows + assert float(str(classifier[0]["spend"])) > 0, classifier + assert all(float(str(row["spend"])) > 0 for row in rows), rows + + +@pytest.mark.parametrize("model", BRIDGES) +def test_text_only_reply_is_available_with_no_verdicts_and_no_classifier_call(rig: Rig, model: str) -> None: + request_nonce: Final = nonce() + rig.register(request_nonce, backend(model, tool_id(request_nonce), stream=False, text="all clean")) + body: Final = json_body(rig.post(auto_mode_body(model, request_nonce))) + assert tool_use_ids(body) == (), body + assert body["safeguard_results"] == available({}), body + received: Final = rig.received(request_nonce) + assert rig.classifier_calls(received) == (), [request.target for request in received] + assert len(rig.backend_calls(received)) == 1, [request.target for request in received] + + +def test_truncated_chat_tool_call_is_unavailable_without_a_classifier_call(rig: Rig) -> None: + request_nonce: Final = nonce() + identity: Final = tool_id(request_nonce) + rig.register(request_nonce, chat_backend(identity, stream=False, finish_reason="length")) + body: Final = json_body(rig.post(auto_mode_body(CHAT_MODEL, request_nonce))) + ids: Final = tool_use_ids(body) + assert body["stop_reason"] == "max_tokens", body + assert body["safeguard_results"] == unavailable(ids, "truncated"), body + assert rig.classifier_calls(rig.received(request_nonce)) == () + + +def test_truncated_responses_tool_call_is_unavailable_without_a_classifier_call(rig: Rig) -> None: + request_nonce: Final = nonce() + rig.register(request_nonce, responses_backend(tool_id(request_nonce), stream=False, incomplete=True)) + body: Final = json_body(rig.post(auto_mode_body(RESPONSES_MODEL, request_nonce))) + ids: Final = tool_use_ids(body) + assert body["stop_reason"] == "max_tokens", body + assert body["safeguard_results"] == unavailable(ids, "truncated"), body + assert rig.classifier_calls(rig.received(request_nonce)) == () + + +@pytest.mark.parametrize("model", BRIDGES) +def test_identical_twin_carries_the_verdict_its_cached_classifier_call_returned(rig: Rig, model: str) -> None: + request_nonce: Final = nonce() + rig.register(request_nonce, backend(model, tool_id(request_nonce), stream=False)) + request: Final[Mapping[str, JsonValue]] = auto_mode_body(model, request_nonce) + first: Final = json_body(rig.post(request)) + second: Final = json_body(rig.post(request)) + ids: Final = tool_use_ids(first) + assert tool_use_ids(second) == ids, second + assert first["safeguard_results"] == flagged(ids) == second["safeguard_results"], (first, second) + received: Final = rig.received(request_nonce) + assert len(rig.classifier_calls(received)) == 1, [request.target for request in received] + assert len(rig.backend_calls(received)) == 1, [request.target for request in received] + + +def test_chat_bridge_cache_hit_after_a_request_without_safeguards_is_still_classified(rig: Rig) -> None: + request_nonce: Final = nonce() + rig.register(request_nonce, backend(CHAT_MODEL, tool_id(request_nonce), stream=False)) + plain: Final = json_body(rig.post(messages_body(CHAT_MODEL, request_nonce), beta=None)) + assert "safeguard_results" not in plain, plain + with_safeguards: Final = json_body(rig.post(auto_mode_body(CHAT_MODEL, request_nonce))) + ids: Final = tool_use_ids(with_safeguards) + assert ids == tool_use_ids(plain), (plain, with_safeguards) + assert with_safeguards["safeguard_results"] == flagged(ids), with_safeguards + received: Final = rig.received(request_nonce) + assert len(rig.classifier_calls(received)) == 1, [request.target for request in received] + assert len(rig.backend_calls(received)) == 1, [request.target for request in received] diff --git a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_caller_checks.py b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_caller_checks.py new file mode 100644 index 00000000000..6b8ad1d3360 --- /dev/null +++ b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_caller_checks.py @@ -0,0 +1,352 @@ +"""The classifier call runs as the caller: a key, team, member, user, end user, agent, or budget that cannot use the +classifier model turns the whole safeguard into ``unsupported`` instead of charging or bypassing the caller's limits, +the router keeps the call on the deployments the caller's access groups grant, and a limit that only holds for this +window (a full RPM or TPM window, or counters Redis cannot read) makes this turn's tool uses ``unavailable`` and leaves +the next turn on the server check, while the slot the turn itself took never denies its classifier call. A native call +that sends its own ``metadata`` is charged to the same per-model budgets as the classifier call.""" + +from __future__ import annotations + +import uuid +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import Scenario, eventually, object_value, string_value +from integration._support.process import graceful_stop_seconds +from integration._support.redis_process import owned_redis +from integration._support.safeguards_rig import ( + CHAT_MODEL, + CLASSIFIER_MODEL, + NATIVE_MODEL, + UNSUPPORTED, + Rig, + auto_mode_body, + chat_backend, + classifier_answering, + deployment, + flagged, + json_body, + messages_body, + native_backend, + native_message, + nonce, + safeguards_rig, + text_block, + tool_id, + tool_use_ids, + unavailable, +) +from integration._support.wire import Request +from pydantic import JsonValue + +pytestmark = pytest.mark.timeout(graceful_stop_seconds() + 240) + +TINY_BUDGET: Final = {CLASSIFIER_MODEL: {"max_budget": 1e-7, "budget_duration": "1d"}} +TINY_NATIVE_BUDGET: Final = {NATIVE_MODEL: {"max_budget": 1e-7, "budget_duration": "1d"}} + + +def classify(rig: Rig, key: str) -> tuple[dict[str, JsonValue], str]: + request_nonce: Final = nonce() + rig.register(request_nonce, chat_backend(tool_id(request_nonce), stream=False)) + return json_body(rig.post(auto_mode_body(CHAT_MODEL, request_nonce), key=key)), request_nonce + + +def assert_unsupported_without_a_classifier_call(rig: Rig, key: str) -> None: + body, request_nonce = classify(rig, key) + assert body["safeguard_results"] == list(UNSUPPORTED), body + assert len(tool_use_ids(body)) == 1, body + received: Final = rig.received(request_nonce) + assert rig.classifier_calls(received) == (), [request.target for request in received] + assert len(rig.backend_calls(received)) == 1, [request.target for request in received] + + +def test_key_without_the_classifier_model_gets_unsupported(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + assert_unsupported_without_a_classifier_call(rig, scenario.key(models=[CHAT_MODEL])) + + +def test_team_without_the_classifier_model_gets_unsupported(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + team_id: Final = scenario.team(models=[CHAT_MODEL]) + assert_unsupported_without_a_classifier_call(rig, scenario.key(team_id=team_id)) + + +def assert_unavailable_for_the_turn_without_a_classifier_call(rig: Rig, key: str) -> None: + body, request_nonce = classify(rig, key) + ids: Final = tool_use_ids(body) + assert len(ids) == 1, body + assert body["safeguard_results"] == unavailable(ids, "error"), body + received: Final = rig.received(request_nonce) + assert rig.classifier_calls(received) == (), [request.target for request in received] + assert len(rig.backend_calls(received)) == 1, [request.target for request in received] + + +def spend_one_classifier_request(rig: Rig, key: str) -> None: + request_nonce: Final = nonce() + rig.register(request_nonce, chat_backend(tool_id(request_nonce), stream=False), classifier_answering("pong")) + reply: Final = rig.gateway.chat(CLASSIFIER_MODEL, key=key, text=f"spend the window {request_nonce}") + assert "pong" in str(reply), reply + assert len(rig.classifier_calls(rig.received(request_nonce))) == 1 + + +def test_classifier_model_rpm_window_the_caller_spent_is_unavailable_for_the_turn(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_rpm_limit={CLASSIFIER_MODEL: 1}) + spend_one_classifier_request(rig, key) + assert_unavailable_for_the_turn_without_a_classifier_call(rig, key) + + +def test_classifier_call_takes_no_slot_of_the_callers_classifier_model_rpm_limit(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_rpm_limit={CLASSIFIER_MODEL: 1}) + first, _ = classify(rig, key) + assert first["safeguard_results"] == flagged(tool_use_ids(first)), first + second, _ = classify(rig, key) + assert second["safeguard_results"] == flagged(tool_use_ids(second)), second + + +def assert_flagged_after_one_classifier_call(rig: Rig, key: str) -> None: + body, request_nonce = classify(rig, key) + assert body["safeguard_results"] == flagged(tool_use_ids(body)), body + received: Final = rig.received(request_nonce) + assert len(rig.classifier_calls(received)) == 1, [request.target for request in received] + + +def test_key_rpm_limit_that_admitted_the_turn_leaves_its_classifier_call_running(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + assert_flagged_after_one_classifier_call(rig, scenario.key(rpm_limit=1)) + + +def test_user_rpm_limit_that_admitted_the_turn_leaves_its_classifier_call_running(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + user_id: Final = scenario.user(user_role="internal_user", rpm_limit=1) + assert_flagged_after_one_classifier_call(rig, scenario.key(user_id=user_id)) + + +def test_team_rpm_limit_that_admitted_the_turn_leaves_its_classifier_call_running(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + team_id: Final = scenario.team(rpm_limit=1) + assert_flagged_after_one_classifier_call(rig, scenario.key(team_id=team_id)) + + +def test_classifier_model_tpm_window_spent_by_earlier_turns_is_unavailable_for_the_turn(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_tpm_limit={CLASSIFIER_MODEL: 100}) + first, _ = classify(rig, key) + assert first["safeguard_results"] == flagged(tool_use_ids(first)), first + spent: Final = eventually( + lambda: classify(rig, key)[0], + lambda body: body["safeguard_results"] == unavailable(tool_use_ids(body), "error"), + seconds=60, + ) + assert len(tool_use_ids(spent)) == 1, spent + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 240) +def test_unreadable_rate_limit_counters_make_the_turn_unavailable(tmp_path: Path) -> None: + with owned_redis(tmp_path / "redis") as redis: + + def redis_environment(wire_url: str) -> Mapping[str, str]: + return {"REDIS_HOST": redis.host, "REDIS_PORT": str(redis.port)} + + with safeguards_rig( + tmp_path, + environment=redis_environment, + general_settings={"fail_closed_rate_limit_enforcement": True}, + ) as rig: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_rpm_limit={CLASSIFIER_MODEL: 5}) + first, _ = classify(rig, key) + assert first["safeguard_results"] == flagged(tool_use_ids(first)), first + redis.stop() + assert_unavailable_for_the_turn_without_a_classifier_call(rig, key) + redis.start() + + +def test_key_model_budget_for_the_classifier_is_enforced(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(model_max_budget=TINY_BUDGET) + first, _ = classify(rig, key) + assert first["safeguard_results"] == flagged(tool_use_ids(first)), first + exhausted: Final = eventually( + lambda: classify(rig, key)[0], lambda body: body["safeguard_results"] == list(UNSUPPORTED), seconds=60 + ) + assert len(tool_use_ids(exhausted)) == 1, exhausted + + +def test_team_member_model_budget_for_the_classifier_is_enforced(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + team_id: Final = scenario.team(team_member_model_max_budget=TINY_BUDGET) + user_id: Final = scenario.member(team_id) + key: Final = scenario.key(team_id=team_id, user_id=user_id) + first, _ = classify(rig, key) + assert first["safeguard_results"] == flagged(tool_use_ids(first)), first + exhausted: Final = eventually( + lambda: classify(rig, key)[0], lambda body: body["safeguard_results"] == list(UNSUPPORTED), seconds=60 + ) + assert len(tool_use_ids(exhausted)) == 1, exhausted + + +def native_turn_with_client_metadata(rig: Rig, key: str) -> tuple[httpx.Response, str]: + request_nonce: Final = nonce() + message: Final = native_message(tool_id(request_nonce), content=[text_block("done")], stop_reason="end_turn") + rig.register(request_nonce, native_backend(message, stream=False)) + body: Final = messages_body(NATIVE_MODEL, request_nonce, metadata={"user_id": f"session-{request_nonce}"}) + return rig.post(body, key=key, beta=None), request_nonce + + +def test_native_call_sending_its_own_metadata_is_charged_to_the_key_model_budget(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(user_id=scenario.user(), model_max_budget=TINY_NATIVE_BUDGET) + first, first_nonce = native_turn_with_client_metadata(rig, key) + assert first.status_code == 200, first.text + assert len(rig.backend_calls(rig.received(first_nonce))) == 1 + exhausted, exhausted_nonce = eventually( + lambda: native_turn_with_client_metadata(rig, key), lambda turn: turn[0].status_code != 200, seconds=60 + ) + assert f"exceeded budget for model={NATIVE_MODEL}" in exhausted.text, exhausted.text + assert rig.backend_calls(rig.received(exhausted_nonce)) == () + + +def test_key_update_dropping_the_classifier_takes_effect_on_the_next_request(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(models=[CHAT_MODEL, CLASSIFIER_MODEL]) + first, _ = classify(rig, key) + assert first["safeguard_results"] == flagged(tool_use_ids(first)), first + rig.gateway.post("/key/update", {"key": key, "models": [CHAT_MODEL]}) + exhausted: Final = eventually( + lambda: classify(rig, key)[0], lambda body: body["safeguard_results"] == list(UNSUPPORTED), seconds=30 + ) + assert len(tool_use_ids(exhausted)) == 1, exhausted + rig.gateway.post("/key/update", {"key": key, "models": [CHAT_MODEL, CLASSIFIER_MODEL]}) + restored: Final = eventually( + lambda: classify(rig, key)[0], lambda body: body["safeguard_results"] != list(UNSUPPORTED), seconds=30 + ) + assert restored["safeguard_results"] == flagged(tool_use_ids(restored)), restored + + +def direct_classifier_status(rig: Rig, key: str, **extra: JsonValue) -> int: + body: Final = {"model": CLASSIFIER_MODEL, "messages": [{"role": "user", "content": f"direct {nonce()}"}], **extra} + return rig.gateway.request("POST", "/v1/chat/completions", body, key=key).status_code + + +def customer(scenario: Scenario, **fields: JsonValue) -> str: + identity: Final = f"integration-{uuid.uuid4().hex}" + scenario.gateway.post("/customer/new", {"user_id": identity, **fields}) + scenario.cleanups.callback(scenario.gateway.post, "/customer/delete", {"user_ids": [identity]}) + return identity + + +def test_end_user_without_the_classifier_model_gets_unsupported(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key() + end_user: Final = customer(scenario, models=[CHAT_MODEL]) + assert direct_classifier_status(rig, key, user=end_user) == 403 + request_nonce: Final = nonce() + rig.register(request_nonce, chat_backend(tool_id(request_nonce), stream=False)) + body: Final = json_body( + rig.post(auto_mode_body(CHAT_MODEL, request_nonce, metadata={"user_id": end_user}), key=key) + ) + assert body["safeguard_results"] == list(UNSUPPORTED), body + received: Final = rig.received(request_nonce) + assert rig.classifier_calls(received) == (), [request.target for request in received] + assert len(rig.backend_calls(received)) == 1, [request.target for request in received] + + +def agent_key(scenario: Scenario, models: Sequence[str]) -> str: + gateway: Final = scenario.gateway + group: Final = gateway.request( + "POST", + "/v1/access_group", + {"access_group_name": f"integration-{uuid.uuid4().hex}", "access_model_names": list(models)}, + ) + assert group.status_code == 201, group.text + group_id: Final = string_value(object_value(group.json())["access_group_id"]) + scenario.cleanups.callback(gateway.request, "DELETE", f"/v1/access_group/{group_id}") + agent: Final = gateway.post( + "/v1/agents", + { + "agent_name": f"integration-{uuid.uuid4().hex[:10]}", + "agent_card_params": { + "protocolVersion": "0.3.0", + "name": "integration", + "description": "integration agent", + "url": "http://127.0.0.1:1/agent", + "version": "1", + "capabilities": {}, + "defaultInputModes": ["text"], + "defaultOutputModes": ["text"], + "skills": [], + }, + "access_group_ids": [group_id], + }, + ) + agent_id: Final = string_value(agent["agent_id"]) + scenario.cleanups.callback(gateway.request, "DELETE", f"/v1/agents/{agent_id}") + return scenario.key(agent_id=agent_id) + + +def test_agent_whose_access_groups_omit_the_classifier_gets_unsupported(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + key: Final = agent_key(scenario, [CHAT_MODEL]) + assert direct_classifier_status(rig, key) == 403 + assert_unsupported_without_a_classifier_call(rig, key) + + +def test_personal_key_of_a_user_without_the_classifier_model_gets_unsupported(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + user_id: Final = scenario.user(user_role="internal_user", models=[CHAT_MODEL]) + key: Final = scenario.key(user_id=user_id) + assert direct_classifier_status(rig, key) == 403 + assert_unsupported_without_a_classifier_call(rig, key) + + +def test_team_key_reaches_the_classifier_its_team_grants_beyond_its_users_personal_models(rig: Rig) -> None: + with rig.gateway.scenario() as scenario: + team_id: Final = scenario.team(models=[CHAT_MODEL, CLASSIFIER_MODEL]) + user_id: Final = scenario.user(user_role="internal_user", models=[CHAT_MODEL]) + rig.gateway.post("/team/member_add", {"team_id": team_id, "member": {"role": "user", "user_id": user_id}}) + key: Final = scenario.key(team_id=team_id, user_id=user_id) + spend_one_classifier_request(rig, key) + assert_flagged_after_one_classifier_call(rig, key) + + +GROUPED_CLASSIFIER: Final = "safeguards-classifier-grouped" +GRANTED_GROUP: Final = "safeguards-group-granted" +OTHER_GROUP: Final = "safeguards-group-other" +GRANTED_TARGET: Final = "/classifier/granted/v1/chat/completions" + + +def grouped_classifiers(wire_url: str) -> Sequence[Mapping[str, JsonValue]]: + return ( + { + **deployment(GROUPED_CLASSIFIER, "openai/gpt-5.4-mini", f"{wire_url}/classifier/granted/v1", weight=1), + "model_info": {"access_groups": [GRANTED_GROUP]}, + }, + { + **deployment(GROUPED_CLASSIFIER, "openai/gpt-5.4-mini", f"{wire_url}/classifier/other/v1", weight=1000), + "model_info": {"access_groups": [OTHER_GROUP]}, + }, + ) + + +def classifier_targets(requests: Sequence[Request]) -> tuple[str, ...]: + return tuple(request.target for request in requests if request.target.startswith("/classifier/")) + + +@pytest.mark.timeout(graceful_stop_seconds() + 240) +def test_classifier_call_keeps_to_the_deployments_the_callers_access_group_grants(tmp_path: Path) -> None: + with safeguards_rig(tmp_path, classifier_model=GROUPED_CLASSIFIER, extra_models=grouped_classifiers) as rig: + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(models=[CHAT_MODEL, GRANTED_GROUP]) + direct_nonce: Final = nonce() + rig.register(direct_nonce, chat_backend(tool_id(direct_nonce), stream=False), classifier_answering("pong")) + assert "pong" in str(rig.gateway.chat(GROUPED_CLASSIFIER, key=key, text=f"direct {direct_nonce}")) + assert classifier_targets(rig.received(direct_nonce)) == (GRANTED_TARGET,) + for _ in range(3): + body, request_nonce = classify(rig, key) + assert body["safeguard_results"] == flagged(tool_use_ids(body)), body + assert classifier_targets(rig.received(request_nonce)) == (GRANTED_TARGET,) diff --git a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_chaos.py b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_chaos.py new file mode 100644 index 00000000000..0815b9c89de --- /dev/null +++ b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_chaos.py @@ -0,0 +1,185 @@ +"""Auto mode under disruption: a classifier outage in the middle of a mixed burst across the chat bridge, the +Responses bridge, and a native backend, streaming and not, concurrent classifier calls, the proxy stopped mid burst, +and a worker killed while the other keeps serving.""" + +from __future__ import annotations + +import threading +from collections.abc import Sequence +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from itertools import product +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support.client import eventually +from integration._support.process import graceful_stop_seconds, stop_root_process +from integration._support.safeguards_rig import ( + CHAT_MODEL, + NATIVE_MODEL, + RESPONSES_MODEL, + Responder, + Rig, + auto_mode_body, + chat_backend, + classifier_dropping_the_connection, + classifier_flagging_everything, + classifier_slow, + flagged, + json_body, + native_backend, + native_message, + nonce, + responses_backend, + safeguards_rig, + stream_results, + stream_tool_use_ids, + text_block, + tool_id, + tool_use_block, + tool_use_ids, + unavailable, +) +from integration._support.wire import Reply, Request +from pydantic import JsonValue + +pytestmark = pytest.mark.timeout(2 * graceful_stop_seconds() + 300) + +BURST: Final = 10 +MIX: Final = tuple(product((CHAT_MODEL, RESPONSES_MODEL, NATIVE_MODEL), (False, True))) * 4 + + +@dataclass(frozen=True, slots=True) +class Outcome: + request_nonce: str + ids: tuple[str, ...] + results: JsonValue + + +def scripted_backend(model: str, identity: str, stream: bool) -> Responder: + if model == CHAT_MODEL: + return chat_backend(identity, stream=stream) + if model == RESPONSES_MODEL: + return responses_backend(identity, stream=stream) + message: Final = native_message(identity, content=(text_block("Deleting."), tool_use_block(identity))) + return native_backend(message, stream=stream) + + +def send(rig: Rig, model: str, stream: bool, request_nonce: str) -> Outcome: + body: Final = auto_mode_body(model, request_nonce) + if stream: + status, raw = rig.stream(body) + assert status == 200, raw + return Outcome(request_nonce, stream_tool_use_ids(raw), stream_results(raw)) + reply: Final = json_body(rig.post(body)) + return Outcome(request_nonce, tool_use_ids(reply), reply.get("safeguard_results")) + + +def mixed_burst(rig: Rig, classifier: Responder) -> tuple[Outcome, ...]: + plan: Final = tuple((model, stream, nonce()) for model, stream in MIX) + for model, stream, request_nonce in plan: + rig.register(request_nonce, scripted_backend(model, tool_id(request_nonce), stream), classifier) + with ThreadPoolExecutor(max_workers=len(plan)) as pool: + return tuple(pool.map(lambda entry: send(rig, entry[0], entry[1], entry[2]), plan)) + + +def chat_burst(rig: Rig, classifier: Responder) -> tuple[Outcome, ...]: + nonces: Final = tuple(nonce() for _ in range(BURST)) + for request_nonce in nonces: + rig.register(request_nonce, chat_backend(tool_id(request_nonce), stream=False), classifier) + with ThreadPoolExecutor(max_workers=BURST) as pool: + return tuple(pool.map(lambda request_nonce: send(rig, CHAT_MODEL, False, request_nonce), nonces)) + + +def assert_landed_once(rig: Rig, received: Sequence[Request], request_nonce: str) -> None: + own: Final = tuple(request for request in received if request_nonce.encode() in request.body) + assert len(rig.backend_calls(own)) == 1, [request.target for request in own] + assert len(rig.classifier_calls(own)) == 1, [request.target for request in own] + + +def assert_every_id_landed_once(rig: Rig, outcomes: Sequence[Outcome]) -> None: + received: Final = rig.wire.drain() + for outcome in outcomes: + assert outcome.ids == (tool_id(outcome.request_nonce),), outcome + assert_landed_once(rig, received, outcome.request_nonce) + + +def classifier_released_together(parties: int) -> Responder: + barrier: Final = threading.Barrier(parties, timeout=20) + + def respond(request: Request) -> Reply: + try: + barrier.wait() + except threading.BrokenBarrierError: + return Reply(status=503, body=b'{"error": {"message": "classifier calls did not overlap"}}') + return classifier_flagging_everything(request) + + return respond + + +def test_classifier_outage_mid_burst_degrades_then_recovers(rig: Rig) -> None: + outage: Final = mixed_burst(rig, classifier_dropping_the_connection) + for outcome in outage: + assert outcome.results == unavailable(outcome.ids, "error"), outcome + assert_every_id_landed_once(rig, outage) + recovered: Final = mixed_burst(rig, classifier_flagging_everything) + for outcome in recovered: + assert outcome.results == flagged(outcome.ids), outcome + assert_every_id_landed_once(rig, recovered) + + +def test_classifier_calls_of_concurrent_requests_overlap(rig: Rig) -> None: + outcomes: Final = chat_burst(rig, classifier_released_together(BURST)) + for outcome in outcomes: + assert outcome.results == flagged(outcome.ids), outcome + assert_every_id_landed_once(rig, outcomes) + + +def test_proxy_stopped_mid_burst_fails_in_flight_requests_cleanly(tmp_path: Path) -> None: + with safeguards_rig(tmp_path) as rig: + nonces: Final = tuple(nonce() for _ in range(BURST)) + for request_nonce in nonces: + rig.register(request_nonce, chat_backend(tool_id(request_nonce), stream=False), classifier_slow(4)) + + def attempt(request_nonce: str) -> str: + try: + response: Final = rig.post(auto_mode_body(CHAT_MODEL, request_nonce)) + except httpx.HTTPError: + return "error" + return str(response.status_code) + + with ThreadPoolExecutor(max_workers=BURST) as pool: + futures: Final = tuple(pool.submit(attempt, request_nonce) for request_nonce in nonces) + eventually(lambda: rig.wire.received.qsize(), lambda count: count >= 2 * BURST, seconds=60) + assert stop_root_process(rig.owned.process) + outcomes: Final = tuple(future.result(timeout=graceful_stop_seconds() + 60) for future in futures) + assert set(outcomes) <= {"200", "500", "502", "503", "error"}, outcomes + received: Final = rig.wire.drain() + for request_nonce in nonces: + assert_landed_once(rig, received, request_nonce) + with safeguards_rig(tmp_path) as fresh: + request_nonce: Final = nonce() + fresh.register(request_nonce, chat_backend(tool_id(request_nonce), stream=False)) + body: Final = json_body(fresh.post(auto_mode_body(CHAT_MODEL, request_nonce))) + assert body["safeguard_results"] == flagged(tool_use_ids(body)), body + + +def classify_on(rig: Rig) -> Outcome: + request_nonce: Final = nonce() + rig.register(request_nonce, chat_backend(tool_id(request_nonce), stream=False)) + return send(rig, CHAT_MODEL, False, request_nonce) + + +def test_one_worker_killed_while_the_other_keeps_serving(tmp_path: Path) -> None: + with safeguards_rig(tmp_path) as rig: + victim: Final = min(rig.owned.worker_pids()) + with rig.owned.replacing_worker(victim) as kill: + kill() + during: Final = tuple(classify_on(rig) for _ in range(5)) + after: Final = classify_on(rig) + for outcome in (*during, after): + assert outcome.results == flagged(outcome.ids), outcome + assert_every_id_landed_once(rig, (*during, after)) + assert victim not in rig.owned.worker_pids() diff --git a/tests/integration/messages_endpoint/safeguards/test_safeguard_results_classifier_failures.py b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_classifier_failures.py new file mode 100644 index 00000000000..af8269cdcfc --- /dev/null +++ b/tests/integration/messages_endpoint/safeguards/test_safeguard_results_classifier_failures.py @@ -0,0 +1,277 @@ +"""A classifier that fails, stalls, refuses, or answers nonsense degrades every tool use to ``unavailable`` with the +reason Claude Code expects, and the classifier prompt carries only the context the feature documents.""" + +from __future__ import annotations + +import json +from collections.abc import Mapping, Sequence +from pathlib import Path +from typing import Final + +import pytest +from integration._support.client import eventually +from integration._support.process import graceful_stop_seconds +from integration._support.safeguards_rig import ( + CHAT_MODEL, + CLASSIFIER_CONTEXT, + CLASSIFIER_MODEL, + Responder, + Rig, + auto_mode_body, + chat_backend, + chat_completion_reply, + classifier_answering, + classifier_failing, + classifier_flagging_everything, + classifier_prompt, + classifier_slow, + deployment, + flagged, + ids_under_review, + json_body, + messages_body, + nonce, + one_classifier_call, + safeguards_rig, + tool_id, + tool_use_ids, + unavailable, +) +from integration._support.wire import Reply, Request +from pydantic import JsonValue + +pytestmark = pytest.mark.timeout(graceful_stop_seconds() + 240) + +BOUNDED_CLASSIFIER: Final = "safeguards-classifier-bounded" +RESERVED_CLASSIFIER: Final = "safeguards-classifier-reserved" +OMISSION_MARKER: Final = json.dumps({"meta": "earlier turns omitted"}) +TRANSCRIPT_BUDGET: Final = 24_000 + + +def verdicts_shaped(shape: str) -> Responder: + def respond(request: Request) -> Reply: + verdicts: Final = {identity: shape for identity in ids_under_review(request)} + return chat_completion_reply(json.dumps({"verdicts": verdicts})) + + return respond + + +FAILURES: Final[Mapping[str, tuple[Responder, str]]] = { + "server_error": (classifier_failing(500), "error"), + "context_length": ( + classifier_failing( + 400, + message="This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens.", + error_type="invalid_request_error", + code="context_length_exceeded", + ), + "input_too_long", + ), + "content_policy": ( + classifier_failing( + 400, + message="Invalid prompt: your prompt was flagged as potentially violating our usage policy. content_policy_violation", + error_type="invalid_request_error", + code="content_policy_violation", + ), + "refused", + ), + "unparsable": (classifier_answering("I cannot classify that"), "error"), + "missing_verdict": (classifier_answering(json.dumps({"verdicts": {}})), "error"), + "invalid_shape": (verdicts_shaped("flagged"), "error"), +} + + +def classify(rig: Rig, classifier: Responder, **extra: JsonValue) -> tuple[dict[str, JsonValue], str]: + request_nonce: Final = nonce() + rig.register(request_nonce, chat_backend(tool_id(request_nonce), stream=False), classifier) + return json_body(rig.post(auto_mode_body(CHAT_MODEL, request_nonce, **extra))), request_nonce + + +@pytest.mark.parametrize("failure", tuple(FAILURES)) +def test_classifier_failure_degrades_to_unavailable_with_its_reason(rig: Rig, failure: str) -> None: + classifier, reason = FAILURES[failure] + body, request_nonce = classify(rig, classifier) + ids: Final = tool_use_ids(body) + assert len(ids) == 1, body + assert body["safeguard_results"] == unavailable(ids, reason), body + classifier_call, backend_calls = one_classifier_call(rig, request_nonce) + assert ids_under_review(classifier_call) == ids + assert len(backend_calls) == 1 + + +def test_rejected_classifier_key_is_unavailable_with_one_warning_and_no_traceback(rig: Rig) -> None: + rejected_key: Final = classifier_failing( + 401, message="Incorrect API key provided", error_type="invalid_request_error", code="invalid_api_key" + ) + log_offset: Final = rig.owned.log.stat().st_size + body, request_nonce = classify(rig, rejected_key) + ids: Final = tool_use_ids(body) + assert len(ids) == 1, body + assert body["safeguard_results"] == unavailable(ids, "error"), body + one_classifier_call(rig, request_nonce) + named: Final = f"safeguards classifier {CLASSIFIER_MODEL}" + tail: Final = eventually( + lambda: rig.owned.log.read_bytes()[log_offset:].decode(errors="replace"), + lambda text: named in text, + seconds=30, + ) + lines: Final = tuple(line for line in tail.splitlines() if named in line) + assert len(lines) == 1, tail + assert "WARNING" in lines[0], tail + assert "Traceback" not in tail, tail + + +@pytest.mark.timeout(graceful_stop_seconds() + 360) +def test_classifier_stalled_past_the_deadline_is_a_timeout(rig: Rig) -> None: + body, request_nonce = classify(rig, classifier_slow(35)) + ids: Final = tool_use_ids(body) + assert body["safeguard_results"] == unavailable(ids, "timeout"), body + assert len(rig.classifier_calls(rig.received(request_nonce))) == 1 + + +def bounded_classifier(wire_url: str) -> Sequence[Mapping[str, JsonValue]]: + return (deployment(BOUNDED_CLASSIFIER, "openai/gpt-5.4-mini", f"{wire_url}/classifier/v1", timeout=2),) + + +def reserved_for_another_team(start: str, end: str) -> Mapping[str, JsonValue]: + return {"start": start, "end": end, "timezone": "UTC", "team_ids": ["safeguards-other-team"]} + + +def reserved_classifier(wire_url: str) -> Sequence[Mapping[str, JsonValue]]: + whole_day: Final = (reserved_for_another_team("00:00", "12:00"), reserved_for_another_team("12:00", "00:00")) + return ( + { + **deployment(RESERVED_CLASSIFIER, "openai/gpt-5.4-mini", f"{wire_url}/classifier/v1"), + "model_info": {"access_windows": [dict(window) for window in whole_day]}, + }, + ) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 240) +def test_classifier_deployment_timeout_bounds_the_classifier_call(tmp_path: Path) -> None: + with safeguards_rig(tmp_path, classifier_model=BOUNDED_CLASSIFIER, extra_models=bounded_classifier) as rig: + body, request_nonce = classify(rig, classifier_slow(10)) + ids: Final = tool_use_ids(body) + assert len(ids) == 1, body + assert body["safeguard_results"] == unavailable(ids, "timeout"), body + assert len(rig.classifier_calls(rig.received(request_nonce))) == 1 + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 240) +def test_router_refusal_other_than_no_healthy_deployments_is_unavailable(tmp_path: Path) -> None: + with safeguards_rig(tmp_path, classifier_model=RESERVED_CLASSIFIER, extra_models=reserved_classifier) as rig: + body, request_nonce = classify(rig, classifier_flagging_everything) + ids: Final = tool_use_ids(body) + assert len(ids) == 1, body + assert body["safeguard_results"] == unavailable(ids, "error"), body + received: Final = rig.received(request_nonce) + assert rig.classifier_calls(received) == (), [request.target for request in received] + assert len(rig.backend_calls(received)) == 1, [request.target for request in received] + + +def test_classifier_prompt_carries_only_the_documented_context(rig: Rig) -> None: + context: Final[Mapping[str, JsonValue]] = { + **CLASSIFIER_CONTEXT, + "rules": ["never push to main CTX-RULE-MARKER"], + "unknown_key": "CTX-UNKNOWN-MARKER", + "git_state": { + "visibility": "private CTX-VISIBILITY-MARKER", + "remote_url": "CTX-REMOTE-MARKER", + "token": "CTX-TOKEN-MARKER", + }, + } + safeguards: Final[list[JsonValue]] = [{"type": "dangerous_tool_use", "classifier_context": dict(context)}] + request_nonce: Final = nonce() + rig.register(request_nonce, chat_backend(tool_id(request_nonce), stream=False)) + body: Final = json_body(rig.post(messages_body(CHAT_MODEL, request_nonce, safeguards=safeguards))) + assert body["safeguard_results"] == flagged(tool_use_ids(body)), body + prompt: Final = classifier_prompt(one_classifier_call(rig, request_nonce)[0]) + assert "CTX-RULE-MARKER" in prompt and "CTX-VISIBILITY-MARKER" in prompt, prompt + for hidden in ("CTX-UNKNOWN-MARKER", "CTX-REMOTE-MARKER", "CTX-TOKEN-MARKER", "never-forwarded"): + assert hidden not in prompt, (hidden, prompt) + + +def test_long_transcript_keeps_the_ends_and_marks_the_omission(rig: Rig) -> None: + request_nonce: Final = nonce() + turns: Final[list[JsonValue]] = [ + {"role": "user" if index % 2 == 0 else "assistant", "content": f"TURN-{index:02d} " + "lorem ipsum " * 82} + for index in range(39) + ] + rig.register(request_nonce, chat_backend(tool_id(request_nonce), stream=False)) + request: Final = { + **auto_mode_body(CHAT_MODEL, request_nonce), + "messages": [*turns, {"role": "user", "content": f"TURN-39 clean up {request_nonce}"}], + } + body: Final = json_body(rig.post(request)) + assert body["safeguard_results"] == flagged(tool_use_ids(body)), body + prompt: Final = classifier_prompt(one_classifier_call(rig, request_nonce)[0]) + assert "earlier turns omitted" in prompt, prompt + assert "TURN-00" in prompt and "TURN-38" in prompt and "TURN-39" in prompt, prompt + assert "TURN-05" not in prompt, prompt + assert len(prompt) < 40_000, len(prompt) + + +def text_blocks(*texts: str) -> Sequence[JsonValue]: + return tuple({"type": "text", "text": text} for text in texts) + + +def later_turns(count: int, filler: int) -> Sequence[JsonValue]: + return tuple( + {"role": "assistant" if index % 2 == 1 else "user", "content": f"LATER-{index:02d} " + "x" * filler} + for index in range(1, count + 1) + ) + + +def transcript_sent_for(rig: Rig, opening: Sequence[JsonValue], later: Sequence[JsonValue]) -> str: + request_nonce: Final = nonce() + rig.register(request_nonce, chat_backend(tool_id(request_nonce), stream=False)) + request: Final = { + **auto_mode_body(CHAT_MODEL, request_nonce), + "messages": [ + {"role": "user", "content": list(opening)}, + *later, + {"role": "user", "content": f"NEWEST-TURN now clean up {request_nonce}"}, + ], + } + body: Final = json_body(rig.post(request)) + assert body["safeguard_results"] == flagged(tool_use_ids(body)), body + prompt: Final = classifier_prompt(one_classifier_call(rig, request_nonce)[0]) + transcript: Final = prompt.split("one JSON line per turn:\n", 1)[1].split("\n\nTool calls to classify:", 1)[0] + assert len(transcript) <= TRANSCRIPT_BUDGET, len(transcript) + return transcript + + +def test_long_session_keeps_the_whole_opening_message_and_the_newest_turns(rig: Rig) -> None: + opening: Final = text_blocks( + "AccessDeniedSignatureDoesNotMatchSignatureDoesNotMatch
+ {row.model}
+
+
+ {routedRows.length} decision {plural(routedRows.length, "model", "models")} across {routedProviderCount}{" "}
+ {plural(routedProviderCount, "provider", "providers")} {plural(routedRows.length, "answers", "answer")} on{" "}
+ /v1/decisions.
+ {selfHostedCount > 0 &&
+ ` ${selfHostedCount} more ${plural(selfHostedCount, "is self-hosted and answers on its own route", "are self-hosted and answer on their own routes")}.`}
+ {needsUpgradeCount > 0 &&
+ ` ${needsUpgradeCount} more ${plural(needsUpgradeCount, "needs", "need")} a newer LiteLLM version.`}
+
{model.model_group}
- {model.mode &&{announcement.title}
+{announcement.description}