diff --git a/litellm/integrations/arize/_utils.py b/litellm/integrations/arize/_utils.py index f719dacd148..9ec4922c41f 100644 --- a/litellm/integrations/arize/_utils.py +++ b/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") diff --git a/litellm/integrations/langfuse/langfuse_otel.py b/litellm/integrations/langfuse/langfuse_otel.py index ee75e6501e8..2f136329c0a 100644 --- a/litellm/integrations/langfuse/langfuse_otel.py +++ b/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) diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index a9bd93e6420..88e656b0550 100644 --- a/litellm/integrations/opentelemetry.py +++ b/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 ( diff --git a/tests/unit/integrations/arize/test_arize_utils.py b/tests/unit/integrations/arize/test_arize_utils.py index 31699359bbe..3b3d8d7eaf0 100644 --- a/tests/unit/integrations/arize/test_arize_utils.py +++ b/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() diff --git a/tests/unit/integrations/test_langfuse_otel.py b/tests/unit/integrations/test_langfuse_otel.py index a2803719c75..2e69628d81f 100644 --- a/tests/unit/integrations/test_langfuse_otel.py +++ b/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))