diff --git a/litellm/constants.py b/litellm/constants.py index b4dcac764f9..7b83f6e2d6c 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -450,6 +450,7 @@ MCP_PEEKED_BODY_SCOPE_KEY: Final = "litellm_mcp_peeked_body" MCP_ADMISSION_BODY_PEEK_TIMEOUT_SECONDS: Final = 5.0 REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" REALTIME_METERED_USAGE_KEY: Final = "realtime_metered_usage" +FAILED_PROMPT_TOKEN_ESTIMATE_KEY: Final = "failed_prompt_token_estimate" OPENAI_LIVE_SESSION_START_TIMEOUT_SECONDS: Final = 30 OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS: Final = 5 OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES: Final = frozenset( diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index b582d6b8c4f..ae1651352fe 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -66,6 +66,7 @@ _BASE64_INLINE_PATTERN: Final = re.compile( class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callback#callback-class # Class variables or attributes server_fulfilled_tool_names: ClassVar[frozenset[str]] = frozenset() + wants_failed_prompt_token_estimate: ClassVar[bool] = False enforces_request_content: bool = False """ diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index dc14d3904c0..829e70b2cb8 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -6,14 +6,14 @@ from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass, field from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Protocol, cast +from typing import TYPE_CHECKING, Any, ClassVar, Final, Protocol, cast from pydantic import ConfigDict, TypeAdapter from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger -from litellm.constants import KUBERNETES_POD_ROUTING_KEY +from litellm.constants import FAILED_PROMPT_TOKEN_ESTIMATE_KEY, KUBERNETES_POD_ROUTING_KEY from litellm.integrations._types.open_inference import ( OpenInferenceSpanKindValues, ) @@ -181,6 +181,8 @@ HTTP_RESPONSE_STATUS_CODE_ATTRIBUTE: Final = "http.response.status_code" HTTP_ROUTE_ATTRIBUTE: Final = "http.route" URL_PATH_ATTRIBUTE: Final = "url.path" PREPROCESSING_DURATION_MS_ATTRIBUTE: Final = "litellm.preprocessing.duration_ms" +USAGE_PROVIDER_STATUS_ATTRIBUTE: Final = "litellm.usage.provider_status" +ESTIMATED_INPUT_TOKENS_ATTRIBUTE: Final = "litellm.usage.estimated_input_tokens" TEAM_METADATA_ATTRIBUTE: Final = "litellm.team.metadata" MODEL_GROUP_ATTRIBUTE: Final = "litellm.model_group" PROVIDER_MODEL_ATTRIBUTE: Final = "litellm.provider.model" @@ -455,6 +457,8 @@ class OpenTelemetryConfig: class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): + wants_failed_prompt_token_estimate: ClassVar[bool] = True + def __init__( self, config: OpenTelemetryConfig | None = None, @@ -2240,6 +2244,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # Always create if no parent span exists (backward compatibility) # OR if USE_OTEL_LITELLM_REQUEST_SPAN is explicitly enabled should_create_primary_span: Final = parent_otel_span is None or get_secret_bool("USE_OTEL_LITELLM_REQUEST_SPAN") + failure_kwargs: Final = _JSON_OBJECT.validate_python(kwargs) span = None if should_create_primary_span: @@ -2254,6 +2259,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): span_kwargs["kind"] = self.span_kind.CLIENT span = otel_tracer.start_span(**span_kwargs) self.set_attributes(span, kwargs, response_obj) + self._set_failure_usage_attributes(span, failure_kwargs) # Record exception information using OTEL standard method span_error_message: Final = self._record_exception_on_span(span=span, kwargs=kwargs) @@ -2267,6 +2273,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # Note: parent_otel_span is guaranteed to be not None here if parent_otel_span is not None and parent_otel_span.is_recording(): self.set_attributes(parent_otel_span, kwargs, response_obj) + self._set_failure_usage_attributes(parent_otel_span, failure_kwargs) parent_error_message: Final = self._record_exception_on_span(span=parent_otel_span, kwargs=kwargs) parent_otel_span.set_status(Status(StatusCode.ERROR, parent_error_message)) @@ -2286,6 +2293,14 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ): parent_otel_span.end(end_time=self._to_ns(end_time)) + def _set_failure_usage_attributes(self, span: Span, kwargs: Mapping[str, object]) -> None: + if FAILED_PROMPT_TOKEN_ESTIMATE_KEY not in kwargs: + return + self.safe_set_attribute(span=span, key=USAGE_PROVIDER_STATUS_ATTRIBUTE, value="unknown") + estimate: Final = kwargs[FAILED_PROMPT_TOKEN_ESTIMATE_KEY] + if isinstance(estimate, int): + self.safe_set_attribute(span=span, key=ESTIMATED_INPUT_TOKENS_ATTRIBUTE, value=estimate) + def _record_exception_on_span(self, span: Span, kwargs: _FailureKwargsView) -> str | None: """ Record exception information on the span using OTEL standard methods. diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 97f3ca68e42..0ca09140f40 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -29,6 +29,7 @@ from typing_extensions import TypedDict import litellm from litellm._internal_context import in_mirrored_call, outside_mirrored_call, post_response_phase +from litellm._lazy_imports import get_token_counter_new from litellm._logging import ( is_debugging_on, redact_string, @@ -47,6 +48,7 @@ from litellm.constants import ( DEFAULT_MOCK_RESPONSE_COMPLETION_TOKEN_COUNT, DEFAULT_MOCK_RESPONSE_PROMPT_TOKEN_COUNT, EMPTY_MAPPING, + FAILED_PROMPT_TOKEN_ESTIMATE_KEY, INTERNAL_CALLBACK_NAMES, KUBERNETES_POD_ROUTING_KEY, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, @@ -674,6 +676,9 @@ class UsageObservation(Enum): ESTIMATED = "estimated" +_CHAT_CALL_TYPES: Final = frozenset({CallTypes.completion.value, CallTypes.acompletion.value}) + + class _UsageFields(TypedDict, total=False, extra_items=object): pass @@ -828,7 +833,7 @@ class Logging(LiteLLMLoggingBaseClass): # Shallow copy of the outer list only (inner message dicts are shared). # Safe because the logging layer does not mutate individual message dicts. _copy_start: Final = time.time() - self.messages = copy.copy(messages) if messages is not None else None + self.messages: Sequence[object] | None = copy.copy(messages) if messages is not None else None self.message_copy_duration_ms: float = (time.time() - _copy_start) * 1000 self.callback_duration_ms: float = 0.0 self.stream = stream @@ -4683,7 +4688,9 @@ class Logging(LiteLLMLoggingBaseClass): except Exception as e: verbose_logger.debug("Error in _handle_callback_failure: %s", e) - def _failure_handler_helper_fn(self, exception, traceback_exception, start_time=None, end_time=None): + def _failure_handler_helper_fn( + self, exception, traceback_exception, start_time=None, end_time=None, *, estimate_prompt_tokens: bool = True + ): if start_time is None: start_time = self.start_time if end_time is None: @@ -4714,6 +4721,8 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["response_cost"] = self.finalized_cost.amount elif self.model_call_details.get("combined_usage_object") is None: self.model_call_details["response_cost"] = 0 + if estimate_prompt_tokens and self._should_estimate_failed_prompt(): + self.model_call_details[FAILED_PROMPT_TOKEN_ESTIMATE_KEY] = self._prompt_token_estimate() raised_exception: Final[object] = exception exception_headers: Final = ( @@ -4748,6 +4757,43 @@ class Logging(LiteLLMLoggingBaseClass): ) return start_time, end_time + def _should_estimate_failed_prompt(self) -> bool: + return ( + self._model_call_details_view.get("call_type") in _CHAT_CALL_TYPES + and not isinstance(self.model_call_details.get("combined_usage_object"), Usage) + and self._failure_callbacks_want_prompt_token_estimate() + ) + + def _failure_callbacks_want_prompt_token_estimate(self) -> bool: + callbacks: Final = ( + *(self.dynamic_failure_callbacks or ()), + *(self.dynamic_async_failure_callbacks or ()), + *litellm.callbacks, + *litellm.logging_callback_manager.get_failure_callbacks(), + ) + return any(isinstance(cb, CustomLogger) and cb.wants_failed_prompt_token_estimate for cb in callbacks) + + def _prompt_token_estimate(self) -> int | None: + from litellm.litellm_core_utils.token_counter import has_only_text_content + + messages: Final = self.messages + if litellm.disable_token_counter or not messages: + return None + try: + if not has_only_text_content(messages): + return None + return get_token_counter_new()(model=self.model, messages=messages, use_default_image_token_count=True) + except (ValueError, TypeError, LookupError, AttributeError, RuntimeError, OSError): + return None + + async def _offloaded_prompt_token_estimate(self) -> int | None: + from litellm.litellm_core_utils.token_counter import offload_token_count + + try: + return await offload_token_count(self._prompt_token_estimate)() + except (RuntimeError, OSError): + return None + async def special_failure_handlers(self, exception: Exception): """ Custom events, emitted for specific failures. @@ -4826,6 +4872,7 @@ class Logging(LiteLLMLoggingBaseClass): traceback_exception=traceback_exception, start_time=start_time, end_time=end_time, + estimate_prompt_tokens=not for_async_request, ) callbacks: Final = self.get_combined_callback_list( dynamic_success_callbacks=self.dynamic_failure_callbacks, @@ -5025,7 +5072,10 @@ class Logging(LiteLLMLoggingBaseClass): traceback_exception=traceback_exception, start_time=start_time, end_time=end_time, + estimate_prompt_tokens=False, ) + if self._should_estimate_failed_prompt(): + self.model_call_details[FAILED_PROMPT_TOKEN_ESTIMATE_KEY] = await self._offloaded_prompt_token_estimate() if self._on_failure_payload_built is not None: self._on_failure_payload_built() callbacks: Final = self.get_combined_callback_list( diff --git a/litellm/litellm_core_utils/logging_callback_manager.py b/litellm/litellm_core_utils/logging_callback_manager.py index 56c6436eae6..70aca83e70f 100644 --- a/litellm/litellm_core_utils/logging_callback_manager.py +++ b/litellm/litellm_core_utils/logging_callback_manager.py @@ -428,6 +428,9 @@ class LoggingCallbackManager: _get_all_callbacks = get_all_callbacks + def get_failure_callbacks(self) -> tuple[CustomLogger | Callable[..., object] | str, ...]: + return (*litellm.failure_callback, *litellm._async_failure_callback) + def remove_callback_from_all_lists(self, obj: object, require_self: bool = False) -> None: """ Remove a callback object from every callback list it may have been @@ -489,7 +492,7 @@ class LoggingCallbackManager: """ # Get callback lists success_callbacks: Final = set(litellm.success_callback + litellm._async_success_callback) - failure_callbacks: Final = set(litellm.failure_callback + litellm._async_failure_callback) + failure_callbacks: Final = set(self.get_failure_callbacks()) general_callbacks: Final = set(litellm.callbacks) # Get all unique callbacks diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index f19ea60fc76..11b02b9272b 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -17,7 +17,7 @@ import anyio import anyio.lowlevel import httpx import tiktoken -from pydantic import TypeAdapter +from pydantic import BaseModel, TypeAdapter from tiktoken.model import MODEL_TO_ENCODING from typing_extensions import ParamSpec, TypeVar @@ -48,7 +48,7 @@ from litellm.constants import ( TOKEN_COUNTER_MAX_EXACT_CHARS, ) from litellm.litellm_core_utils.asyncify import asyncify -from litellm.litellm_core_utils.core_helpers import as_str_mapping +from litellm.litellm_core_utils.core_helpers import as_object_list, as_str_mapping from litellm.litellm_core_utils.prompt_templates.chat_template_tools import ( ChatTemplateToolSection, chat_template_tool_section, @@ -929,7 +929,6 @@ def token_counter( if litellm.disable_token_counter is True: return 0 - verbose_logger.debug("messages in token_counter: %s, text in token_counter: %s", messages, text) if text is not None and messages is not None: raise ValueError("text and messages cannot both be set") if use_default_image_token_count is None: @@ -1872,6 +1871,29 @@ _TEXT_FIELD_BY_BLOCK_TYPE: Final = MappingProxyType( ) +def _is_text_block(block: object) -> bool: + mapping: Final = as_str_mapping(block) + block_type: Final = mapping.get("type") if mapping is not None else None + return isinstance(block, str) or (isinstance(block_type, str) and block_type in _TEXT_FIELD_BY_BLOCK_TYPE) + + +def _has_only_text_content(message: object) -> bool: + normalized: Final[object] = message.model_dump(exclude_none=True) if isinstance(message, BaseModel) else message + mapping: Final = as_str_mapping(normalized) + if mapping is None: + return False + content: Final = mapping.get("content") + if content is None or isinstance(content, str): + return True + blocks: Final = as_object_list(content) + return blocks is not None and all(_is_text_block(block) for block in blocks) + + +def has_only_text_content(messages: Sequence[object]) -> bool: + """Whether counting ``messages`` stays on the tokenizer, never decoding images, files or documents.""" + return all(_has_only_text_content(message) for message in messages) + + def _count_content_list( count_function: TokenCounterFunction, content_list: str diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index 6f2424e74d2..36e74cea05e 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -40,7 +40,9 @@ import litellm from litellm.integrations import opentelemetry as otel_module from litellm.integrations.arize.arize_phoenix import ArizePhoenixLogger from litellm.integrations.opentelemetry import ( + ESTIMATED_INPUT_TOKENS_ATTRIBUTE, LITELLM_REQUEST_SPAN_NAME, + USAGE_PROVIDER_STATUS_ATTRIBUTE, OpenTelemetry, OpenTelemetryConfig, OTELMetricAttributeFilter, @@ -54,13 +56,30 @@ from litellm.proxy import proxy_server from litellm.proxy._types import ManagementEndpointLoggingPayload, ProxyException, SpanAttributes from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth from litellm.types.services import ServiceLoggerPayload, ServiceTypes -from collections.abc import AsyncIterator, Callable, Mapping -from litellm.constants import KUBERNETES_POD_ROUTING_KEY, LOGGING_WORKER_MAX_TIME_PER_COROUTINE +from collections.abc import AsyncIterator, Callable, Iterator, Mapping +from litellm.constants import ( + FAILED_PROMPT_TOKEN_ESTIMATE_KEY, + KUBERNETES_POD_ROUTING_KEY, + LOGGING_WORKER_MAX_TIME_PER_COROUTINE, + REDACTED_BY_LITELLM, +) +from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.litellm_core_utils import litellm_logging as litellm_logging_module +from litellm.litellm_core_utils import token_counter as token_counter_module +from litellm.litellm_core_utils.token_counter import token_counter from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.hidden_params import set_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper -from litellm.types.utils import ChatCompletionMessageToolCall, Delta, Function, ModelResponseStream, StreamingChoices +from litellm.types.utils import ( + ChatCompletionMessageToolCall, + Delta, + Function, + ModelResponseStream, + Message, + StreamingChoices, + Usage, +) from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome exporter = InMemorySpanExporter() @@ -7940,3 +7959,428 @@ class TestOpentelemetryUnitTests: # Assert: error.message should be set from error_str using ErrorAttributes constant mock_span.set_attribute.assert_called_with(ErrorAttributes.ERROR_MESSAGE, "Fallback error message") + + +_PROMPT_SECRET: Final = "confidential-prompt-marker" +_SECRET_PROMPT: Final = ({"role": "user", "content": " ".join([_PROMPT_SECRET] * 200)},) +_UNCOUNTABLE_PROMPT: Final = ({"role": "assistant", "content": None, "tool_calls": [{"id": "call_without_function"}]},) +_STRUCTURED_TEXT_PROMPT: Final = ( + {"role": "system", "content": [{"type": "text", "text": "Answer briefly."}]}, + {"role": "user", "content": ["plain part", {"type": "text", "text": _PROMPT_SECRET}]}, +) +_PDF_BASE64: Final = "JVBERi0xLjQK" +_ANTHROPIC_PDF_DOCUMENT: Final = { + "type": "document", + "source": {"type": "base64", "media_type": "application/pdf", "data": _PDF_BASE64}, +} +_FAILED_MODEL: Final = "hosted_vllm/test-model" +_SYNC_THEN_ASYNC: Final = "sync_redaction_first" +_ASYNC_THEN_SYNC: Final = "async_callback_first" +_GEN_AI_USAGE_KEYS: Final = ( + SpanAttributes.GEN_AI_USAGE_INPUT_TOKENS.value, + SpanAttributes.GEN_AI_USAGE_OUTPUT_TOKENS.value, + SpanAttributes.GEN_AI_USAGE_TOTAL_TOKENS.value, +) + + +@pytest.fixture +def tracer_provider() -> Iterator[TracerProvider]: + provider: Final = TracerProvider() + yield provider + provider.shutdown() + + +@pytest.fixture +def in_memory_otel(tracer_provider: TracerProvider) -> tuple[OpenTelemetry, InMemorySpanExporter]: + exporter: Final = InMemorySpanExporter() + tracer_provider.add_span_processor(SimpleSpanProcessor(exporter)) + otel: Final = OpenTelemetry(tracer_provider=tracer_provider) + otel.tracer = tracer_provider.get_tracer("litellm") + return otel, exporter + + +def _redacted_logging( + otel: OpenTelemetry, messages: list[dict[str, object]] | None, call_type: str = "acompletion" +) -> LitellmLogging: + logging_obj: Final = LitellmLogging( + model=_FAILED_MODEL, + messages=messages, + stream=False, + call_type=call_type, + start_time=_FIXED_START, + litellm_call_id="failed-request-call-id", + function_id="failed-request-function-id", + dynamic_failure_callbacks=[otel], + dynamic_async_failure_callbacks=[otel], + kwargs={"turn_off_message_logging": True}, + ) + logging_obj.update_environment_variables(model=_FAILED_MODEL, user="", optional_params={}, litellm_params={}) + return logging_obj + + +async def _fail(logging_obj: LitellmLogging, exporter: InMemorySpanExporter, order: str) -> ReadableSpan: + error: Final = litellm.Timeout(message="timed out", model="test-model", llm_provider="hosted_vllm") + if order == _SYNC_THEN_ASYNC: + logging_obj.failure_handler(error, "tb", for_async_request=True) + await logging_obj.async_failure_handler(error, "tb") + else: + await logging_obj.async_failure_handler(error, "tb") + logging_obj.failure_handler(error, "tb", for_async_request=True) + assert logging_obj.model_call_details["messages"] == [{"role": "user", "content": REDACTED_BY_LITELLM}] + (span,) = exporter.get_finished_spans() + return span + + +def _span_text(span: ReadableSpan) -> str: + event_attributes: Final = tuple(dict(event.attributes or {}) for event in span.events) + return safe_dumps({"attributes": dict(span.attributes or {}), "events": event_attributes}) + + +def _assert_unknown_usage(attributes: Mapping[str, object]) -> None: + assert attributes[USAGE_PROVIDER_STATUS_ATTRIBUTE] == "unknown" + assert not any(key in attributes for key in _GEN_AI_USAGE_KEYS) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("order", [_SYNC_THEN_ASYNC, _ASYNC_THEN_SYNC]) +async def test_failed_chat_span_flags_unknown_usage_with_pre_redaction_prompt_estimate( + in_memory_otel: tuple[OpenTelemetry, InMemorySpanExporter], order: str +) -> None: + otel, exporter = in_memory_otel + expected_estimate: Final = token_counter(model=_FAILED_MODEL, messages=list(_SECRET_PROMPT)) + redacted_estimate: Final = token_counter( + model=_FAILED_MODEL, messages=[{"role": "user", "content": REDACTED_BY_LITELLM}] + ) + assert expected_estimate != redacted_estimate + + span: Final = await _fail(_redacted_logging(otel, list(_SECRET_PROMPT)), exporter, order) + + attributes: Final = dict(span.attributes or {}) + _assert_unknown_usage(attributes) + assert attributes[ESTIMATED_INPUT_TOKENS_ATTRIBUTE] == expected_estimate + assert _PROMPT_SECRET not in _span_text(span) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "messages", + [ + None, + [], + list(_UNCOUNTABLE_PROMPT), + [{"role": "user", "content": [{"type": [], "text": _PROMPT_SECRET}]}], + [{"role": "user", "content": [{"type": {"name": "text"}, "text": _PROMPT_SECRET}]}], + [{"role": "user", "content": {"type": "text", "text": _PROMPT_SECRET}}], + ], + ids=[ + "no_messages", + "empty_messages", + "tokenizer_failure", + "list_block_type", + "dict_block_type", + "mapping_content", + ], +) +async def test_failed_chat_span_flags_unknown_usage_without_inventing_an_estimate( + in_memory_otel: tuple[OpenTelemetry, InMemorySpanExporter], messages: list[dict[str, object]] | None +) -> None: + otel, exporter = in_memory_otel + + span: Final = await _fail(_redacted_logging(otel, messages), exporter, _SYNC_THEN_ASYNC) + + attributes: Final = dict(span.attributes or {}) + _assert_unknown_usage(attributes) + assert ESTIMATED_INPUT_TOKENS_ATTRIBUTE not in attributes + assert span.status.status_code == trace.StatusCode.ERROR + assert any("Timeout" in str((event.attributes or {}).get("exception.type")) for event in span.events) + assert _PROMPT_SECRET not in _span_text(span) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("call_type", ["acompletion", "completion"], ids=["async_request", "sync_request"]) +async def test_failed_prompt_estimate_keeps_the_prompt_out_of_debug_logs( + in_memory_otel: tuple[OpenTelemetry, InMemorySpanExporter], caplog: pytest.LogCaptureFixture, call_type: str +) -> None: + otel, exporter = in_memory_otel + logging_obj: Final = _redacted_logging(otel, list(_SECRET_PROMPT), call_type=call_type) + error: Final = litellm.Timeout(message="timed out", model="test-model", llm_provider="hosted_vllm") + + with caplog.at_level("DEBUG", logger="LiteLLM"): + if call_type == "acompletion": + logging_obj.failure_handler(error, "tb", for_async_request=True) + await logging_obj.async_failure_handler(error, "tb") + else: + logging_obj.failure_handler(error, "tb") + + assert caplog.records + assert logging_obj.model_call_details[FAILED_PROMPT_TOKEN_ESTIMATE_KEY] == token_counter( + model=_FAILED_MODEL, messages=list(_SECRET_PROMPT) + ) + assert _PROMPT_SECRET not in caplog.text + assert not any(_PROMPT_SECRET in _span_text(span) for span in exporter.get_finished_spans()) + + +async def _raise_offload_error(error: BaseException) -> int: + raise error + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("seam", "error"), + [ + ("offload", RuntimeError("can't start new thread")), + ("offload", OSError("worker unavailable")), + ("counter", RuntimeError("tokenizer failed to load")), + ], + ids=["offload_runtime_error", "offload_os_error", "counter_runtime_error"], +) +async def test_failed_chat_span_survives_estimate_errors( + in_memory_otel: tuple[OpenTelemetry, InMemorySpanExporter], + monkeypatch: pytest.MonkeyPatch, + seam: str, + error: Exception, +) -> None: + otel, exporter = in_memory_otel + + def failing_counter(**_kwargs: object) -> int: + raise error + + if seam == "offload": + monkeypatch.setattr( + token_counter_module, "offload_token_count", lambda _function: lambda: _raise_offload_error(error) + ) + else: + monkeypatch.setattr(litellm_logging_module, "get_token_counter_new", lambda: failing_counter) + + span: Final = await _fail(_redacted_logging(otel, list(_SECRET_PROMPT)), exporter, _SYNC_THEN_ASYNC) + + attributes: Final = dict(span.attributes or {}) + _assert_unknown_usage(attributes) + assert ESTIMATED_INPUT_TOKENS_ATTRIBUTE not in attributes + assert span.status.status_code == trace.StatusCode.ERROR + assert any("Timeout" in str((event.attributes or {}).get("exception.type")) for event in span.events) + + +@pytest.mark.asyncio +async def test_cancelled_prompt_estimate_cancels_async_failure_logging( + in_memory_otel: tuple[OpenTelemetry, InMemorySpanExporter], monkeypatch: pytest.MonkeyPatch +) -> None: + otel, exporter = in_memory_otel + monkeypatch.setattr( + token_counter_module, + "offload_token_count", + lambda _function: lambda: _raise_offload_error(asyncio.CancelledError()), + ) + logging_obj: Final = _redacted_logging(otel, list(_SECRET_PROMPT)) + + with pytest.raises(asyncio.CancelledError): + await logging_obj.async_failure_handler( + litellm.Timeout(message="timed out", model="test-model", llm_provider="hosted_vllm"), "tb" + ) + + assert exporter.get_finished_spans() == () + + +@pytest.mark.asyncio +async def test_failed_chat_span_estimates_sdk_message_objects_like_dicts( + in_memory_otel: tuple[OpenTelemetry, InMemorySpanExporter], +) -> None: + otel, exporter = in_memory_otel + text: Final = " ".join([_PROMPT_SECRET] * 50) + + message_span: Final = await _fail( + _redacted_logging(otel, [Message(role="user", content=text)]), exporter, _SYNC_THEN_ASYNC + ) + exporter.clear() + dict_span: Final = await _fail( + _redacted_logging(otel, [{"role": "user", "content": text}]), exporter, _SYNC_THEN_ASYNC + ) + + message_estimate: Final = dict(message_span.attributes or {}).get(ESTIMATED_INPUT_TOKENS_ATTRIBUTE) + assert isinstance(message_estimate, int) + assert message_estimate == dict(dict_span.attributes or {}).get(ESTIMATED_INPUT_TOKENS_ATTRIBUTE) + assert _PROMPT_SECRET not in _span_text(message_span) + + +@pytest.mark.asyncio +async def test_failed_chat_span_estimates_structured_text_prompts( + in_memory_otel: tuple[OpenTelemetry, InMemorySpanExporter], +) -> None: + otel, exporter = in_memory_otel + + span: Final = await _fail(_redacted_logging(otel, list(_STRUCTURED_TEXT_PROMPT)), exporter, _SYNC_THEN_ASYNC) + + attributes: Final = dict(span.attributes or {}) + _assert_unknown_usage(attributes) + assert attributes[ESTIMATED_INPUT_TOKENS_ATTRIBUTE] == token_counter( + model=_FAILED_MODEL, messages=list(_STRUCTURED_TEXT_PROMPT) + ) + assert _PROMPT_SECRET not in _span_text(span) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "attachment", + [ + {"type": "file", "file": {"filename": "report.pdf", "file_data": f"data:application/pdf;base64,{_PDF_BASE64}"}}, + _ANTHROPIC_PDF_DOCUMENT, + {"type": "tool_result", "tool_use_id": "toolu_report", "content": [_ANTHROPIC_PDF_DOCUMENT]}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="}}, + ], + ids=["openai_file", "anthropic_document", "document_inside_tool_result", "image_url"], +) +async def test_failed_chat_span_skips_the_estimate_for_prompts_with_attachments( + in_memory_otel: tuple[OpenTelemetry, InMemorySpanExporter], + monkeypatch: pytest.MonkeyPatch, + attachment: dict[str, object], +) -> None: + otel, exporter = in_memory_otel + parsed_pdfs: Final[list[str]] = [] + + def record_pdf_parse(data_url: str, *_args: object) -> None: + parsed_pdfs.append(data_url) + + monkeypatch.setattr(token_counter_module, "_count_inline_pdf_tokens", record_pdf_parse) + messages: Final = [{"role": "user", "content": [{"type": "text", "text": _PROMPT_SECRET}, attachment]}] + + span: Final = await _fail(_redacted_logging(otel, messages), exporter, _SYNC_THEN_ASYNC) + + attributes: Final = dict(span.attributes or {}) + _assert_unknown_usage(attributes) + assert ESTIMATED_INPUT_TOKENS_ATTRIBUTE not in attributes + assert parsed_pdfs == [] + assert _PROMPT_SECRET not in _span_text(span) + + +@pytest.mark.asyncio +async def test_async_failure_counts_the_prompt_off_the_event_loop( + in_memory_otel: tuple[OpenTelemetry, InMemorySpanExporter], monkeypatch: pytest.MonkeyPatch +) -> None: + otel, exporter = in_memory_otel + counting_started: Final = threading.Event() + release_counting: Final = threading.Event() + + def blocked_token_counter(**_kwargs: object) -> int: + counting_started.set() + release_counting.wait(timeout=30) + return 4321 + + monkeypatch.setattr(litellm_logging_module, "get_token_counter_new", lambda: blocked_token_counter) + failure: Final = asyncio.create_task( + _fail(_redacted_logging(otel, list(_SECRET_PROMPT)), exporter, _SYNC_THEN_ASYNC) + ) + try: + assert await asyncio.to_thread(counting_started.wait, 30) + other_coroutine_ran: Final = asyncio.Event() + + async def other_request() -> None: + other_coroutine_ran.set() + + await asyncio.create_task(other_request()) + assert other_coroutine_ran.is_set() + assert not failure.done() + finally: + release_counting.set() + await asyncio.gather(failure, return_exceptions=True) + span: Final = failure.result() + + attributes: Final = dict(span.attributes or {}) + _assert_unknown_usage(attributes) + assert attributes[ESTIMATED_INPUT_TOKENS_ATTRIBUTE] == 4321 + + +@pytest.mark.asyncio +async def test_failed_chat_span_flags_unknown_usage_without_estimate_when_token_counter_is_disabled( + in_memory_otel: tuple[OpenTelemetry, InMemorySpanExporter], monkeypatch: pytest.MonkeyPatch +) -> None: + otel, exporter = in_memory_otel + monkeypatch.setattr(litellm, "disable_token_counter", True) + + span: Final = await _fail(_redacted_logging(otel, list(_SECRET_PROMPT)), exporter, _SYNC_THEN_ASYNC) + + attributes: Final = dict(span.attributes or {}) + _assert_unknown_usage(attributes) + assert ESTIMATED_INPUT_TOKENS_ATTRIBUTE not in attributes + + +@pytest.mark.asyncio +async def test_failure_inside_active_parent_span_records_unknown_usage_on_the_parent( + in_memory_otel: tuple[OpenTelemetry, InMemorySpanExporter], + tracer_provider: TracerProvider, + monkeypatch: pytest.MonkeyPatch, +) -> None: + otel, exporter = in_memory_otel + monkeypatch.delenv("USE_OTEL_LITELLM_REQUEST_SPAN", raising=False) + logging_obj: Final = _redacted_logging(otel, list(_SECRET_PROMPT)) + error: Final = litellm.Timeout(message="timed out", model="test-model", llm_provider="hosted_vllm") + + with trace.use_span(tracer_provider.get_tracer("caller").start_span("caller_request"), end_on_exit=True): + logging_obj.failure_handler(error, "tb") + await logging_obj.async_failure_handler(error, "tb") + + (parent,) = exporter.get_finished_spans() + attributes: Final = dict(parent.attributes or {}) + assert parent.name == "caller_request" + _assert_unknown_usage(attributes) + assert attributes[ESTIMATED_INPUT_TOKENS_ATTRIBUTE] == token_counter( + model=_FAILED_MODEL, messages=list(_SECRET_PROMPT) + ) + assert _PROMPT_SECRET not in _span_text(parent) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "recovered", + [ + Usage(prompt_tokens=0, completion_tokens=0, total_tokens=0), + Usage(prompt_tokens=120, completion_tokens=7, total_tokens=127), + ], + ids=["recovered_zero", "recovered_partial_stream"], +) +async def test_failed_chat_span_adds_no_usage_attributes_when_usage_was_recovered( + in_memory_otel: tuple[OpenTelemetry, InMemorySpanExporter], recovered: Usage +) -> None: + otel, exporter = in_memory_otel + logging_obj: Final = _redacted_logging(otel, list(_SECRET_PROMPT)) + logging_obj.record_partial_usage_for_failure(recovered, 0.0) + + span: Final = await _fail(logging_obj, exporter, _SYNC_THEN_ASYNC) + + attributes: Final = dict(span.attributes or {}) + assert USAGE_PROVIDER_STATUS_ATTRIBUTE not in attributes + assert ESTIMATED_INPUT_TOKENS_ATTRIBUTE not in attributes + assert not any(key in attributes for key in _GEN_AI_USAGE_KEYS) + + +@pytest.mark.asyncio +async def test_failed_non_chat_span_gets_no_failure_usage_flags( + in_memory_otel: tuple[OpenTelemetry, InMemorySpanExporter], +) -> None: + otel, exporter = in_memory_otel + logging_obj: Final = _redacted_logging(otel, list(_SECRET_PROMPT), call_type="aresponses") + + span: Final = await _fail(logging_obj, exporter, _SYNC_THEN_ASYNC) + + attributes: Final = dict(span.attributes or {}) + assert USAGE_PROVIDER_STATUS_ATTRIBUTE not in attributes + assert ESTIMATED_INPUT_TOKENS_ATTRIBUTE not in attributes + + +def test_successful_span_reports_provider_usage_without_failure_usage_flags( + in_memory_otel: tuple[OpenTelemetry, InMemorySpanExporter], +) -> None: + otel, exporter = in_memory_otel + data_dir: Final = Path(__file__).parent / "open_telemetry" / "data" + kwargs: Final = { + **json.loads((data_dir / "captured_kwargs.json").read_text()), + FAILED_PROMPT_TOKEN_ESTIMATE_KEY: 999, + } + response_obj: Final = json.loads((data_dir / "captured_response.json").read_text()) + + otel.log_success_event(kwargs, response_obj, _FIXED_START, _FIXED_START + timedelta(seconds=1)) + + span_attributes: Final = tuple(dict(span.attributes or {}) for span in exporter.get_finished_spans()) + input_tokens: Final = {attrs.get(SpanAttributes.GEN_AI_USAGE_INPUT_TOKENS.value) for attrs in span_attributes} + assert response_obj["usage"]["prompt_tokens"] in input_tokens + assert not any(USAGE_PROVIDER_STATUS_ATTRIBUTE in attrs for attrs in span_attributes) + assert not any(ESTIMATED_INPUT_TOKENS_ATTRIBUTE in attrs for attrs in span_attributes) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 80782fa89d4..41cd17b72db 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -31,6 +31,7 @@ from litellm._logging import session_id_var, trace_id_var, verbose_logger from litellm._service_logger import ServiceLogging from litellm.caching.caching import DualCache from litellm.constants import ( + FAILED_PROMPT_TOKEN_ESTIMATE_KEY, KUBERNETES_POD_ROUTING_KEY, LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL, LOGGING_WORKER_MAX_TIME_PER_COROUTINE, @@ -41,6 +42,7 @@ from litellm.cost_calculator import ocr_batch_cost from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging +from litellm.litellm_core_utils.token_counter import token_counter from litellm.litellm_core_utils.litellm_logging import ( BilledRoundsSpend, Logging, @@ -14938,3 +14940,125 @@ async def test_a_response_discarded_after_its_success_log_was_parked_is_billed_o assert kwargs["response_cost"] == pytest.approx(both_attempts) assert prepared is not None assert prepared.amount == pytest.approx(both_attempts) + + +class _FailureKwargsRecorder(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.failure_kwargs: asyncio.Queue[dict[str, object]] = asyncio.Queue() + + async def async_log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: datetime.datetime, + end_time: datetime.datetime, + ) -> None: + self.failure_kwargs.put_nowait(dict(kwargs)) + + +class _EstimateWantingRecorder(_FailureKwargsRecorder): + wants_failed_prompt_token_estimate = True + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("recorder_type", "expects_estimate"), + [(_FailureKwargsRecorder, False), (_EstimateWantingRecorder, True)], + ids=["plain_logger_skips_counting", "opted_in_logger_gets_estimate"], +) +async def test_failed_chat_prompt_is_counted_only_for_loggers_that_want_the_estimate( + monkeypatch: pytest.MonkeyPatch, + recorder_type: type[_FailureKwargsRecorder], + expects_estimate: bool, +) -> None: + for callback_list in ("callbacks", "success_callback", "_async_success_callback", "failure_callback"): + monkeypatch.setattr(litellm, callback_list, []) + monkeypatch.setattr(litellm, "_async_failure_callback", []) + recorder: Final = recorder_type() + messages: Final = [{"role": "user", "content": "count these prompt words"}] + logging_obj: Final = LitellmLogging( + model="hosted_vllm/test-model", + messages=messages, + stream=False, + call_type="acompletion", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="failed-request-call-id", + function_id="failed-request-function-id", + dynamic_async_failure_callbacks=[recorder], + ) + logging_obj.update_environment_variables( + model="hosted_vllm/test-model", user="", optional_params={}, litellm_params={} + ) + + await logging_obj.async_failure_handler( + litellm.Timeout(message="timed out", model="test-model", llm_provider="hosted_vllm"), "tb" + ) + + received: Final = recorder.failure_kwargs.get_nowait() + expected: Final = { + FAILED_PROMPT_TOKEN_ESTIMATE_KEY: token_counter(model="hosted_vllm/test-model", messages=messages) + } + assert {k: v for k, v in received.items() if k == FAILED_PROMPT_TOKEN_ESTIMATE_KEY} == ( + expected if expects_estimate else {} + ) + + +_GLOBAL_CALLBACK_LISTS: Final = ( + "callbacks", + "success_callback", + "_async_success_callback", + "failure_callback", + "_async_failure_callback", +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("registration", "expects_estimate"), + [ + ("success_callback", False), + ("_async_success_callback", False), + ("callbacks", True), + ("failure_callback", True), + ("_async_failure_callback", True), + ("dynamic_failure_callbacks", True), + ("dynamic_async_failure_callbacks", True), + ], +) +async def test_failed_chat_prompt_is_counted_only_when_an_opted_in_logger_receives_failures( + monkeypatch: pytest.MonkeyPatch, registration: str, expects_estimate: bool +) -> None: + for callback_list in _GLOBAL_CALLBACK_LISTS: + monkeypatch.setattr(litellm, callback_list, []) + recorder: Final = _EstimateWantingRecorder() + if registration in _GLOBAL_CALLBACK_LISTS: + monkeypatch.setattr(litellm, registration, [recorder]) + counted_prompts: Final[list[object]] = [] + + def recording_token_counter(**kwargs: object) -> int: + counted_prompts.append(kwargs["messages"]) + return 7 + + monkeypatch.setattr(litellm_logging, "get_token_counter_new", lambda: recording_token_counter) + messages: Final = [{"role": "user", "content": "count these prompt words"}] + logging_obj: Final = LitellmLogging( + model="hosted_vllm/test-model", + messages=messages, + stream=False, + call_type="acompletion", + start_time=datetime.datetime(2026, 1, 1), + litellm_call_id="failed-request-call-id", + function_id="failed-request-function-id", + **{registration: [recorder]} if registration.startswith("dynamic_") else {}, + ) + logging_obj.update_environment_variables( + model="hosted_vllm/test-model", user="", optional_params={}, litellm_params={} + ) + + await logging_obj.async_failure_handler( + litellm.Timeout(message="timed out", model="test-model", llm_provider="hosted_vllm"), "tb" + ) + + assert counted_prompts == ([messages] if expects_estimate else []) + assert logging_obj.model_call_details.get(FAILED_PROMPT_TOKEN_ESTIMATE_KEY) == (7 if expects_estimate else None) diff --git a/tests/unit/litellm_core_utils/test_logging_callback_manager.py b/tests/unit/litellm_core_utils/test_logging_callback_manager.py index 25a52137d65..d7cad81f179 100644 --- a/tests/unit/litellm_core_utils/test_logging_callback_manager.py +++ b/tests/unit/litellm_core_utils/test_logging_callback_manager.py @@ -344,3 +344,28 @@ def test_duplicate_multiple_loggers_test(): assert langfuse_count == 1, "Should have exactly one LangfusePromptManagement instance" assert otel_count == 1, "Should have exactly one OpenTelemetry instance" + + +def test_get_callbacks_by_type_treats_both_failure_lists_as_failure(monkeypatch: pytest.MonkeyPatch) -> None: + for callback_list in ( + "callbacks", + "success_callback", + "_async_success_callback", + "failure_callback", + "_async_failure_callback", + ): + monkeypatch.setattr(litellm, callback_list, []) + manager = LoggingCallbackManager() + manager.add_litellm_callback("general_logger") + manager.add_litellm_success_callback("sync_success_logger") + manager.add_litellm_async_success_callback("async_success_logger") + manager.add_litellm_failure_callback("sync_failure_logger") + manager.add_litellm_async_failure_callback("async_failure_logger") + manager.add_litellm_success_callback("sync_success_and_async_failure_logger") + manager.add_litellm_async_failure_callback("sync_success_and_async_failure_logger") + + callbacks_by_type = manager.get_callbacks_by_type() + + assert set(callbacks_by_type["success"]) == {"sync_success_logger", "async_success_logger"} + assert set(callbacks_by_type["failure"]) == {"sync_failure_logger", "async_failure_logger"} + assert set(callbacks_by_type["success_and_failure"]) == {"general_logger", "sync_success_and_async_failure_logger"} diff --git a/tests/unit/litellm_core_utils/test_token_counter.py b/tests/unit/litellm_core_utils/test_token_counter.py index 7f42a39bc95..a6102a6e829 100644 --- a/tests/unit/litellm_core_utils/test_token_counter.py +++ b/tests/unit/litellm_core_utils/test_token_counter.py @@ -46,6 +46,7 @@ from litellm.litellm_core_utils.token_counter import ( _get_tiktoken_count_function, calculate_img_tokens, estimate_image_tokens, + has_only_text_content, high_detail_image_token_upper_bound, is_openai_model, messages_reach_token_count, @@ -3929,3 +3930,35 @@ def test_openai_model_check_holds_while_openai_models_are_registered_concurrentl registering.result() finally: litellm.open_ai_chat_completion_models.difference_update(registered) + + +@pytest.mark.parametrize( + ("messages", "text_only"), + [ + ([{"role": "user", "content": "hi"}, {"role": "assistant", "content": None}], True), + ([{"role": "user", "content": ["part", {"type": "text", "text": "hi"}, {"type": "thinking"}]}], True), + ([{"role": "user", "content": "hi"}, "a bare string message"], False), + ([{"role": "user", "content": "hi"}, 42], False), + ([None], False), + ([{"role": "user", "content": {"type": "text", "text": "hi"}}], False), + ( + [{"role": "user", "content": [{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA=="}}]}], + False, + ), + ([{"role": "user", "content": [{"type": ["text"], "text": "hi"}]}], False), + ([{"role": "user", "content": [{"text": "hi"}]}], False), + ], + ids=[ + "string_and_empty_content", + "text_blocks", + "string_message", + "number_message", + "none_message", + "mapping_content", + "image_block", + "non_string_block_type", + "untyped_block", + ], +) +def test_has_only_text_content_accepts_only_plain_text_chat_messages(messages: list[object], text_only: bool) -> None: + assert has_only_text_content(messages) is text_only