diff --git a/litellm/integrations/langfuse/field_cap.py b/litellm/integrations/langfuse/field_cap.py new file mode 100644 index 00000000000..1308691206e --- /dev/null +++ b/litellm/integrations/langfuse/field_cap.py @@ -0,0 +1,156 @@ +"""Per-field caps that keep huge payloads out of Langfuse's ClickHouse rows. + +Every string bound for a span attribute is capped before export, mirroring the +agent-runtime capture discipline: 32 KiB per field, cut on a code-point boundary, +marked with ``...(truncated)``. The budget is the ``ensure_ascii``-escaped width, +which is what ``safe_dumps`` actually emits, so CJK and emoji content cannot +inflate a capped field past the limit. Cycles and over-deep nesting resolve to +the same markers ``safe_dumps`` would have produced, so a hostile payload loses +its content but not its event. The 413 backstop in ``langfuse_sdk`` stays as the +last resort; this makes it rare. +""" + +import json +import os +from collections.abc import Mapping, Sequence +from itertools import accumulate +from typing import ( + Final, + cast, # noqa: TID251 # identity-preserving type laundering for isinstance-narrowed user payloads +) + +from litellm._logging import verbose_logger +from litellm.constants import DEFAULT_MAX_RECURSE_DEPTH + +DEFAULT_MAX_FIELD_BYTES: Final = 32 * 1024 +MAX_FIELD_BYTES_METADATA_KEY: Final = "langfuse_max_field_bytes" +_MAX_FIELD_BYTES_ENV: Final = "LANGFUSE_MAX_FIELD_BYTES" +_TRUNCATION_SUFFIX: Final = "...(truncated)" +_MARKER_SERIALIZED_WIDTH: Final = len(_TRUNCATION_SUFFIX) + 2 +_CYCLE_MARKER: Final = "CircularReference Detected" +_DEPTH_MARKER: Final = "MaxDepthExceeded" + + +def resolve_max_field_bytes(override: object = None) -> int: + """Bytes allowed per string field; a value of ``0`` or less disables capping. + + Precedence: the per-request ``langfuse_max_field_bytes`` metadata value, then + ``LANGFUSE_MAX_FIELD_BYTES``, then the default. + """ + raw: Final = override if override is not None else os.environ.get(_MAX_FIELD_BYTES_ENV) + parsed: Final = _parse_byte_count(raw) + if parsed is not None: + return parsed + if raw is not None: + verbose_logger.warning("Invalid langfuse max field bytes %r; falling back to %d", raw, DEFAULT_MAX_FIELD_BYTES) + return DEFAULT_MAX_FIELD_BYTES + + +def _parse_byte_count(raw: object) -> int | None: + if isinstance(raw, bool) or raw is None: + return None + if isinstance(raw, int): + return raw + if isinstance(raw, float) and raw.is_integer(): + return int(raw) + if isinstance(raw, str): + try: + return int(raw.strip()) + except ValueError: + return None + return None + + +def cap_payload(value: object, max_field_bytes: int) -> object: + """Return ``value`` with every string capped at ``max_field_bytes`` serialized bytes. + + Plain containers are rebuilt capped; anything else passes through untouched. A + disabled limit (``max_field_bytes <= 0``) returns ``value`` itself. + """ + if max_field_bytes <= 0: + return value + return _cap(value, max_field_bytes, frozenset(), 0) + + +def _cap(value: object, limit: int, seen: frozenset[int], depth: int) -> object: + if depth > DEFAULT_MAX_RECURSE_DEPTH: + return _DEPTH_MARKER + if isinstance(value, str): + if _serialized_width(value) <= limit: + return value + if limit < _MARKER_SERIALIZED_WIDTH: + return "" + return _serialized_prefix(value, limit - _MARKER_SERIALIZED_WIDTH) + _TRUNCATION_SUFFIX + if id(value) in seen: + return _CYCLE_MARKER + rebuilt: Final = _capped_container(value, limit, seen, depth) + return value if rebuilt is None else rebuilt + + +def _capped_container(value: object, limit: int, seen: frozenset[int], depth: int) -> object | None: + """``value`` rebuilt with every string capped, or None when it is not a plain container.""" + identity: Final = id(value) + descended: Final = seen | {identity} + child_depth: Final = depth + 1 + mapping: Final = _as_mapping(value) + if mapping is not None: + return {key: _cap(item, limit, descended, child_depth) for key, item in mapping.items()} + items: Final = _as_object_list(value) + if items is not None: + return [_cap(item, limit, descended, child_depth) for item in items] + members: Final = _as_object_tuple(value) + if members is not None: + return tuple(_cap(item, limit, descended, child_depth) for item in members) + return None + + +def _serialized_width(value: str) -> int: + """Bytes ``safe_dumps`` emits for this string: the quotes plus the ``ensure_ascii`` escapes.""" + return len(json.dumps(value)) + + +def _serialized_prefix(value: str, budget: int) -> str: + """Longest head of ``value`` whose ``ensure_ascii``-escaped form fits ``budget`` bytes. + + ``safe_dumps`` escapes non-ASCII as ``\\uXXXX`` (6 bytes, 12 per astral code point), so a raw + byte cap understates what reaches ClickHouse; the cut budgets the escaped form instead. + Escapes are self-contained per code point, so cutting between code points is always safe. + """ + running_widths: Final = accumulate(_escaped_width(char) for char in value) + cut: Final = next((index for index, total in enumerate(running_widths) if total > budget), len(value)) + return value[:cut] + + +def _escaped_width(char: str) -> int: + """Bytes ``json.dumps`` emits for one code point under ``ensure_ascii``.""" + point: Final[int] = ord(char) + if point in (0x08, 0x09, 0x0A, 0x0C, 0x0D) or point in (0x22, 0x5C): + return 2 + if point < 0x20: + return 6 + if point < 0x7F: + return 1 + if point <= 0xFFFF: + return 6 + return 12 + + +def _as_mapping(value: object) -> Mapping[object, object] | None: + """``value`` as a plain mapping, or None.""" + if isinstance(value, Mapping): + return cast(Mapping[object, object], value) # cast-ok: isinstance-checked above; preserves the payload identity + return None + + +def _as_object_list(value: object) -> Sequence[object] | None: + """``value`` as a list, or None.""" + if isinstance(value, list): + return cast("Sequence[object]", value) # cast-ok: isinstance-checked above; read-only view for iteration + return None + + +def _as_object_tuple(value: object) -> tuple[object, ...] | None: + """``value`` as a tuple, or None.""" + if isinstance(value, tuple): + return cast(tuple[object, ...], value) # cast-ok: isinstance-checked above; preserves the payload identity + return None diff --git a/litellm/integrations/langfuse/langfuse.py b/litellm/integrations/langfuse/langfuse.py index bd68ffe8cd6..4aeacc8c89f 100644 --- a/litellm/integrations/langfuse/langfuse.py +++ b/litellm/integrations/langfuse/langfuse.py @@ -15,6 +15,11 @@ from packaging.version import Version import litellm from litellm._logging import verbose_logger from litellm.constants import MAX_LANGFUSE_INITIALIZED_CLIENTS +from litellm.integrations.langfuse.field_cap import ( + MAX_FIELD_BYTES_METADATA_KEY, + cap_payload, + resolve_max_field_bytes, +) from litellm.integrations.langfuse.langfuse_mock_client import ( create_mock_langfuse_client, should_use_langfuse_mock, @@ -727,6 +732,9 @@ class LangFuseLogger: litellm_call_id=litellm_call_id, proxy_server_request=litellm_params.get("proxy_server_request"), ) + + max_field_bytes: Final = resolve_max_field_bytes(clean_metadata.pop(MAX_FIELD_BYTES_METADATA_KEY, None)) + session_id: Final = clean_metadata.pop("session_id", None) trace_name = cast(str | None, clean_metadata.pop("trace_name", None)) clean_metadata.pop("trace_id", None) @@ -745,15 +753,23 @@ class LangFuseLogger: ) # Apply custom masking function if provided - masked_input: Final[object] = ( - self._apply_masking_function(input, masking_function) - if masking_function is not None and callable(masking_function) - else input + masked_input: Final[object] = cap_payload( + cast( # cast-ok: unknown-typed logging payload; passes through cap_payload unchanged + "object", + self._apply_masking_function(input, masking_function) + if masking_function is not None and callable(masking_function) + else input, + ), + max_field_bytes, ) - masked_output: Final[object] = ( - self._apply_masking_function(output, masking_function) - if masking_function is not None and callable(masking_function) - else output + masked_output: Final[object] = cap_payload( + cast( # cast-ok: unknown-typed logging payload; passes through cap_payload unchanged + "object", + self._apply_masking_function(output, masking_function) + if masking_function is not None and callable(masking_function) + else output, + ), + max_field_bytes, ) clean_metadata = redact_user_api_key_info(metadata=clean_metadata) @@ -922,11 +938,17 @@ class LangFuseLogger: "input": masked_input if not mask_input else "redacted-by-litellm", "output": masked_output if not mask_output else "redacted-by-litellm", "cost_details": {"total": cost} if usage is not None and isinstance(cost, (int, float)) else None, - "metadata": { - **log_requester_metadata(redact_user_api_key_info(metadata=allowlisted_metadata)), # pyright: ignore[reportArgumentType] # TypedDict in, plain metadata dict out - **enrichments, - **_lookup_ids(litellm_call_id, response_obj), - }, + "metadata": cap_payload( + cast( # cast-ok: unknown-typed logging payload; passes through cap_payload unchanged + "object", + { + **log_requester_metadata(redact_user_api_key_info(metadata=allowlisted_metadata)), # pyright: ignore[reportArgumentType] # TypedDict in, plain metadata dict out + **enrichments, + **_lookup_ids(litellm_call_id, response_obj), + }, + ), + max_field_bytes, + ), "version": _optional_str(clean_metadata.pop("version", None)), } @@ -959,15 +981,47 @@ class LangFuseLogger: name=trace_params.get("name"), user_id=trace_params.get("user_id"), session_id=trace_params.get("session_id"), - version=trace_params.get("version"), + version=cap_payload( + cast( # cast-ok: unknown-typed payload + "object", + trace_params.get("version"), + ), + max_field_bytes, + ), release=trace_params.get("release"), - tags=trace_params.get("tags"), - metadata=trace_params.get("metadata"), + tags=cap_payload( + cast( # cast-ok: unknown-typed payload + "object", + trace_params.get("tags"), + ), + max_field_bytes, + ), + metadata=cap_payload( + cast( # cast-ok: unknown-typed payload + "object", + trace_params.get("metadata"), + ), + max_field_bytes, + ), public=trace_public, - input=None if generation_is_trace_root and trace_input == generation_params["input"] else trace_input, + input=None + if generation_is_trace_root and trace_input == generation_params["input"] + else cap_payload( + cast( # cast-ok: unknown-typed payload + "object", + trace_input, + ), + max_field_bytes, + ), output=None if generation_is_trace_root and trace_output == generation_params["output"] - else trace_output, + else cap_payload( + cast( # cast-ok: unknown-typed payload + "object", + trace_output, + ), + max_field_bytes, + ), ) generation_attributes: Final = observation_attributes( observation_type="generation", @@ -978,7 +1032,13 @@ class LangFuseLogger: status_message=generation_params.get("status_message"), version=generation_params["version"], model=model_name, - model_parameters=optional_params, + model_parameters=cap_payload( + cast( # cast-ok: unknown-typed payload + "object", + optional_params, + ), + max_field_bytes, + ), usage_details=usage_details, cost_details=generation_params["cost_details"], completion_start_time=kwargs.get("completion_start_time", None), @@ -997,10 +1057,24 @@ class LangFuseLogger: ) try: log_provider_specific_information_as_span( - tracing=self.tracing, parent=generation, enrichments=enrichments + tracing=self.tracing, + parent=generation, + enrichments=cast( # cast-ok: cap_payload returns object; parameter wants Mapping + "Mapping[str, Any]", + cap_payload( + cast( # cast-ok: unknown-typed payload + "object", + enrichments, + ), + max_field_bytes, + ), + ), ) self._log_guardrail_information_as_span( - tracing=self.tracing, parent=generation, standard_logging_object=standard_logging_object + tracing=self.tracing, + parent=generation, + standard_logging_object=standard_logging_object, + max_field_bytes=max_field_bytes, ) finally: generation.end(end_time) @@ -1148,6 +1222,7 @@ class LangFuseLogger: tracing: "LangfuseTracing", parent: "LangfuseObservation", standard_logging_object: StandardLoggingPayload | None, + max_field_bytes: int, ): """ Log guardrail information as a span @@ -1185,8 +1260,20 @@ class LangFuseLogger: start_time=guardrail_entry.get("start_time", None), attributes=observation_attributes( observation_type="span", - input=guardrail_entry.get("guardrail_request", None), - output=guardrail_entry.get("guardrail_response", None), + input=cap_payload( + cast( # cast-ok: unknown-typed payload + "object", + guardrail_entry.get("guardrail_request", None), + ), + max_field_bytes, + ), + output=cap_payload( + cast( # cast-ok: unknown-typed payload + "object", + guardrail_entry.get("guardrail_response", None), + ), + max_field_bytes, + ), metadata=MappingProxyType( { "guardrail_name": guardrail_entry.get("guardrail_name", None), diff --git a/tests/unit/integrations/langfuse/test_langfuse_field_cap.py b/tests/unit/integrations/langfuse/test_langfuse_field_cap.py new file mode 100644 index 00000000000..237276d070d --- /dev/null +++ b/tests/unit/integrations/langfuse/test_langfuse_field_cap.py @@ -0,0 +1,228 @@ +"""Per-field byte caps: no string bound for a Langfuse span attribute may exceed the cap. + +The wiring tests are the regression guard for the ingestion guarantee: a multi-megabyte +message must reach the exported span attributes cut at ``LANGFUSE_MAX_FIELD_BYTES`` (default +32 KiB) with a ``...(truncated)`` marker, never whole. +""" + +import json +from datetime import datetime, timezone +from typing import Final + +import pytest +from langfuse import LangfuseOtelSpanAttributes as A + +import litellm +import litellm.integrations.langfuse.langfuse as langfuse_module +from litellm.integrations.langfuse import langfuse_sdk +from litellm.integrations.langfuse.field_cap import ( + _TRUNCATION_SUFFIX, + DEFAULT_MAX_FIELD_BYTES, + MAX_FIELD_BYTES_METADATA_KEY, + cap_payload, + resolve_max_field_bytes, +) +from litellm.integrations.langfuse.langfuse_prompt_management import ( + LangfusePromptManagement, + langfuse_client_init, +) + + +def test_default_cap_is_32_kib(): + assert DEFAULT_MAX_FIELD_BYTES == 32 * 1024 + + +def test_resolve_prefers_request_metadata_over_env(monkeypatch): + monkeypatch.setenv("LANGFUSE_MAX_FIELD_BYTES", "4096") + assert resolve_max_field_bytes(1024) == 1024 + + +def test_resolve_falls_back_to_env(monkeypatch): + monkeypatch.setenv("LANGFUSE_MAX_FIELD_BYTES", "4096") + assert resolve_max_field_bytes() == 4096 + + +def test_resolve_defaults_when_neither_set(monkeypatch): + monkeypatch.delenv("LANGFUSE_MAX_FIELD_BYTES", raising=False) + assert resolve_max_field_bytes() == DEFAULT_MAX_FIELD_BYTES + + +def test_resolve_env_string_zero_disables_capping(monkeypatch): + monkeypatch.setenv("LANGFUSE_MAX_FIELD_BYTES", "0") + assert resolve_max_field_bytes() == 0 + + +@pytest.mark.parametrize("raw", ["abc", "", True, [], 3.5, object()]) +def test_resolve_invalid_values_fall_back_to_default(monkeypatch, raw): + monkeypatch.delenv("LANGFUSE_MAX_FIELD_BYTES", raising=False) + assert resolve_max_field_bytes(raw) == DEFAULT_MAX_FIELD_BYTES + + +def test_short_string_passes_through_unchanged(): + assert cap_payload("hello", DEFAULT_MAX_FIELD_BYTES) == "hello" + + +def test_long_ascii_string_is_cut_at_the_cap_with_marker(): + assert cap_payload("a" * 100, 32) == "a" * 16 + _TRUNCATION_SUFFIX + + +def test_multibyte_string_never_splits_a_rune(): + capped: Final = cap_payload("漢" * 100, 32) + assert capped == "漢" * 2 + _TRUNCATION_SUFFIX + + +def test_capped_string_serialized_width_stays_within_cap(): + capped: Final = cap_payload("x" * 10_000, 256) + assert len(json.dumps(capped)) <= 256 + + +def test_cjk_string_wire_width_stays_within_cap_despite_uXXXX_escaping(): + capped: Final = cap_payload("漢" * 100_000, DEFAULT_MAX_FIELD_BYTES) + assert capped == "漢" * 5_458 + _TRUNCATION_SUFFIX + assert len(json.dumps(capped)) <= DEFAULT_MAX_FIELD_BYTES + + +def test_astral_code_points_count_double_width(): + capped: Final = cap_payload("🙂" * 1_000, 32) + assert capped == "🙂" + _TRUNCATION_SUFFIX + assert len(json.dumps(capped)) <= 32 + + +def test_nested_containers_are_capped_recursively(): + payload: Final = { + "messages": [{"role": "user", "content": "x" * 100}], + "tools": ("y" * 100,), + "count": 5, + "nothing": None, + } + capped: Final = cap_payload(payload, 64) + assert capped["messages"][0]["content"] == "x" * 48 + _TRUNCATION_SUFFIX + assert capped["tools"][0] == "y" * 48 + _TRUNCATION_SUFFIX + assert capped["count"] == 5 + assert capped["nothing"] is None + + +def test_cap_below_marker_size_leaves_only_the_marker(): + assert cap_payload("x" * 100, 16) == _TRUNCATION_SUFFIX + + +def test_cap_below_marker_width_returns_empty_string(): + assert cap_payload("x" * 100, 15) == "" + assert cap_payload("x" * 100, 1) == "" + + +def test_circular_reference_resolves_to_the_safe_dumps_marker(): + payload: Final[dict[str, object]] = {"messages": []} + payload["messages"].append({"role": "user", "content": payload}) + capped: Final = cap_payload(payload, DEFAULT_MAX_FIELD_BYTES) + assert capped["messages"][0]["content"] == "CircularReference Detected" + + +def test_over_deep_nesting_resolves_to_the_safe_dumps_marker(): + from litellm.litellm_core_utils.safe_json_dumps import safe_dumps + + payload: Final[dict[str, object]] = {"leaf": "x" * 10} + for _ in range(200): + payload = {"nested": payload} + capped: Final = cap_payload(payload, DEFAULT_MAX_FIELD_BYTES) + assert "MaxDepthExceeded" in repr(capped) + assert safe_dumps(capped) == safe_dumps(payload) + + +def test_non_container_objects_pass_through(): + sentinel: Final = object() + assert cap_payload(sentinel, 16) is sentinel + + +def test_zero_limit_returns_value_unchanged(): + payload: Final = {"a": "z" * 100} + assert cap_payload(payload, 0) is payload + + +class _CapturedGeneration: + id: Final = "f" * 16 + + def __init__(self, attributes): + self.attributes = attributes + + def end(self, end_time=None): + return None + + +def _mock_mode(monkeypatch): + monkeypatch.setenv("LANGFUSE_MOCK", "true") + monkeypatch.setenv("LANGFUSE_HOST", "http://127.0.0.1:1") + monkeypatch.setenv("LANGFUSE_PUBLIC_KEY", "pk-field-cap") + monkeypatch.setenv("LANGFUSE_SECRET_KEY", "sk-field-cap") + monkeypatch.delenv("LANGFUSE_MAX_FIELD_BYTES", raising=False) + langfuse_client_init.cache_clear() + + +def _log_with_capture(monkeypatch, metadata, content_size): + captured = {} + original_start_generation: Final = langfuse_sdk.start_generation + + def fake_start_generation(**kwargs): + captured["attributes"] = dict(kwargs["attributes"]) + return original_start_generation(**kwargs) + + monkeypatch.setattr(langfuse_sdk, "start_generation", fake_start_generation) + monkeypatch.setattr(langfuse_module, "log_provider_specific_information_as_span", lambda **kwargs: None) + huge: Final = "x" * content_size + now: Final = datetime.now(timezone.utc) + logger = LangfusePromptManagement() + try: + logger.log_event_on_langfuse( + kwargs={ + "litellm_call_id": "call-field-cap", + "call_type": "completion", + "model": "gpt-5.4", + "litellm_params": {"metadata": metadata}, + "messages": [{"role": "user", "content": huge}], + "optional_params": {"temperature": 0.5}, + }, + response_obj=litellm.ModelResponse( + choices=[{"message": {"role": "assistant", "content": "y" * content_size}}] + ), + start_time=now, + end_time=now, + ) + finally: + langfuse_sdk.release_langfuse_tracing(logger.tracing, grace_seconds=0.0) + langfuse_client_init.cache_clear() + return captured["attributes"] + + +def test_exported_generation_fields_are_capped_before_ingestion(monkeypatch): + _mock_mode(monkeypatch) + monkeypatch.setattr(litellm, "langfuse_default_tags", ["huge_tag"]) + attributes: Final = _log_with_capture( + monkeypatch, + metadata={"trace_id": "b" * 32, "trace_version": "v" * 100_000, "huge_tag": "t" * 100_000}, + content_size=DEFAULT_MAX_FIELD_BYTES * 4, + ) + + envelope_slack: Final = 256 # the JSON envelope around the capped string + for key in (A.OBSERVATION_INPUT, A.OBSERVATION_OUTPUT): + value = attributes[key] + assert _TRUNCATION_SUFFIX in value, f"{key} was not truncated" + assert len(value.encode("utf-8")) <= (DEFAULT_MAX_FIELD_BYTES + len(_TRUNCATION_SUFFIX) + envelope_slack), ( + f"{key} exceeded the cap" + ) + + for key in (A.VERSION, A.TRACE_TAGS): + assert len(str(attributes[key])) <= DEFAULT_MAX_FIELD_BYTES + envelope_slack, f"{key} exceeded the cap" + + assert json.loads(attributes[A.OBSERVATION_MODEL_PARAMETERS]) == {"temperature": 0.5} + + +def test_metadata_override_can_disable_capping(monkeypatch): + _mock_mode(monkeypatch) + attributes: Final = _log_with_capture( + monkeypatch, + metadata={"trace_id": "c" * 32, MAX_FIELD_BYTES_METADATA_KEY: 0}, + content_size=DEFAULT_MAX_FIELD_BYTES * 4, + ) + + assert len(attributes[A.OBSERVATION_INPUT]) > DEFAULT_MAX_FIELD_BYTES + assert attributes[A.OBSERVATION_INPUT].find(_TRUNCATION_SUFFIX) == -1