diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 36e5c95c79c..bfc7175cc65 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -723,29 +723,35 @@ def _ran_or_runs_outside_pipelines( def _post_call_steps_named_before_headers( - data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth + data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: str ) -> frozenset[str]: """ Post_call pipeline steps the executor names, as it runs them, before the response headers go out: every step of a whole response, and of a streamed one the steps of - the pipelines that hold the stream until they finish. + the pipelines that hold the stream until they finish. The LLM pass-through route builds + its headers before any response step runs, so none of its steps. """ + if call_type == CallTypes.allm_passthrough_route.value: + return frozenset() if data.get("stream") is not True and data.get("method") not in A2A_STREAMING_METHODS: return pipeline_managed_guardrail_names(data, "post_call") gated: Final = _stream_gated_pipelines(data, user_api_key_dict) return frozenset() if _pipelines_stream_live(gated) else _pipeline_step_guardrail_names(gated) -def _pipeline_steps_that_have_not_run(data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth) -> frozenset[str]: +def _pipeline_steps_that_have_not_run( + data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: str +) -> frozenset[str]: """ Guardrails only policy pipelines run whose pre_call hook has not run for this request: pre_call steps a pipeline skipped and the post_call steps the executor names when it runs them. A stream runs a pipeline it cannot hold as guardrails on their own, and one - whose steps all stream live after the headers go out, so those steps are named up front - like any other guardrail. + whose steps all stream live after the headers go out, and the LLM pass-through route runs + its response steps after its headers are built, so those steps are named up front like + any other guardrail. """ pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call") - post_call: Final = _post_call_steps_named_before_headers(data, user_api_key_dict) + post_call: Final = _post_call_steps_named_before_headers(data, user_api_key_dict, call_type) return frozenset( name for name in pre_call | post_call @@ -2694,7 +2700,7 @@ class ProxyLogging: data.pop("prompt_version", None) data.pop("prompt_environment", None) - def _process_guardrail_metadata(self, data: dict, user_api_key_dict: UserAPIKeyAuth) -> None: + def _process_guardrail_metadata(self, data: dict, user_api_key_dict: UserAPIKeyAuth, call_type: str) -> None: """Process guardrails from metadata and add to applied_guardrails.""" from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, @@ -2722,6 +2728,7 @@ class ProxyLogging: not_run: Final = _pipeline_steps_that_have_not_run( cast(dict[str, object], data), # cast-ok: request payload user_api_key_dict, + call_type, ) for guardrail_name in guardrails_in_metadata: if ( @@ -3015,7 +3022,7 @@ class ProxyLogging: and (guardrails_only or not caps.has_pre_call_override) ): if data is not None: - self._process_guardrail_metadata(data, user_api_key_dict) + self._process_guardrail_metadata(data, user_api_key_dict, call_type) return data parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = ( @@ -3119,17 +3126,17 @@ class ProxyLogging: ) if data is not None: - self._process_guardrail_metadata(data, user_api_key_dict) + self._process_guardrail_metadata(data, user_api_key_dict, call_type) return data except SensitiveDataRouteException as e: data = await self._handle_sensitive_data_route_exception(e, data, user_api_key_dict) if data is not None: - self._process_guardrail_metadata(data, user_api_key_dict) + self._process_guardrail_metadata(data, user_api_key_dict, call_type) return data except Exception: if data is not None: - self._process_guardrail_metadata(data, user_api_key_dict) + self._process_guardrail_metadata(data, user_api_key_dict, call_type) raise async def _run_parallel_pre_call_guardrails( diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py index 58e46141cac..174a88ac306 100644 --- a/tests/integration/observability/test_guardrail_pipeline_applied_header.py +++ b/tests/integration/observability/test_guardrail_pipeline_applied_header.py @@ -20,6 +20,7 @@ import yaml from tests.integration._support.client import Gateway, eventually, gateway_from_environment, object_value from tests.integration._support.otlp_sink import Span, SpanSinks, recorded_spans from tests.integration._support.process import owned_proxy +from tests.integration._support.upstream import aws_event_stream_frame from tests.integration._support.wire import Reply, Request, Wire, wire_server _HEADER: Final = "x-litellm-applied-guardrails" @@ -50,6 +51,8 @@ _CANNOT_HOLD_STREAM: Final = "cannot-hold-stream" _ERRORS_PAST_FIRST: Final = "errors-past-first" _JUDGES_THEN_CHECKS: Final = "judges-then-checks" _ENDPOINTS: Final = ("/v1/chat/completions", "/v1/responses", "/v1/messages") +_BEDROCK_MODEL: Final = "pipeline-bedrock" +_BEDROCK_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" @dataclass(frozen=True, slots=True) @@ -126,7 +129,58 @@ def _responses_reply(request: Request) -> Reply: ) +def _bedrock_event(event_type: str, payload: dict[str, object]) -> bytes: + return aws_event_stream_frame( + {":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, + json.dumps(payload).encode(), + ) + + +def _bedrock_reply(request: Request) -> Reply: + action: Final = request.target.rsplit("/", 1)[1] + if action == "converse-stream": + return Reply( + content_type="application/vnd.amazon.eventstream", + chunks=( + _bedrock_event("messageStart", {"role": "assistant"}), + _bedrock_event("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": "guarded"}}), + _bedrock_event("contentBlockStop", {"contentBlockIndex": 0}), + _bedrock_event("messageStop", {"stopReason": "end_turn"}), + _bedrock_event( + "metadata", {"usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4}, "metrics": {}} + ), + ), + ) + if action == "invoke": + return Reply( + body=json.dumps( + { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "model": _BEDROCK_MODEL_ID, + "content": [{"type": "text", "text": "guarded"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 3, "output_tokens": 1}, + } + ).encode() + ) + assert action == "converse", request.target + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "guarded"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4}, + "metrics": {"latencyMs": 1}, + } + ).encode() + ) + + def _completion(request: Request) -> Reply: + if request.target.startswith("/model/"): + return _bedrock_reply(request) if request.target == "/v1/responses": return _responses_reply(request) assert request.target == "/v1/chat/completions", request.target @@ -208,11 +262,21 @@ def _policy(mode: str, steps: tuple[tuple[str, str, str], ...], *, on_error: str } -def _config(directory: Path, vendor: Wire) -> Path: +def _config(directory: Path, vendor: Wire, upstream: Wire) -> Path: configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) configuration["litellm_settings"] = {**configuration.get("litellm_settings", {}), "callbacks": ["otel"]} configuration["model_list"] = [ *configuration.get("model_list", []), + { + "model_name": _BEDROCK_MODEL, + "litellm_params": { + "model": f"bedrock/{_BEDROCK_MODEL_ID}", + "api_base": upstream.url, + "aws_access_key_id": "AKIASYNTHETICPROVIDER", + "aws_secret_access_key": "synthetic-provider-secret", + "aws_region_name": "us-east-1", + }, + }, { "model_name": _JUDGE_MODEL, "litellm_params": { @@ -316,7 +380,7 @@ def rig(tmp_path_factory: pytest.TempPathFactory, audit_sinks: SpanSinks) -> Ite shared, directory, {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": audit_sinks.operator, "OTEL_BSP_SCHEDULE_DELAY": "100"}, - config=_config(directory, vendor), + config=_config(directory, vendor, upstream), workers=2, ) as owned, owned.scenario() as scenario, @@ -572,6 +636,45 @@ def test_a_response_pipeline_names_a_judge_step_that_records_its_own_results(rig assert sorted(object_value(rows[0]["metadata"]).get("applied_guardrails") or []) == [_JUDGE, _RESPONSE_SECOND] +def _bedrock_body(action: str) -> dict[str, object]: + prompt: Final = f"hello {uuid.uuid4().hex}" + if action == "invoke": + return { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 20, + "messages": [{"role": "user", "content": [{"type": "text", "text": prompt}]}], + } + return {"messages": [{"role": "user", "content": [{"text": prompt}]}]} + + +@pytest.mark.parametrize("action", ["converse", "invoke", "converse-stream"]) +def test_a_pass_through_route_names_its_response_pipeline_up_front_because_its_headers_are_built_first( + rig: _Rig, action: str +) -> None: + rig.vendor.drain() + rig.upstream.drain() + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"tags": [_RESPONSE_STOPS_AFTER_FIRST]}) + response: Final = rig.gateway.request( + "POST", f"/bedrock/model/{_BEDROCK_MODEL}/{action}", _bedrock_body(action), key=scenario.key(team_id=team) + ) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST,) + assert [request.target.rsplit("/", 1)[1] for request in rig.upstream.drain()] == [action] + assert sorted(response.headers.get(_HEADER, "").split(",")) == [_RESPONSE_FIRST, _RESPONSE_SECOND], dict( + response.headers + ) + rows: Final = eventually( + lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)), + lambda found: len(found) == 1, + seconds=70, + ) + assert sorted(object_value(rows[0]["metadata"]).get("applied_guardrails") or []) == [ + _RESPONSE_FIRST, + _RESPONSE_SECOND, + ], rows + + def _agent_card(name: str, url: str) -> dict[str, object]: return { "protocolVersion": "0.3", diff --git a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py index bd0480182b7..ddeaf45d32b 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -280,7 +280,7 @@ def test_process_guardrail_metadata_calls_header_helper(proxy_logging, monkeypat monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) data = {"metadata": {"guardrails": ["g1", "g2"]}} - proxy_logging._process_guardrail_metadata(data, UserAPIKeyAuth()) + proxy_logging._process_guardrail_metadata(data, UserAPIKeyAuth(), "completion") snapshot = { "call_count": len(calls), "first_name": calls[0]["name"], @@ -305,7 +305,7 @@ def test_process_guardrail_metadata_skips_already_applied(proxy_logging, monkeyp monkeypatch.setattr(callback_utils, "add_guardrail_to_applied_guardrails_header", fake_add) data = {"metadata": {"guardrails": ["g1", "g2"], "applied_guardrails": ["g1"]}} - proxy_logging._process_guardrail_metadata(data, UserAPIKeyAuth()) + proxy_logging._process_guardrail_metadata(data, UserAPIKeyAuth(), "completion") assert calls == ["g2"] @@ -317,12 +317,12 @@ def test_process_guardrail_metadata_no_metadata_is_noop(proxy_logging, monkeypat "add_guardrail_to_applied_guardrails_header", MagicMock(side_effect=AssertionError("should not be called")), ) - proxy_logging._process_guardrail_metadata({}, UserAPIKeyAuth()) + proxy_logging._process_guardrail_metadata({}, UserAPIKeyAuth(), "completion") def test_process_guardrail_metadata_invalid_data_raises(proxy_logging): with pytest.raises(AttributeError): - proxy_logging._process_guardrail_metadata(None, UserAPIKeyAuth()) # type: ignore[arg-type] + proxy_logging._process_guardrail_metadata(None, UserAPIKeyAuth(), "completion") # type: ignore[arg-type] # --------------------------------------------------------------------------- @@ -1185,7 +1185,7 @@ class _PipelineStepGuardrail(CustomGuardrail): def _two_step_pipeline_data( mode: str, first_on_pass: str, first_on_fail: str = "block", second_on_fail: str = "block" ) -> dict[str, object]: - pipeline = GuardrailPipeline( + pipeline: Final = GuardrailPipeline( mode=mode, steps=[ PipelineStep(guardrail="first", on_pass=first_on_pass, on_fail=first_on_fail), @@ -1214,11 +1214,11 @@ async def test_pre_call_hook_names_only_the_pre_call_pipeline_steps_that_ran( expected_calls: tuple[int, int], expected_header: list[str], ) -> None: - first = _PipelineStepGuardrail("first", "pre_call") - second = _PipelineStepGuardrail("second", "pre_call") + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _PipelineStepGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) - result = await proxy_logging.pre_call_hook( + result: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=_two_step_pipeline_data("pre_call", first_on_pass), call_type="completion", @@ -1232,11 +1232,11 @@ async def test_pre_call_hook_names_only_the_pre_call_pipeline_steps_that_ran( async def test_pre_call_hook_names_a_skipped_step_whose_guardrail_also_runs_outside_the_pipeline( proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, ) -> None: - first = _PipelineStepGuardrail("first", "pre_call") - second = _PipelineStepGuardrail("second", ["pre_call", "post_call"]) + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _PipelineStepGuardrail("second", ["pre_call", "post_call"]) monkeypatch.setattr(litellm, "callbacks", [first, second]) - result = await proxy_logging.pre_call_hook( + result: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=_two_step_pipeline_data("pre_call", "allow"), call_type="completion", @@ -1270,11 +1270,11 @@ async def test_pre_call_hook_names_a_step_that_raised_before_the_pipeline_moved_ proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, error: Exception, ) -> None: - first = _RaisingPipelineStepGuardrail("first", error) - second = _PipelineStepGuardrail("second", "pre_call") + first: Final = _RaisingPipelineStepGuardrail("first", error) + second: Final = _PipelineStepGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) - result = await proxy_logging.pre_call_hook( + result: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=_two_step_pipeline_data("pre_call", "next", first_on_fail="next"), call_type="completion", @@ -1290,10 +1290,10 @@ async def test_deployment_hook_still_runs_a_pipeline_step_that_raised( proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, error: Exception, ) -> None: - first = _RaisingPipelineStepGuardrail("first", error) - second = _PipelineStepGuardrail("second", "pre_call") + first: Final = _RaisingPipelineStepGuardrail("first", error) + second: Final = _PipelineStepGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) - result = await proxy_logging.pre_call_hook( + result: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=_two_step_pipeline_data("pre_call", "next", first_on_fail="next"), call_type="completion", @@ -1324,11 +1324,11 @@ class _OpensLitellmMetadataGuardrail(_PipelineStepGuardrail): async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_later_step_opened( proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, ) -> None: - first = _PipelineStepGuardrail("first", "pre_call") - second = _OpensLitellmMetadataGuardrail("second", "pre_call") + first: Final = _PipelineStepGuardrail("first", "pre_call") + second: Final = _OpensLitellmMetadataGuardrail("second", "pre_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) - result = await proxy_logging.pre_call_hook( + result: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=_two_step_pipeline_data("pre_call", "next"), call_type="completion", @@ -1344,12 +1344,12 @@ async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_lat async def test_post_call_pipeline_steps_are_named_only_once_they_run( proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, ) -> None: - first = _PipelineStepGuardrail("first", "post_call") - second = _PipelineStepGuardrail("second", "post_call") + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) - data = await proxy_logging.pre_call_hook( + data: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=_two_step_pipeline_data("post_call", "allow"), call_type="completion", @@ -1376,12 +1376,12 @@ class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail): async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_step_opens_litellm_metadata( proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, ) -> None: - first = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call") - second = _PipelineStepGuardrail("second", "post_call") + first: Final = _OpensLitellmMetadataOnResponseGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", "post_call") monkeypatch.setattr(litellm, "callbacks", [first, second]) monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) - data = await proxy_logging.pre_call_hook( + data: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), data=_two_step_pipeline_data("post_call", "next"), call_type="completion", @@ -1416,14 +1416,14 @@ async def test_pre_call_hook_names_up_front_the_post_call_steps_a_stream_runs_af streams_live: bool, expected_header: list[str] | None, ) -> None: - guardrails = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second")] + guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second")] for guardrail in guardrails: guardrail.streaming_buffer_until_moderated = not streams_live monkeypatch.setattr(litellm, "callbacks", guardrails) - data = _two_step_pipeline_data("post_call", "next", first_on_fail="next", second_on_fail="next") + data: Final = _two_step_pipeline_data("post_call", "next", first_on_fail="next", second_on_fail="next") data.update(asks_for) - result = await proxy_logging.pre_call_hook( + result: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(request_route=request_route), data=data, call_type="completion", @@ -1432,6 +1432,28 @@ async def test_pre_call_hook_names_up_front_the_post_call_steps_a_stream_runs_af assert result["metadata"].get("applied_guardrails") == expected_header +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True], ids=["whole-response", "stream"]) +async def test_pre_call_hook_names_up_front_the_post_call_steps_the_llm_pass_through_route_runs_after_its_headers( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + stream: bool, +) -> None: + guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second")] + monkeypatch.setattr(litellm, "callbacks", guardrails) + data: Final = _two_step_pipeline_data("post_call", "allow") + data["stream"] = stream + + result: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(request_route="/bedrock/model/m/converse"), + data=data, + call_type="allm_passthrough_route", + ) + + assert result["metadata"].get("applied_guardrails") == ["first", "second"] + + @pytest.mark.asyncio async def test_post_call_pipeline_replacement_response_reaches_caller( proxy_logging, make_user_api_key_auth, monkeypatch diff --git a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py index 48dac33d010..89a59d44380 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py +++ b/tests/unit/proxy/utils/proxy_logging/test_pre_call_hook.py @@ -170,7 +170,7 @@ async def test_pre_call_hook_processes_guardrail_metadata_when_no_overrides(prox proxy_logging.slack_alerting_instance = MagicMock(alerting=None) invoked = {} - def fake_process(d: dict[str, object], user_api_key_dict: UserAPIKeyAuth) -> None: + def fake_process(d: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: str) -> None: invoked["data"] = d proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment]