mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
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 <yucheng@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
ae6a4a2f2a
commit
09b694894d
5 changed files with 249 additions and 24 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue