diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f7679b31f69..009089e0a8e 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/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 diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b078a65759e..62710d570db 100644 --- a/litellm/proxy/utils.py +++ b/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 diff --git a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py b/tests/test_litellm/litellm_core_utils/test_litellm_logging.py index 999adbdd935..626a13c8061 100644 --- a/tests/test_litellm/litellm_core_utils/test_litellm_logging.py +++ b/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 diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py index 6fb000b4fa7..5132aeb02e8 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/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 # ---------------------------------------------------------------------------