Merge 60f469d63f into 527d882508

litellm/integrations/arize/_utils.py

@@ -10,10 +10,12 @@ from litellm.integrations.opentelemetry_utils.base_otel_llm_obs_attributes impor
     BaseLLMObsOTELAttributes,
     safe_set_attribute,
 )
+from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.litellm_core_utils.redact_messages import (
     should_redact_message_logging,
 )
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
+from litellm.types.decisions import DecisionsResponse
 from litellm.types.utils import CallTypes, StandardLoggingMCPToolCall, StandardLoggingPayload
 
 if TYPE_CHECKING:
@@ -33,6 +35,23 @@ from litellm.integrations._types.open_inference import (
 _JSON_OBJECT: Final = TypeAdapter(dict[str, object], config=ConfigDict(hide_input_in_errors=True))
 
 
+def should_redact_span_content(kwargs: Mapping[str, object]) -> bool:
+    params: Final = as_str_mapping(kwargs.get("litellm_params")) or {}
+    routed_metadata: Final = as_str_mapping(params.get("litellm_metadata"))
+    if routed_metadata is None or "headers" in routed_metadata:
+        return should_redact_message_logging(dict(kwargs))
+    metadata: Final = as_str_mapping(params.get("metadata")) or {}
+    return should_redact_message_logging(
+        {
+            **kwargs,
+            "litellm_params": {
+                **params,
+                "litellm_metadata": {**routed_metadata, "headers": metadata.get("headers", {})},
+            },
+        }
+    )
+
+
 class ArizeOTELAttributes(BaseLLMObsOTELAttributes):
     @staticmethod
     @override
@@ -253,11 +272,17 @@ def _set_usage_outputs(span: "Span", response_obj, span_attrs):
         return
 
     safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_TOTAL, _safe_get(usage, "total_tokens"))
-    completion_tokens: Final = _safe_get(usage, "completion_tokens") or _safe_get(usage, "output_tokens")
-    if completion_tokens:
+    reported_completion_tokens: Final = _safe_get(usage, "completion_tokens")
+    completion_tokens: Final = (
+        reported_completion_tokens if reported_completion_tokens is not None else _safe_get(usage, "output_tokens")
+    )
+    if completion_tokens is not None:
         safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_COMPLETION, completion_tokens)
-    prompt_tokens: Final = _safe_get(usage, "prompt_tokens") or _safe_get(usage, "input_tokens")
-    if prompt_tokens:
+    reported_prompt_tokens: Final = _safe_get(usage, "prompt_tokens")
+    prompt_tokens: Final = (
+        reported_prompt_tokens if reported_prompt_tokens is not None else _safe_get(usage, "input_tokens")
+    )
+    if prompt_tokens is not None:
         safe_set_attribute(span, span_attrs.LLM_TOKEN_COUNT_PROMPT, prompt_tokens)
 
     # Reasoning tokens live in `completion_tokens_details` for Chat Completions
@@ -604,6 +629,10 @@ def _coerce_response_obj_for_attrs(response_obj):
       `id`, `model`, and `usage`. On failure the original object is returned
       so behavior is no worse than today.
     """
+    if isinstance(response_obj, DecisionsResponse):
+        from litellm.cost_calculator import get_usage_object
+
+        return {**response_obj.model_dump(), "usage": get_usage_object(response_obj)}
     if response_obj is None or hasattr(response_obj, "get"):
         return response_obj
     dumped: Final = _to_plain_dict(response_obj)
@@ -1117,7 +1146,7 @@ def _maybe_set_mcp_tool_attrs(
     if tool_name:
         safe_set_attribute(span, SpanAttributes.TOOL_NAME, tool_name)
 
-    if should_redact_message_logging(kwargs):  # pyright: ignore[reportArgumentType]  # reads, never mutates
+    if should_redact_span_content(kwargs):
         return
 
     arguments: Final[object] = mcp_meta.get("arguments")

litellm/integrations/langfuse/langfuse_otel.py

@@ -27,6 +27,7 @@ from litellm.litellm_core_utils.core_helpers import as_str_mapping
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
 from litellm.secret_managers.main import str_to_bool
+from litellm.types.decisions import DecisionsResponse
 from litellm.types.integrations.langfuse_otel import (
     LangfuseSpanAttributes,
 )
@@ -157,7 +158,7 @@ class LangfuseOtelLogger(OpenTelemetry):
         super().__init__(config=config, *args, **kwargs)
 
     @staticmethod
-    def set_langfuse_otel_attributes(span: Span, kwargs, response_obj):
+    def set_langfuse_otel_attributes(span: Span, kwargs, response_obj, *, capture_decisions_content: bool = False):
         """
         Sets OpenTelemetry span attributes for Langfuse observability.
         Uses the same attribute setting logic as Arize Phoenix for consistency.
