mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-04 07:27:51 +00:00
fix(guardrails): name the LLM pass-through route's response pipeline up front
The LLM pass-through route builds its response headers before the response pipeline runs, so the pre_call sweep claims those steps as on main. Adds Final to the new test locals.
This commit is contained in:
yucheng
parent
80b95eb8c0
commit
5434ce4675
4 changed files with 175 additions and 43 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue