Merge 25e9f8edc2 into 3f4dc93b40

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

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),

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