@@ -169,7 +170,9 @@ class LangfuseOtelLogger(OpenTelemetry):
         #########################################################
         # Set Langfuse specific attributes
         #########################################################
-        LangfuseOtelLogger._set_langfuse_specific_attributes(span=span, kwargs=kwargs, response_obj=response_obj)
+        LangfuseOtelLogger._set_langfuse_specific_attributes(
+            span=span, kwargs=kwargs, response_obj=response_obj, capture_decisions_content=capture_decisions_content
+        )
 
     @staticmethod
     def _extract_langfuse_metadata(kwargs: Mapping[str, Any]) -> dict:
@@ -244,7 +247,11 @@ class LangfuseOtelLogger(OpenTelemetry):
                 safe_set_attribute(span, enum_attr.value, value)
 
     @staticmethod
-    def _observation_output(response_obj: _Gettable | None) -> str | None:
+    def _observation_output(
+        response_obj: _Gettable | DecisionsResponse | None, *, capture_decisions_content: bool = False
+    ) -> str | None:
+        if isinstance(response_obj, DecisionsResponse):
+            return safe_dumps(dict(response_obj.answers)) if capture_decisions_content else None
         if not response_obj or not hasattr(response_obj, "get"):
             return None
         return _extract_output_items(response_obj) or _extract_choices_output(response_obj)
@@ -282,7 +289,9 @@ class LangfuseOtelLogger(OpenTelemetry):
         return tuple(dict.fromkeys(candidates))
 
     @staticmethod
-    def _set_langfuse_specific_attributes(span: Span, kwargs: Mapping[str, object], response_obj):
+    def _set_langfuse_specific_attributes(
+        span: Span, kwargs: Mapping[str, object], response_obj, *, capture_decisions_content: bool = False
+    ):
         """
         Sets Langfuse specific metadata attributes onto the OTEL span.
 
@@ -291,7 +300,7 @@ class LangfuseOtelLogger(OpenTelemetry):
         values (lists/dicts) are serialised to JSON strings for OTEL
         compatibility.
         """
-        from litellm.integrations.arize._utils import safe_set_attribute
+        from litellm.integrations.arize._utils import safe_set_attribute, should_redact_span_content
         from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 
         dynamic_params: Final = as_str_mapping(kwargs.get("standard_callback_dynamic_params"))
@@ -323,7 +332,10 @@ class LangfuseOtelLogger(OpenTelemetry):
                 serialize_langfuse_datetime(completion_start_time),
             )
 
-        output_json: Final = LangfuseOtelLogger._observation_output(response_obj)
+        output_json: Final = LangfuseOtelLogger._observation_output(
+            response_obj,
+            capture_decisions_content=capture_decisions_content and not should_redact_span_content(kwargs),
+        )
         if output_json is not None:
             safe_set_attribute(span, LangfuseSpanAttributes.OBSERVATION_OUTPUT.value, output_json)
 

litellm/integrations/opentelemetry.py

@@ -2434,7 +2434,14 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
                     LangfuseOtelLogger,
                 )
 
