mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(otel): distinguish missing usage from prompt estimates on chat failures (#51443)
* fix(otel): distinguish unknown timeout usage from prompt estimates * fix(otel): handle expected tokenizer failures and cover parent spans * fix(otel): validate failure usage inputs at typed boundaries * fix(otel): type Logging.messages for the failed prompt token estimate Declare Logging.messages as Sequence[object] | None at its single __init__ assignment and drop the list TypeAdapter, so the failure estimate reads a typed value without new unknown-type diagnostics. * fix(otel): skip the failure prompt estimate for prompts with attachments Only estimate the failed prompt when every message is plain text, so files, documents and images never reach PDF parsing or media decoding on the failure path. Those requests still report unknown provider usage without an estimate. * fix(otel): count the failed prompt off the event loop for async requests Async failure logging now awaits the existing offload_token_count limiter for the optional prompt estimate, and the sync half of an async request no longer counts inline, so a slow tokenizer load cannot stall other requests. * fix(otel): only count the failed prompt for loggers that receive failures The estimate check now reads general callbacks, the global failure lists through a typed get_failure_callbacks accessor, and the dynamic failure lists, so a logger registered only for success never triggers failure token counting. * fix(otel): return an immutable failure callback snapshot * fix(otel): keep malformed content blocks from breaking failure logging Only compare string block types against the text block table, and run the text-only check inside the estimate's narrow try, so a malformed prompt leaves usage unknown without an estimate instead of interrupting failure callbacks. * fix(otel): stop the token counter from debug-logging raw prompts The failure estimate counts the original prompt, so the counter's debug line could write prompt text to application logs while message logging was off. Drop that line and assert the prompt stays out of debug logs on both routes. * test(otel): pin text-only eligibility and failure callback categories Cover the public has_only_text_content contract for malformed and non-text messages, and check get_callbacks_by_type puts sync and async failure registrations in the failure categories through the public manager APIs. * fix(otel): keep failure logging when the prompt estimate cannot run Return no estimate when the offloaded count hits a worker or tokenizer RuntimeError or OSError, so failure callbacks still run, while cancellation still propagates. Treat SDK Message objects like their dict form when checking whether a prompt is text only. --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
joshua-berri
GitHub
parent
5ebc284053
commit
dc1284d3d2
10 changed files with 729 additions and 11 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue