From 09b694894d4b71fab0b6a194ada64689b0f59e08 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 11 Sep 2026 13:12:14 -0700 Subject: [PATCH] fix(datadog_llm_obs): keep tool call and result structure under redaction and emit tool output tokens (#40666) * fix(datadog_llm_obs): keep tool call and result structure under redaction and emit tool output tokens Under datadog_llm_observability_params.turn_off_message_logging the span kept only one role plus "redacted-by-litellm" per message, so Datadog showed Tool Call 0, Tool Result 0 and no tool output token data. The shared CustomLogger hook collapsed the messages before the callback ran, and the Datadog redaction then dropped tool_calls and tool_results. The Datadog callback now opts out of the shared message collapse (redacts_messages_itself) and redacts its own normalized messages, keeping roles, tool names, ids and types while replacing content, arguments and results. Tool result tokens are counted with litellm.token_counter before redaction and shipped as the tool_output_tokens metric. Other callbacks keep the inherited behavior. Resolves LIT-7545 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(datadog_llm_obs): drop explanatory docstrings from the redaction change Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(ui): regenerate schema.d.ts for the classifier descriptions changed in #40655 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yucheng Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/custom_logger.py | 8 +- .../integrations/datadog/datadog_llm_obs.py | 65 +++++-- litellm/types/integrations/datadog_llm_obs.py | 1 + .../datadog/test_datadog_llm_obs.py | 164 ++++++++++++++++-- .../test_redact_messages.py | 35 ++++ 5 files changed, 249 insertions(+), 24 deletions(-) diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index d445a3adf14..62ca6b0254e 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -886,12 +886,16 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac return LITELLM_METADATA_FIELD return OLD_LITELLM_METADATA_FIELD + def redacts_messages_itself(self) -> bool: + return False + def redact_standard_logging_payload_from_model_call_details(self, model_call_details: dict) -> dict: """ Redacts or excludes fields from StandardLoggingPayload before callbacks receive it. This method handles two features: - 1. turn_off_message_logging: When True, redacts messages and responses + 1. turn_off_message_logging: When True, redacts messages and responses (unless the callback + redacts them itself, see `redacts_messages_itself`) 2. standard_logging_payload_excluded_fields: Removes specified fields entirely Return a modified copy of the provided logging payload. @@ -921,7 +925,7 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac } # Handle turn_off_message_logging - redact messages and responses (if not already excluded) - if turn_off_message_logging: + if turn_off_message_logging and not self.redacts_messages_itself(): redacted_str: Final = "redacted-by-litellm" if "messages" not in (excluded_fields or ()) and standard_logging_object_copy.get("messages") is not None: diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py index 728bf41856f..c64a12c6d75 100644 --- a/litellm/integrations/datadog/datadog_llm_obs.py +++ b/litellm/integrations/datadog/datadog_llm_obs.py @@ -165,18 +165,54 @@ def _metadata_without_prompt_carriers(standard_logging_metadata: Mapping[str, An ) -def _redact_messages(messages: Sequence[Message]) -> tuple[Message, ...]: - """Each message's shape with its content replaced and tool payloads dropped; no message is invented.""" - return tuple( - { - "role": role if isinstance(role, str) and role in _SAFE_REDACTED_MESSAGE_ROLES else "", - "content": REDACTED_BY_LITELLM, - } - for message in messages - for role in (message.get("role", ""),) +def _safe_identifier(value: object) -> str: + return value if isinstance(value, str) else "" + + +def _redact_tool_call(tool_call: ToolCall) -> ToolCall: + return ToolCall( + name=_safe_identifier(tool_call.get("name")), + arguments=REDACTED_BY_LITELLM, + tool_id=_safe_identifier(tool_call.get("tool_id")), + type=_safe_identifier(tool_call.get("type")), ) +def _redact_tool_result(tool_result: ToolResult) -> ToolResult: + return ToolResult( + name=_safe_identifier(tool_result.get("name")), + result=REDACTED_BY_LITELLM, + tool_id=_safe_identifier(tool_result.get("tool_id")), + type=_safe_identifier(tool_result.get("type")), + ) + + +def _redact_message(message: Message) -> Message: + role: Final = message.get("role", "") + tool_calls: Final = message.get("tool_calls", ()) + tool_results: Final = message.get("tool_results", ()) + redacted: Final[Message] = { + "role": role if isinstance(role, str) and role in _SAFE_REDACTED_MESSAGE_ROLES else "", + "content": REDACTED_BY_LITELLM, + **({"tool_calls": tuple(_redact_tool_call(call) for call in tool_calls)} if tool_calls else {}), + **({"tool_results": tuple(_redact_tool_result(result) for result in tool_results)} if tool_results else {}), + } + return redacted + + +def _redact_messages(messages: Sequence[Message]) -> tuple[Message, ...]: + return tuple(_redact_message(message) for message in messages) + + +def _tool_output_tokens(messages: Sequence[Message], model: str) -> float | None: + results: Final = tuple( + result.get("result", "") for message in messages for result in message.get("tool_results", ()) + ) + if not results: + return None + return float(sum(litellm.token_counter(model=model, text=result) for result in results)) + + def _cost_dimension_tags( standard_logging_payload: StandardLoggingPayload, router_fields: Mapping[str, object] ) -> tuple[str, ...]: @@ -583,6 +619,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): standard_logging_payload=standard_logging_payload, call_type=standard_logging_payload.get("call_type"), ) + tool_output_tokens: Final = _tool_output_tokens(input_messages, standard_logging_payload.get("model") or "") input_meta: Final = InputMeta(messages=_redact_messages(input_messages) if redact_payload else input_messages) output_meta: Final = OutputMeta( messages=_redact_messages(output_messages) if redact_payload else output_messages @@ -618,7 +655,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): **({"tool_definitions": tool_definitions} if tool_definitions else {}), } - metrics: Final = self._assemble_metrics(standard_logging_payload) + metrics: Final = self._assemble_metrics(standard_logging_payload, tool_output_tokens) payload: Final[LLMObsPayload] = LLMObsPayload( parent_id=metadata_parent_id if metadata_parent_id else "undefined", @@ -676,6 +713,9 @@ class DataDogLLMObsLogger(CustomBatchLogger): ) return error_info + def redacts_messages_itself(self) -> bool: + return True + def _payload_logging_is_off(self, kwargs: Mapping[str, Any]) -> bool: return ( bool(self.turn_off_message_logging) @@ -683,7 +723,9 @@ class DataDogLLMObsLogger(CustomBatchLogger): or should_redact_message_logging(dict(kwargs)) ) - def _assemble_metrics(self, standard_logging_payload: StandardLoggingPayload) -> LLMMetrics: + def _assemble_metrics( + self, standard_logging_payload: StandardLoggingPayload, tool_output_tokens: float | None + ) -> LLMMetrics: """ Build the span metrics, including the prompt-cache counts LLM Obs charts cache savings from. @@ -721,6 +763,7 @@ class DataDogLLMObsLogger(CustomBatchLogger): else {} ), **({"reasoning_output_tokens": reasoning_output_tokens} if reasoning_output_tokens else {}), + **({"tool_output_tokens": tool_output_tokens} if tool_output_tokens is not None else {}), } return metrics diff --git a/litellm/types/integrations/datadog_llm_obs.py b/litellm/types/integrations/datadog_llm_obs.py index 17cf5831c96..f4fcabf53ed 100644 --- a/litellm/types/integrations/datadog_llm_obs.py +++ b/litellm/types/integrations/datadog_llm_obs.py @@ -87,6 +87,7 @@ class LLMMetrics(TypedDict, total=False): cache_write_input_tokens: ReadOnly[float] non_cached_input_tokens: ReadOnly[float] reasoning_output_tokens: ReadOnly[float] + tool_output_tokens: ReadOnly[float] class LLMObsPayload(TypedDict, total=False): diff --git a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py index 0555447e34f..a2f81091893 100644 --- a/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py +++ b/tests/test_litellm/integrations/datadog/test_datadog_llm_obs.py @@ -541,24 +541,166 @@ def _span_json(logger_under_test: DataDogLLMObsLogger, payload: dict[str, Any]) return json.loads(safe_dumps(span)) +SECRET_TOOL_RESULT: Final = '{"city": "Paris", "temp_c": 18, "account_secret": "SECRET-7545"}' +TOOL_CONVERSATION: Final[list[dict[str, Any]]] = [ + {"role": "user", "content": "secret prompt"}, + {"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]}, + {"role": "tool", "tool_call_id": "call_abc123", "content": SECRET_TOOL_RESULT}, +] + + +def _redacted_span_as_the_proxy_builds_it(payload: dict[str, Any]) -> dict[str, Any]: + logger_under_test = _redacting_logger(turn_off_message_logging=True) + return _span_json( + logger_under_test, logger_under_test.redact_standard_logging_payload_from_model_call_details(payload) + ) + + def test_redaction_keeps_the_conversation_shape_without_its_content() -> None: - """Roles and message count survive so the trace stays legible; contents and tool payloads do not.""" - result = _span_json( - _redacting_logger(turn_off_message_logging=True), + result = _redacted_span_as_the_proxy_builds_it( + build_payload( + messages=TOOL_CONVERSATION, + response_message={"role": "assistant", "content": "secret response", "tool_calls": [ASSISTANT_TOOL_CALL]}, + ) + ) + + redacted_call = { + "name": "get_weather", + "arguments": "redacted-by-litellm", + "tool_id": "call_abc123", + "type": "function", + } + assert result["meta"]["input"]["messages"] == [ + {"role": "user", "content": "redacted-by-litellm"}, + {"role": "assistant", "content": "redacted-by-litellm", "tool_calls": [redacted_call]}, + { + "role": "tool", + "content": "redacted-by-litellm", + "tool_results": [ + {"name": "get_weather", "result": "redacted-by-litellm", "tool_id": "call_abc123", "type": "function"} + ], + }, + ] + assert result["meta"]["output"]["messages"] == [ + {"role": "assistant", "content": "redacted-by-litellm", "tool_calls": [redacted_call]} + ] + serialized = safe_dumps(result) + assert "SECRET-7545" not in serialized + assert "Paris" not in serialized + assert "secret" not in serialized + + +def test_redaction_counts_tool_result_tokens_before_replacing_them() -> None: + payload = build_payload(messages=TOOL_CONVERSATION) + payload["standard_logging_object"]["model"] = "claude-sonnet-5" + + result = _redacted_span_as_the_proxy_builds_it(payload) + + expected_tokens = litellm.token_counter(model="claude-sonnet-5", text=SECRET_TOOL_RESULT) + assert expected_tokens > 0 + assert result["metrics"]["tool_output_tokens"] == float(expected_tokens) + assert result["metrics"]["input_tokens"] == 4447.0 + + +def test_tool_output_tokens_sum_every_result_in_the_request(logger: DataDogLLMObsLogger) -> None: + payload = build( + logger, + messages=[ + {"role": "tool", "tool_call_id": "call_1", "content": "one two three"}, + {"role": "tool", "tool_call_id": "call_2", "content": "four five six seven"}, + ], + ) + + assert payload["metrics"]["tool_output_tokens"] == float( + litellm.token_counter(text="one two three") + litellm.token_counter(text="four five six seven") + ) + + +def test_a_request_without_tool_results_reports_no_tool_output_tokens(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, messages=[{"role": "user", "content": "hi"}]) + + assert "tool_output_tokens" not in payload["metrics"] + assert "tool_output_tokens" not in _redacted_span_as_the_proxy_builds_it(build_payload())["metrics"] + + +def test_a_tool_that_returned_nothing_still_counts_as_zero_tool_output_tokens(logger: DataDogLLMObsLogger) -> None: + payload = build(logger, messages=[{"role": "tool", "tool_call_id": "call_1", "content": ""}]) + + assert payload["metrics"]["tool_output_tokens"] == 0.0 + + +def test_redaction_keeps_anthropic_tool_blocks_as_structure_only() -> None: + result = _redacted_span_as_the_proxy_builds_it( build_payload( messages=[ - {"role": "user", "content": "secret prompt"}, - {"role": "assistant", "content": None, "tool_calls": [ASSISTANT_TOOL_CALL]}, - ], - response_message={"role": "assistant", "content": "secret response"}, - ), + { + "role": "assistant", + "content": [ + {"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Paris"}} + ], + }, + { + "role": "user", + "content": [{"type": "tool_result", "tool_use_id": "toolu_1", "content": SECRET_TOOL_RESULT}], + }, + ] + ) ) assert result["meta"]["input"]["messages"] == [ - {"role": "user", "content": "redacted-by-litellm"}, - {"role": "assistant", "content": "redacted-by-litellm"}, + { + "role": "assistant", + "content": "redacted-by-litellm", + "tool_calls": [ + {"name": "get_weather", "arguments": "redacted-by-litellm", "tool_id": "toolu_1", "type": "tool_use"} + ], + }, + { + "role": "user", + "content": "redacted-by-litellm", + "tool_results": [ + {"name": "get_weather", "result": "redacted-by-litellm", "tool_id": "toolu_1", "type": "function"} + ], + }, ] - assert result["meta"]["output"]["messages"] == [{"role": "assistant", "content": "redacted-by-litellm"}] + assert "Paris" not in safe_dumps(result) + + +def test_redaction_blanks_tool_identifiers_that_are_not_strings() -> None: + result = _redacted_span_as_the_proxy_builds_it( + build_payload( + messages=[ + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": {"leak": "SECRET-7545"}, "type": ["SECRET-7545"], "function": {"name": ["SECRET-7545"]}} + ], + } + ] + ) + ) + + assert result["meta"]["input"]["messages"][0]["tool_calls"] == [ + {"name": "", "arguments": "redacted-by-litellm", "tool_id": "", "type": ""} + ] + assert "SECRET-7545" not in safe_dumps(result) + + +def test_the_shared_hook_still_strips_what_redaction_governs_besides_messages() -> None: + payload = build_payload(messages=TOOL_CONVERSATION) + payload["standard_logging_object"]["classifier_input"] = {"system": "SECRET-7545"} + logger_under_test = _redacting_logger(turn_off_message_logging=True) + + with patch.object( # test-quality-ok: the hook reads this module global with no injection seam + litellm, "standard_logging_payload_excluded_fields", ["response"] + ): + redacted = logger_under_test.redact_standard_logging_payload_from_model_call_details(payload) + + assert "classifier_input" not in redacted["standard_logging_object"] + assert "response" not in redacted["standard_logging_object"] + assert redacted["standard_logging_object"]["messages"] == TOOL_CONVERSATION + assert payload["standard_logging_object"]["classifier_input"] == {"system": "SECRET-7545"} def test_redaction_drops_unrecognized_and_malformed_message_roles() -> None: diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index 76092e96307..584a3ac471c 100644 --- a/tests/test_litellm/litellm_core_utils/test_redact_messages.py +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -926,3 +926,38 @@ def test_classifier_callback_redaction_preserves_exclusions(monkeypatch: pytest. assert failure_payload["response"]["choices"][0]["message"]["content"] == "redacted-by-litellm" assert payload["classifier_input"] == {"system": "private rubric"} assert payload["response"]["choices"][0]["message"]["content"] == "private answer" + + +class _SelfRedactingLogger(CustomLogger): + def redacts_messages_itself(self) -> bool: + return True + + +@pytest.mark.parametrize("logger", [CustomLogger(), _SelfRedactingLogger()], ids=["default", "redacts_itself"]) +def test_field_exclusion_alone_leaves_messages_and_responses_intact(monkeypatch: pytest.MonkeyPatch, logger: CustomLogger) -> None: + monkeypatch.setattr(litellm, "standard_logging_payload_excluded_fields", ["model"]) + payload: Final = { + "messages": [{"role": "user", "content": "private prompt"}], + "response": {"choices": [{"message": {"content": "private answer"}}]}, + "model": "classifier", + } + stored: Final = logger.redact_standard_logging_payload_from_model_call_details({"standard_logging_object": payload})[ + "standard_logging_object" + ] + assert stored == {"messages": payload["messages"], "response": payload["response"]} + + +def test_a_callback_that_redacts_itself_keeps_its_messages_but_not_the_classifier_audit() -> None: + payload: Final = { + "classifier_input": {"system": "private rubric"}, + "messages": [{"role": "user", "content": "private prompt"}], + "response": {"choices": [{"message": {"content": "private answer"}}]}, + } + logger: Final = _SelfRedactingLogger() + logger.turn_off_message_logging = True + stored: Final = logger.redact_standard_logging_payload_from_model_call_details({"standard_logging_object": payload})[ + "standard_logging_object" + ] + assert "classifier_input" not in stored + assert stored["messages"] == payload["messages"] + assert stored["response"] == payload["response"]