-                LangfuseOtelLogger.set_langfuse_otel_attributes(span, kwargs, response_obj)
+                LangfuseOtelLogger.set_langfuse_otel_attributes(
+                    span,
+                    kwargs,
+                    response_obj,
+                    capture_decisions_content=(
+                        self._capture_in_span() and self.message_logging and not self.turn_off_message_logging
+                    ),
+                )
                 return
             elif self.callback_name == "weave_otel":
                 from litellm.integrations.weave.weave_otel import (

tests/unit/integrations/arize/test_arize_utils.py

@@ -16,9 +16,15 @@ from litellm.integrations._types.open_inference import (
     SpanAttributes,
     ToolCallAttributes,
 )
-from litellm.integrations.arize._utils import _coerce_response_obj_for_attrs, _parse_passthrough_response
+from litellm.integrations.arize._utils import (
+    ArizeOTELAttributes,
+    _coerce_response_obj_for_attrs,
+    _parse_passthrough_response,
+    set_attributes,
+)
 from litellm.integrations.arize.arize import ArizeLogger
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.types.decisions import DecisionsResponse
 from litellm.types.utils import Choices, ModelResponse, StandardCallbackDynamicParams
 
 
@@ -1605,3 +1611,34 @@ def test_parse_passthrough_response_reads_only_json_objects_from_text(
     raw: object, coerced: object, kwargs: dict[str, object], expected: dict[str, object] | None
 ):
     assert _parse_passthrough_response(raw, coerced, kwargs) == expected
+
+
+@pytest.mark.parametrize(("input_tokens", "output_tokens"), [(12, 1), (12, 0), (0, 0)])
+def test_systemone_usage_emits_total_and_zero_token_counts(input_tokens: int, output_tokens: int) -> None:
+    response: Final = DecisionsResponse.model_validate(
+        {
+            "answers": {"is_allowed": {"type": "noul", "noul": 0.91}},
+            "usage": {"input_tokens": input_tokens, "output_tokens": output_tokens},
+        }
+    )
+    original: Final = response.model_dump()
+    provider: Final = TracerProvider()
+    span: Final = provider.get_tracer("systemone-usage-tests").start_span("litellm_request")
+    try:
+        set_attributes(
+            span,
+            {"call_type": "asystemone", "standard_logging_object": {"metadata": {}}, "litellm_params": {}},
+            response,
+            ArizeOTELAttributes,
+        )
+        attributes: Final = dict(getattr(span, "attributes", None) or {})
+
+        assert {key: value for key, value in attributes.items() if key.startswith("llm.token_count.")} == {
+            "llm.token_count.prompt": input_tokens,
+            "llm.token_count.completion": output_tokens,
+            "llm.token_count.total": input_tokens + output_tokens,
+        }
+        assert response.model_dump() == original
+    finally:
+        span.end()
+        provider.shutdown()

tests/unit/integrations/test_langfuse_otel.py

@@ -26,6 +26,7 @@ from litellm.integrations.opentelemetry import LITELLM_PROXY_REQUEST_SPAN_NAME,
 from litellm.integrations.otel.model.destination import OtelDestination
 from litellm.integrations.otel.plumbing.context import set_request_destinations
 from litellm.proxy._types import UserAPIKeyAuth
+from litellm.types.decisions import DecisionsResponse
 from litellm.types.llms.openai import ResponsesAPIResponse
 from litellm.types.utils import Choices, ModelResponse, StandardCallbackDynamicParams
 
@@ -1587,3 +1588,96 @@ def test_a_request_without_a_langfuse_destination_still_exports_through_the_lega
     assert _spans_the_legacy_logger_exports((_LANGFUSE_DESTINATION.model_copy(update={"callback_name": "arize"}),)) == (
         logged
     )
+
+
+def _systemone_attributes(
+    response: DecisionsResponse,
+    kwargs: Mapping[str, object],
+    *,
+    capture_mode: str | None = None,
+    turn_off_message_logging: bool = False,
+    message_logging: bool = True,
+) -> Mapping[str, object]:
+    exporter: Final = InMemorySpanExporter()
+    provider: Final = TracerProvider()
+    provider.add_span_processor(SimpleSpanProcessor(exporter))
+    logger: Final = LangfuseOtelLogger(
+        config=OpenTelemetryConfig(skip_set_global=True, capture_message_content=capture_mode),
+        callback_name="langfuse_otel",
+        tracer_provider=provider,
+        turn_off_message_logging=turn_off_message_logging,
+        message_logging=message_logging,
+    )
+    span: Final = provider.get_tracer("systemone-langfuse-tests").start_span("litellm_request")
+    try:
+        logger.set_attributes(span, dict(kwargs), response)
+        span.end()
+        return dict(exporter.get_finished_spans()[0].attributes or {})
+    finally:
+        provider.shutdown()
+
+
+@pytest.mark.parametrize(
+    ("capture_mode", "captures_answers"),
+    [(None, True), ("span_only", True), ("span_and_event", True), ("no_content", False), ("event_only", False)],
+)
+@pytest.mark.parametrize("turn_off_message_logging", [False, True])
+def test_systemone_answers_respect_message_content_capture(
+    capture_mode: str | None,
+    captures_answers: bool,
+    turn_off_message_logging: bool,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.setattr(litellm, "turn_off_message_logging", turn_off_message_logging)
+    response: Final = DecisionsResponse.model_validate(
+        {"answers": {"is_allowed": {"type": "noul", "noul": 0.91}}, "usage": {"input_tokens": 12, "output_tokens": 0}}
+    )
+    original: Final = response.model_dump()
+    attributes: Final = _systemone_attributes(
+        response,
+        {"call_type": "asystemone", "standard_logging_object": {"metadata": {}}, "litellm_params": {}},
+        capture_mode=capture_mode,
+    )
+
+    if captures_answers and not turn_off_message_logging:
+        assert json.loads(str(attributes["langfuse.observation.output"])) == {
+            "is_allowed": {"type": "noul", "noul": 0.91}
+        }
+        assert attributes["langfuse.trace.output"] == attributes["langfuse.observation.output"]
+    else:
+        assert "langfuse.observation.output" not in attributes
+        assert "langfuse.trace.output" not in attributes
+        assert "is_allowed" not in json.dumps(dict(attributes))
+    assert {key: value for key, value in attributes.items() if key.startswith("llm.token_count.")} == {
+        "llm.token_count.prompt": 12,
+        "llm.token_count.completion": 0,
+        "llm.token_count.total": 12,
+    }
+    assert response.model_dump() == original
+
+
+@pytest.mark.parametrize("opt_out", ["request", "header", "callback", "legacy"])
+@pytest.mark.parametrize("routed_metadata", [False, True])
+def test_systemone_answers_respect_message_logging_opt_outs(opt_out: str, routed_metadata: bool) -> None:
+    response: Final = DecisionsResponse.model_validate(
+        {"answers": {"is_allowed": {"type": "noul", "noul": 0.91}}, "usage": {"input_tokens": 12, "output_tokens": 0}}
+    )
+    attributes: Final = _systemone_attributes(
+        response,
+        {
+            "call_type": "asystemone",
+            "standard_logging_object": {"metadata": {}},
+            "standard_callback_dynamic_params": {"turn_off_message_logging": True} if opt_out == "request" else {},
+            "litellm_params": {
+                "metadata": {"headers": {"litellm-enable-message-redaction": "true"}} if opt_out == "header" else {},
+                **({"litellm_metadata": {"model_group": "decision-test"}} if routed_metadata else {}),
+            },
+        },
+        capture_mode="span_only",
+        turn_off_message_logging=opt_out == "callback",
+        message_logging=opt_out != "legacy",
+    )
+
+    assert "langfuse.observation.output" not in attributes
+    assert "langfuse.trace.output" not in attributes
+    assert "is_allowed" not in json.dumps(dict(attributes))