scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

fix(logging): drop bodies after masking hook timeouts

Browse source
This commit is contained in:
adc77 2026-10-20 01:27:25 +05:30
parent 2182e8590c
commit dce9f01e7d
3 changed files with 169 additions and 22 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

40
litellm/litellm_core_utils/litellm_logging.py
View file

@ -116,7 +116,6 @@ from litellm.litellm_core_utils.logging_worker import optional_callback_budget
from litellm.litellm_core_utils.model_param_helper import ModelParamHelper
from litellm.litellm_core_utils.ptu_pricing import is_spilled_over_ptu_request
from litellm.litellm_core_utils.redact_messages import (
perform_redaction,
redact_message_input_output_from_custom_logger,
redact_message_input_output_from_logging,
redact_streaming_responses_for_custom_logger,
@ -3574,7 +3573,11 @@ class Logging(LiteLLMLoggingBaseClass):
return result
async def _prepare_accounting_success_hooks(
self, callbacks: Sequence[str | Callable[..., object] | CustomLogger], result: object
self,
callbacks: Sequence[str | Callable[..., object] | CustomLogger],
result: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
) -> object:
from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE
@ -3582,14 +3585,30 @@ class Logging(LiteLLMLoggingBaseClass):
with fail_after(optional_callback_budget(LOGGING_WORKER_MAX_TIME_PER_COROUTINE, fraction=0.5)):
return await self._run_async_success_hooks(callbacks, result)
except TimeoutError:
from litellm.proxy.spend_tracking.spend_event import (
SpendEventBuildError,
build_spend_event_payload,
spend_event_callback_args,
)
self.shadow_eval_request_snapshot = None
self.model_call_details = redact_streaming_responses_for_custom_logger(
self.model_call_details, CustomLogger(message_logging=False)
assembled_stream: Final = "async_complete_streaming_response" in self.model_call_details
accounting_response: Final[object] = (
self.model_call_details["async_complete_streaming_response"] if assembled_stream else result
)
verbose_logger.warning("Logging hooks timed out; recording spend with request and response bodies redacted")
return cast( # cast-ok: legacy redaction returns heterogeneous response objects
object, perform_redaction(self.model_call_details, result)
event: Final = build_spend_event_payload(
self.model_call_details, accounting_response, start_time, end_time, store_bodies=False
)
if isinstance(event, SpendEventBuildError):
verbose_logger.error("Cannot prepare body-free spend accounting after a logging hook timeout")
return event
args: Final = spend_event_callback_args(event)
self.model_call_details = {
**args.kwargs,
**({"async_complete_streaming_response": args.response_obj} if assembled_stream else {}),
}
verbose_logger.warning("Logging hooks timed out, recording spend without request or response bodies")
return args.response_obj
async def async_success_handler(
self,
@ -3816,10 +3835,15 @@ class Logging(LiteLLMLoggingBaseClass):
}.values()
)
result = (
await self._prepare_accounting_success_hooks(hook_callbacks, redacted_result)
await self._prepare_accounting_success_hooks(hook_callbacks, redacted_result, start_time, end_time)
if accounting_callbacks
else await self._run_async_success_hooks(hook_callbacks, redacted_result)
)
if accounting_callbacks:
from litellm.proxy.spend_tracking.spend_event import SpendEventBuildError
if isinstance(result, SpendEventBuildError):
return
ordered_callbacks: Final = (
*accounting_callbacks,
*(

16
litellm/proxy/spend_tracking/spend_event.py
View file

@ -342,10 +342,10 @@ def _function_tool(name: str) -> _FunctionTool:
return tool
def build_spend_event(
def build_spend_event_payload(
raw_kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime, store_bodies: bool
) -> bytes | SpendEventBuildError:
"""Validate the success callback's kwargs and serialize the event once, as a single JSON line."""
) -> SpendEvent | SpendEventBuildError:
"""Validate and project the identity, usage, cost, and optional bodies of a completed call."""
try:
kwargs: Final = _SUCCESS_KWARGS.validate_python(raw_kwargs)
except ValidationError as error:
@ -379,6 +379,16 @@ def build_spend_event(
realtime_tool_calls=kwargs.get("realtime_tool_calls"),
request_tool_names=_request_tool_names(kwargs),
)
return event
def build_spend_event(
raw_kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime, store_bodies: bool
) -> bytes | SpendEventBuildError:
"""Validate the success callback's kwargs and serialize the event once, as a single JSON line."""
event: Final = build_spend_event_payload(raw_kwargs, response_obj, start_time, end_time, store_bodies)
if isinstance(event, SpendEventBuildError):
return event
return event.model_dump_json(fallback=_json_fallback).encode() + b"\n"

135
tests/unit/litellm_core_utils/test_litellm_logging.py
View file

@ -8009,12 +8009,14 @@ async def test_streaming_success_callbacks_survive_guardrail_logging_hook_failur
releasing.async_log_success_event.assert_awaited_once()
def _async_request_logging_obj(call_id: str, callbacks: list[CustomLogger], stream: bool = False) -> LitellmLogging:
def _async_request_logging_obj(
call_id: str, callbacks: list[CustomLogger], stream: bool = False, call_type: str = "acompletion"
) -> LitellmLogging:
logging_obj: Final = LitellmLogging(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
stream=stream,
call_type="acompletion",
call_type=call_type,
start_time=datetime.datetime.now(),
litellm_call_id=call_id,
function_id=call_id,
@ -12959,14 +12961,18 @@ def test_legacy_s3_logs_a_streamed_anthropic_messages_call_handed_off_from_the_a
assert s3_writer.log_event.call_count == expected_s3_writes
if expected_s3_writes:
assert s3_writer.log_event.call_args.kwargs["response_obj"] is complete_response
@pytest.mark.asyncio
@pytest.mark.parametrize("stream", [False, True])
@pytest.mark.parametrize("response_kind", ["chat", "responses"])
@pytest.mark.parametrize("accounting_first", [False, True])
@pytest.mark.parametrize("duplicate_accounting", [False, True])
@pytest.mark.parametrize("blocked_phase", ["hook", "delivery", "neither", "redacted"])
async def test_slow_optional_logging_cannot_cancel_accounting(
monkeypatch: pytest.MonkeyPatch,
stream: bool,
response_kind: str,
accounting_first: bool,
duplicate_accounting: bool,
blocked_phase: str,
@ -12979,6 +12985,7 @@ async def test_slow_optional_logging_cannot_cancel_accounting(
recorded: Final[asyncio.Queue[tuple[bytes, str]]] = asyncio.Queue()
hook_context: Final = contextvars.ContextVar("accounting_hook_context", default="outside")
delivered: Final = asyncio.Event()
deliveries: Final[asyncio.Queue[str]] = asyncio.Queue()
async def store_event(line: bytes) -> None:
recorded.put_nowait((line, hook_context.get()))
@ -13015,6 +13022,7 @@ async def test_slow_optional_logging_cannot_cancel_accounting(
end_time: datetime.datetime,
) -> None:
delivered.set()
deliveries.put_nowait(json.dumps(kwargs, default=str) + str(response_obj))
if blocked_phase == "delivery":
try:
await never.wait()
@ -13026,13 +13034,60 @@ async def test_slow_optional_logging_cannot_cancel_accounting(
accounting_callbacks: Final = [accounting, _ProxyDBLogger(producer)] if duplicate_accounting else [accounting]
callbacks: Final = [*accounting_callbacks, optional] if accounting_first else [optional, *accounting_callbacks]
monkeypatch.setattr(litellm, "_async_success_callback", [])
logging_obj: Final = _async_request_logging_obj("accounting-timeout", callbacks, stream=stream)
response: Final = ModelResponse(
choices=[{"message": {"role": "assistant", "content": "private-marker"}}],
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
logging_obj: Final = _async_request_logging_obj(
"accounting-timeout",
callbacks,
stream=stream,
call_type="aresponses" if response_kind == "responses" else "acompletion",
)
worker: Final = LoggingWorker(timeout=0.1)
worker.ensure_initialized_and_enqueue(logging_obj.async_success_handler(result=response))
response: Final = (
ModelResponse(
choices=[{"message": {"role": "assistant", "content": "private-marker"}}],
usage=Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
)
if response_kind == "chat"
else ResponsesAPIResponse(
id="resp-private-instructions",
created_at=1,
model="gpt-4o-mini",
status="completed",
instructions="private-instructions-canary" if blocked_phase != "redacted" else None,
output=[
{
"type": "message",
"id": "msg-private",
"role": "assistant",
"status": "completed",
"content": [{"type": "output_text", "text": "private-marker", "annotations": []}],
}
],
usage=ResponseAPIUsage(input_tokens=10, output_tokens=5, total_tokens=15),
)
)
original_response: Final = response.model_dump(exclude={"usage"})
logging_obj.update_environment_variables(
litellm_params={
"acompletion": True,
"metadata": {},
"proxy_server_request": {
"body": {
"instructions": "private-request-instructions",
"system": "private-system-canary",
"tools": [{"description": "private-tool-canary"}],
}
},
}
if blocked_phase == "hook"
else {"acompletion": True, "metadata": {}},
optional_params={},
)
logging_result: Final = (
ResponseCompletedEvent(type="response.completed", sequence_number=1, response=response)
if stream and isinstance(response, ResponsesAPIResponse)
else response
)
worker: Final = LoggingWorker(timeout=0.5)
worker.ensure_initialized_and_enqueue(logging_obj.async_success_handler(result=logging_result))
try:
await worker.flush()
line, context = recorded.get_nowait()
@ -13041,11 +13096,69 @@ async def test_slow_optional_logging_cannot_cancel_accounting(
event: Final = decode_spend_event(line)
assert isinstance(event, SpendEvent)
assert event.response is not None
assert Usage.model_validate(event.response["usage"]).total_tokens == 15
assert isinstance(event.response["usage"], dict)
assert event.response["usage"]["total_tokens"] == 15
assert recorded.empty()
assert cancelled.is_set() is (blocked_phase in ("hook", "delivery"))
assert ("private-marker" in line.decode()) is (blocked_phase not in ("hook", "redacted"))
assert (REDACTED_BY_LITELLM in line.decode()) is (blocked_phase in ("hook", "redacted"))
assert response.choices[0].message.content == "private-marker"
assert (REDACTED_BY_LITELLM in line.decode()) is (blocked_phase == "redacted")
assert ("private-instructions-canary" in line.decode()) is (
response_kind == "responses" and blocked_phase in ("neither", "delivery")
)
delivery: Final = deliveries.get_nowait()
assert not any(
value in delivery
for value in (
"private-request-instructions",
"private-system-canary",
"private-tool-canary",
)
)
assert ("private-instructions-canary" in delivery) is (
response_kind == "responses" and blocked_phase in ("neither", "delivery")
)
assert response.model_dump(exclude={"usage"}) == original_response
finally:
await worker.stop()
@pytest.mark.asyncio
async def test_invalid_accounting_after_a_hook_timeout_never_exports_unmasked_data(
monkeypatch: pytest.MonkeyPatch,
caplog: pytest.LogCaptureFixture,
) -> None:
from litellm.proxy.hooks.proxy_track_cost_callback import _ProxyDBLogger
exported: Final = asyncio.Event()
never: Final = asyncio.Event()
class InvalidMaskingHook(CustomLogger):
async def async_logging_hook(
self,
kwargs: dict[str, object],
result: object,
call_type: str,
) -> tuple[dict[str, object], object]:
kwargs["stream"] = "invalid"
await never.wait()
return kwargs, result
async def async_log_success_event(
self,
kwargs: dict[str, object],
response_obj: object,
start_time: datetime.datetime,
end_time: datetime.datetime,
) -> None:
exported.set()
monkeypatch.setattr(litellm, "_async_success_callback", [])
logging_obj: Final = _async_request_logging_obj("invalid-accounting", [InvalidMaskingHook(), _ProxyDBLogger()])
worker: Final = LoggingWorker(timeout=0.5)
worker.ensure_initialized_and_enqueue(logging_obj.async_success_handler(result=ModelResponse()))
try:
await worker.flush()
assert not exported.is_set()
assert "Cannot prepare body-free spend accounting after a logging hook timeout" in caplog.text
finally:
await worker.stop()