diff --git a/litellm/integrations/posthog.py b/litellm/integrations/posthog.py index db9610a5a3c..dda9ac6420e 100644 --- a/litellm/integrations/posthog.py +++ b/litellm/integrations/posthog.py @@ -285,9 +285,11 @@ class PostHogLogger(CustomBatchLogger): end_user: Final = self._safe_get(standard_logging_object, "end_user") if end_user: return str(end_user) - trace_id: Final = self._safe_get(standard_logging_object, "trace_id") - if trace_id: - return str(trace_id) + session_id: Final = self._safe_get(standard_logging_object, "session_id") or self._safe_get( + standard_logging_object, "trace_id" + ) + if session_id: + return str(session_id) return self._safe_uuid() diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 6ba06919e00..bb30eb5b324 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -5050,33 +5050,47 @@ class StandardLoggingPayloadSetup: else: return end_time_float - start_time_float + @staticmethod + def _get_standard_logging_payload_session_id( + logging_obj: Logging, + litellm_params: dict, + ) -> str: + """ + Returns the session id for this request, preferring session-scoped values. + + This links multiple requests made in a single session + """ + metadata: Final = litellm_params.get("metadata") or {} + for candidate in ( + litellm_params.get("litellm_session_id"), + metadata.get("session_id"), + litellm_params.get("litellm_trace_id"), + metadata.get("trace_id"), + ): + if candidate: + return str(candidate) + return logging_obj.litellm_trace_id + @staticmethod def _get_standard_logging_payload_trace_id( logging_obj: Logging, litellm_params: dict, ) -> str: """ - Returns the `litellm_trace_id` for this request + Returns the trace id for this request, preferring trace-scoped values. - This helps link sessions when multiple requests are made in a single session + This groups the LLM calls belonging to one overall request, such as fallbacks + and retries """ - dynamic_litellm_session_id: Final = litellm_params.get("litellm_session_id") - dynamic_litellm_trace_id: Final = litellm_params.get("litellm_trace_id") - - # Note: we recommend using `litellm_session_id` for session tracking - # `litellm_trace_id` is an internal litellm param - if dynamic_litellm_session_id: - return str(dynamic_litellm_session_id) - elif dynamic_litellm_trace_id: - return str(dynamic_litellm_trace_id) - # Fallback: use metadata.session_id or metadata.trace_id for call chaining metadata: Final = litellm_params.get("metadata") or {} - metadata_session_id: Final = metadata.get("session_id") - metadata_trace_id: Final = metadata.get("trace_id") - if metadata_session_id: - return str(metadata_session_id) - if metadata_trace_id: - return str(metadata_trace_id) + for candidate in ( + litellm_params.get("litellm_trace_id"), + metadata.get("trace_id"), + litellm_params.get("litellm_session_id"), + metadata.get("session_id"), + ): + if candidate: + return str(candidate) return logging_obj.litellm_trace_id @staticmethod @@ -5386,6 +5400,10 @@ def get_standard_logging_object_payload( logging_obj=logging_obj, litellm_params=litellm_params, ), + session_id=StandardLoggingPayloadSetup._get_standard_logging_payload_session_id( + logging_obj=logging_obj, + litellm_params=litellm_params, + ), call_type=call_type or "", cache_hit=cache_hit, stream=stream, diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 8d2569b2229..2e7db4b7e57 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -481,8 +481,10 @@ def _get_session_id_for_spend_log( """ from litellm._uuid import uuid - if standard_logging_payload is not None and standard_logging_payload.get("trace_id") is not None: - return str(standard_logging_payload.get("trace_id")) + if standard_logging_payload is not None: + session_id: Final = standard_logging_payload.get("session_id") or standard_logging_payload.get("trace_id") + if session_id is not None: + return str(session_id) # Users can dynamically set the trace_id for each request by passing `litellm_trace_id` in kwargs if kwargs.get("litellm_trace_id") is not None: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 8b2584b0efe..79700c82d0a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3127,6 +3127,7 @@ class StandardAuditLogPayload(TypedDict): class StandardLoggingPayload(TypedDict): id: str trace_id: str # Trace multiple LLM calls belonging to same overall request (e.g. fallbacks/retries) + session_id: str litellm_call_id: str | None # UUID returned in x-litellm-call-id response header call_type: str stream: bool | None diff --git a/tests/logging_callback_tests/test_standard_logging_payload.py b/tests/logging_callback_tests/test_standard_logging_payload.py index f29b245b3be..7ca6277ee54 100644 --- a/tests/logging_callback_tests/test_standard_logging_payload.py +++ b/tests/logging_callback_tests/test_standard_logging_payload.py @@ -509,6 +509,59 @@ def test_get_standard_logging_payload_trace_id(): assert isinstance(result, str) +def test_trace_id_and_session_id_are_resolved_independently(): + """A caller setting both must get their trace_id as the trace and their session_id as the session.""" + from unittest.mock import MagicMock + + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_trace_id = "fallback-id" + + litellm_params = {"metadata": {"trace_id": "caller-trace", "session_id": "caller-session"}} + + assert ( + StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + == "caller-trace" + ) + assert ( + StandardLoggingPayloadSetup._get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params=litellm_params + ) + == "caller-session" + ) + + +def test_trace_id_and_session_id_fall_back_to_each_other(): + """Either id alone still populates both, so existing single-id callers are unaffected.""" + from unittest.mock import MagicMock + + mock_logging_obj = MagicMock() + mock_logging_obj.litellm_trace_id = "fallback-id" + + session_only = {"litellm_session_id": "s-1"} + assert ( + StandardLoggingPayloadSetup._get_standard_logging_payload_trace_id( + logging_obj=mock_logging_obj, litellm_params=session_only + ) + == "s-1" + ) + assert ( + StandardLoggingPayloadSetup._get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params=session_only + ) + == "s-1" + ) + + trace_only = {"litellm_trace_id": "t-1"} + assert ( + StandardLoggingPayloadSetup._get_standard_logging_payload_session_id( + logging_obj=mock_logging_obj, litellm_params=trace_only + ) + == "t-1" + ) + + def test_truncate_standard_logging_payload(): """ 1. original messages, response, and error_str should NOT BE MODIFIED, since these are from kwargs diff --git a/tests/test_litellm/integrations/test_posthog.py b/tests/test_litellm/integrations/test_posthog.py new file mode 100644 index 00000000000..00f4c2daac1 --- /dev/null +++ b/tests/test_litellm/integrations/test_posthog.py @@ -0,0 +1,13 @@ + + +def test_distinct_id_prefers_session_id_over_trace_id(): + """PostHog identifies a person, so it needs session grouping rather than the trace.""" + from litellm.integrations.posthog import PostHogLogger + + logger = PostHogLogger.__new__(PostHogLogger) + payload = {"trace_id": "per-trace", "session_id": "per-session"} + + assert logger._get_distinct_id(standard_logging_object=payload, kwargs={}) == "per-session" + + legacy = {"trace_id": "only-trace"} + assert logger._get_distinct_id(standard_logging_object=legacy, kwargs={}) == "only-trace" diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 9eb45c399db..ca9a37cb8d2 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2959,3 +2959,16 @@ def test_user_traffic_carries_no_internal_call_origin(): ) metadata = json.loads(payload["metadata"]) assert metadata["internal_call_origin"] is None + + +def test_spend_log_session_id_prefers_session_over_trace(): + """SpendLogs.session_id must keep session semantics now that trace_id is trace-first.""" + from litellm.proxy.spend_tracking.spend_tracking_utils import ( + _get_session_id_for_spend_log, + ) + + payload = {"trace_id": "caller-trace", "session_id": "caller-session"} + assert _get_session_id_for_spend_log(kwargs={}, standard_logging_payload=payload) == "caller-session" + + legacy_payload = {"trace_id": "only-trace"} + assert _get_session_id_for_spend_log(kwargs={}, standard_logging_payload=legacy_payload) == "only-trace"