Merge pull request #42027 from BerriAI/litellm_stream_boundary_attr_forwarding

Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/litellm_core_utils/litellm_logging.py

@@ -5528,6 +5528,10 @@ class StandardLoggingPayloadSetup:
             for key in metadata.keys() & _STANDARD_LOGGING_METADATA_KEYS:
                 clean_metadata[key] = metadata[key]
 
+            recorded_guardrails: Final = metadata.get("applied_guardrails")
+            if applied_guardrails and isinstance(recorded_guardrails, list):
+                clean_metadata["applied_guardrails"] = list(dict.fromkeys([*applied_guardrails, *recorded_guardrails]))
+
             user_api_key: Final = metadata.get("user_api_key")
             if user_api_key and isinstance(user_api_key, str) and is_valid_sha256_hash(user_api_key):
                 clean_metadata["user_api_key_hash"] = user_api_key

litellm/proxy/utils.py

@@ -470,12 +470,16 @@ def _record_raising_guardrail(request_data: Mapping[str, object], callback: obje
 
 
 class _UpstreamStreamBoundary(Generic[_T]):
-    __slots__ = ("_upstream", "failure")
+    __slots__ = ("_source", "_upstream", "failure")
 
     def __init__(self, upstream: AsyncIterable[_T]) -> None:
+        self._source: Final = upstream
         self._upstream: Final = upstream.__aiter__()
         self.failure: BaseException | None = None
 
+    def __getattr__(self, name: str) -> object:
+        return getattr(self._source, name)
+
     def __aiter__(self) -> "_UpstreamStreamBoundary[_T]":
         return self
 

tests/test_litellm/litellm_core_utils/test_litellm_logging.py

@@ -3532,6 +3532,24 @@ def test_function_setup_litellm_metadata_guardrail_writes_visible_after_setup():
     assert merged.get("applied_guardrails") == ["pam-ethical-request"]
 
 
+def test_get_standard_logging_metadata_merges_recorded_applied_guardrails():
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    result = StandardLoggingPayloadSetup.get_standard_logging_metadata(
+        metadata={"applied_guardrails": ["blocker"]},
+        litellm_params={},
+        applied_guardrails=["guard-a", "blocker", "guard-b"],
+    )
+    assert result["applied_guardrails"] == ["guard-a", "blocker", "guard-b"]
+
+    result = StandardLoggingPayloadSetup.get_standard_logging_metadata(
+        metadata={"applied_guardrails": ["blocker"]},
+        litellm_params={},
+        applied_guardrails=["guard-a"],
+    )
+    assert result["applied_guardrails"] == ["guard-a", "blocker"]
+
+
 def test_function_setup_metadata_takes_precedence_over_litellm_metadata():
     """
     Test that when BOTH metadata and litellm_metadata are present (e.g., user sets

tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py

@@ -12,7 +12,7 @@ from __future__ import annotations
 import asyncio
 from collections.abc import AsyncGenerator, AsyncIterator
 from datetime import datetime
-from typing import Any, Dict, List
+from typing import Any, Dict, Final, List
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest
@@ -179,6 +179,29 @@ async def _one_chunk() -> AsyncGenerator[object, None]:
     yield "chunk"
 
 
+class _AttributeStream:
+    _hidden_params = {"model_id": "m-1"}
+    model = "gpt-x"
+
+    def __init__(self) -> None:
+        self._chunks = ("chunk-1", "chunk-2")
+        self._index = 0
+        self.closed = False
+
+    def __aiter__(self) -> "_AttributeStream":
+        return self
+
+    async def __anext__(self) -> str:
+        if self._index >= len(self._chunks):
+            raise StopAsyncIteration
+        chunk = self._chunks[self._index]
+        self._index += 1
+        return chunk
+
+    async def aclose(self) -> None:
+        self.closed = True
+
+
 @pytest.mark.asyncio
 async def test_wrap_streaming_iterator_with_enrichment_passes_through_chunks(proxy_logging):
     async def gen():
@@ -247,6 +270,68 @@ async def test_wrap_streaming_iterator_leaves_upstream_http_exception_unattribut
     assert request_data == {}
 
 
+@pytest.mark.asyncio
+async def test_wrap_streaming_iterator_forwards_response_attributes_to_hook(proxy_logging):
+    async def prefix_hook(*, response: AsyncIterator[object]) -> AsyncGenerator[object, None]:
+        async for chunk in response:
+            yield f"{response._hidden_params['model_id']}:{response.model}:{chunk}"
+
+    source = _AttributeStream()
+    wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(
+        callback=MagicMock(guardrail_name="g", event_hook="post_call"),
+        response=source,
+        hook=prefix_hook,
+        request_data={},
+    )
+
+    assert [chunk async for chunk in wrapped] == [
+        "m-1:gpt-x:chunk-1",
+        "m-1:gpt-x:chunk-2",
+    ]
+
+
+@pytest.mark.asyncio
+async def test_wrap_streaming_iterator_forwards_aclose_to_upstream(proxy_logging):
+    async def close_hook(*, response: AsyncIterator[object]) -> AsyncGenerator[object, None]:
+        first: Final = await response.__anext__()
+        yield first
+        await response.aclose()
+
+    source = _AttributeStream()
+    request_data: dict[str, object] = {}
+    wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(
+        callback=MagicMock(guardrail_name="g", event_hook="post_call"),
+        response=source,
+        hook=close_hook,
+        request_data=request_data,
+    )
+
+    assert [chunk async for chunk in wrapped] == ["chunk-1"]
+    assert source.closed is True
+    assert request_data == {}
+
+
+@pytest.mark.asyncio
+async def test_wrap_streaming_iterator_missing_attribute_still_raises(proxy_logging):
+    async def missing_attribute_hook(*, response: AsyncIterator[object]) -> AsyncGenerator[object, None]:
+        _missing: Final = response.not_there
+        if False:
+            yield
+
+    request_data: dict[str, object] = {}
+    wrapped = proxy_logging._wrap_streaming_iterator_with_enrichment(
+        callback=MagicMock(guardrail_name="hook-bug", event_hook="post_call"),
+        response=_one_chunk(),
+        hook=missing_attribute_hook,
+        request_data=request_data,
+    )
+
+    with pytest.raises(AttributeError):
+        async for _ in wrapped:
+            pass
+    assert request_data["metadata"]["applied_guardrails"] == ["hook-bug"]
+
+
 # ---------------------------------------------------------------------------
 # async_post_call_streaming_hook
 # ---------------------------------------------------------------------------