diff --git a/litellm/proxy/a2a/version_convert.py b/litellm/proxy/a2a/version_convert.py index eb55bff50bc..bde2ff45a88 100644 --- a/litellm/proxy/a2a/version_convert.py +++ b/litellm/proxy/a2a/version_convert.py @@ -24,8 +24,7 @@ The two wire shapes: ``Task`` result is a bare object without ``kind``. """ -from collections.abc import Callable, Mapping -from types import MappingProxyType +from collections.abc import Callable from typing import Final, Literal from pydantic import BaseModel @@ -40,26 +39,6 @@ JsonDict = dict[str, object] _V1_SEND_ENVELOPE_KEYS: Final = frozenset({"message", "task"}) _V1_STREAM_ENVELOPE_KEYS: Final = frozenset({"message", "task", "statusUpdate", "artifactUpdate"}) -A2A_PASCAL_TO_WIRE: Final[Mapping[str, str]] = MappingProxyType( - { - "SendMessage": "message/send", - "SendStreamingMessage": "message/stream", - "GetTask": "tasks/get", - "ListTasks": "tasks/list", - "CancelTask": "tasks/cancel", - "SubscribeToTask": "tasks/resubscribe", - "CreateTaskPushNotificationConfig": "tasks/pushNotificationConfig/set", - "GetTaskPushNotificationConfig": "tasks/pushNotificationConfig/get", - "ListTaskPushNotificationConfigs": "tasks/pushNotificationConfig/list", - "DeleteTaskPushNotificationConfig": "tasks/pushNotificationConfig/delete", - "GetExtendedAgentCard": "agent/getAuthenticatedExtendedCard", - } -) -_STREAMING_WIRE_METHODS: Final = frozenset({"message/stream", "tasks/resubscribe"}) -A2A_STREAMING_METHODS: Final = _STREAMING_WIRE_METHODS | frozenset( - pascal for pascal, wire in A2A_PASCAL_TO_WIRE.items() if wire in _STREAMING_WIRE_METHODS -) - def _dump_03(model: BaseModel) -> JsonDict: """Dump a compat (0.3) pydantic model to its camelCase wire dict.""" diff --git a/litellm/proxy/agent_endpoints/a2a_endpoints.py b/litellm/proxy/agent_endpoints/a2a_endpoints.py index 612988a74e9..4691c7b0bc5 100644 --- a/litellm/proxy/agent_endpoints/a2a_endpoints.py +++ b/litellm/proxy/agent_endpoints/a2a_endpoints.py @@ -28,7 +28,6 @@ from litellm.litellm_core_utils.url_utils import SSRFError, validate_url from litellm.llms.a2a.common_utils import resolve_a2a_hop_auth_header from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.a2a.version_convert import ( - A2A_PASCAL_TO_WIRE, A2AVersion, normalize_agent_card, normalize_jsonrpc_response, @@ -66,6 +65,21 @@ _SSE_KEEPALIVE_HEADERS: Final[Mapping[str, str]] = MappingProxyType( } ) +_PASCAL_TO_WIRE: Final[Mapping[str, str]] = { + "SendMessage": "message/send", + "SendStreamingMessage": "message/stream", + "GetTask": "tasks/get", + "ListTasks": "tasks/list", + "CancelTask": "tasks/cancel", + "SubscribeToTask": "tasks/resubscribe", + "CreateTaskPushNotificationConfig": "tasks/pushNotificationConfig/set", + "GetTaskPushNotificationConfig": "tasks/pushNotificationConfig/get", + "ListTaskPushNotificationConfigs": "tasks/pushNotificationConfig/list", + "DeleteTaskPushNotificationConfig": "tasks/pushNotificationConfig/delete", + "GetExtendedAgentCard": "agent/getAuthenticatedExtendedCard", +} + + _DECODED_JSON: Final = TypeAdapter(object) @@ -118,7 +132,7 @@ def _served_version(agent: "AgentResponse", request: Request, original_method: s configured: Final = (agent.agent_card_params or {}).get("protocolVersion") if configured in ("0.3", "1.0"): return configured - if original_method in A2A_PASCAL_TO_WIRE: + if original_method in _PASCAL_TO_WIRE: return "1.0" return "1.0" if request.headers.get("a2a-version", "").startswith("1.") else "0.3" @@ -685,7 +699,7 @@ async def invoke_agent_a2a( params = body.get("params", {}) if method: - method = A2A_PASCAL_TO_WIRE.get(method, method) + method = _PASCAL_TO_WIRE.get(method, method) if isinstance(params, dict): # extract any litellm params from the params - eg. 'guardrails' @@ -788,7 +802,6 @@ async def invoke_agent_a2a( proxy_config=proxy_config, route_type="asend_message", version=version, - defer_async_logging=method == "message/send", ) request_data = data diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6579afb90ef..cec176ac559 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2279,7 +2279,6 @@ class ProxyBaseLLMRequestProcessing: rate_limited_model: str | None = None, skip_guardrails: bool = False, logging_auth: UserAPIKeyAuth | None = None, - defer_async_logging: bool = False, ) -> tuple[dict, LiteLLMLoggingObj]: start_time: Final = datetime.now() # start before calling guardrail hooks @@ -2445,10 +2444,6 @@ class ProxyBaseLLMRequestProcessing: ) self.data["litellm_logging_obj"] = logging_obj - # Set before pre_call_hook: its guardrail sweep leaves a whole response's pipeline - # steps for the executor to name only when the success log waits for them. - if defer_async_logging: - logging_obj.defer_async_logging = True # Merge model-level guardrails before pre_call_hook so DB/UI-configured # guardrails actually execute on pre_call. Without this, guardrails set @@ -2539,7 +2534,6 @@ class ProxyBaseLLMRequestProcessing: model: str | None, route_type: str, llm_router: Router | None, - defer_async_logging: bool = False, ) -> tuple[dict, LiteLLMLoggingObj]: from litellm.proxy.common_utils.proxy_rate_limit_error import ( PER_MODEL_RATE_LIMIT_DESCRIPTOR_KEYS, @@ -2580,7 +2574,6 @@ class ProxyBaseLLMRequestProcessing: model=model, route_type=route_type, llm_router=llm_router, - defer_async_logging=defer_async_logging, ) except ProxyRateLimitError as original_exc: rate_limited_data: Final = self.data @@ -2637,7 +2630,6 @@ class ProxyBaseLLMRequestProcessing: route_type=route_type, llm_router=llm_router, rate_limited_model=original_model, - defer_async_logging=defer_async_logging, ) except ProxyRateLimitError as fallback_exc: if is_hard_per_model_limit(fallback_exc): @@ -2868,11 +2860,6 @@ class ProxyBaseLLMRequestProcessing: requested_model_from_client: Final[str | None] = client_model if isinstance(client_model, str) else None self._debug_log_request_payload() - # Defer async logging when post-call guardrails are configured so the - # StandardLoggingPayload is built after guardrails write to metadata. - # Cache the result to avoid scanning litellm.callbacks twice. - _post_call_guardrails_active: Final = self._has_post_call_guardrails() - if skip_pre_call_logic: logging_obj = self.data.get("litellm_logging_obj") if logging_obj is None: @@ -2897,11 +2884,14 @@ class ProxyBaseLLMRequestProcessing: model=model, route_type=route_type, llm_router=llm_router, - defer_async_logging=_post_call_guardrails_active - and not self._is_streaming_request(data=self.data, is_streaming_request=is_streaming_request), # pyright: ignore[reportUnknownMemberType] # self.data is the untyped request dict ) phase_event("litellm.request.pre_call_completed") + # Defer async logging when post-call guardrails are configured so the + # StandardLoggingPayload is built after guardrails write to metadata. + # Cache the result to avoid scanning litellm.callbacks twice. + _post_call_guardrails_active: Final = self._has_post_call_guardrails() + # Non-streaming: defer the create_task in wrapper_async so the # SLP is built after guardrails write to metadata. Streaming # uses a separate closure mechanism (see below). diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 8e4ba76381d..732c72c54ae 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -179,7 +179,6 @@ from litellm.proxy._types import ( SpecialModelNames, UserAPIKeyAuth, ) -from litellm.proxy.a2a.version_convert import A2A_STREAMING_METHODS from litellm.proxy.agent_endpoints.auth.agent_access_groups import CeilingResolver, resolve_agent_access_group_ceiling from litellm.proxy.auth.route_checks import RouteChecks from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_guardrails_header @@ -722,56 +721,77 @@ def _ran_or_runs_outside_pipelines( ) -def _success_log_waits_for_response_hooks(data: Mapping[str, object]) -> bool: +def _runs_only_in_pipelines( + name: str, data: Mapping[str, object], pre_call: frozenset[str], post_call: frozenset[str] +) -> bool: + return not _ran_or_runs_outside_pipelines( + PipelineExecutor.find_guardrail_callback(name), + data, + frozenset(mode for mode, names in (("pre_call", pre_call), ("post_call", post_call)) if name in names), + ) + + +def _request_steps_that_have_not_run(data: Mapping[str, object]) -> frozenset[str]: + """Pre_call policy pipeline steps a pipeline skipped, whose guardrail nothing else ran or runs for this request.""" + pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call") + return frozenset(name for name in pre_call if _runs_only_in_pipelines(name, data, pre_call, frozenset())) + + +def _steps_that_ran(result: PipelineExecutionResult) -> frozenset[str]: + return frozenset(step.guardrail_name for step in result.step_results if step.outcome != "skip") + + +def _whole_response_log_waits_for_its_pipelines(data: Mapping[str, object], call_type: str) -> bool: + """ + The success log is parked until the response hooks finish, and the headers are built + after them, except on the LLM pass-through route. + """ logging_obj: Final = data.get("litellm_logging_obj") - return isinstance(logging_obj, Logging) and logging_obj.defer_async_logging + return ( + call_type != CallTypes.allm_passthrough_route.value + and isinstance(logging_obj, Logging) + and logging_obj.enqueue_deferred_logging is not None + ) -def _post_call_steps_named_before_headers( - data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: str -) -> frozenset[str]: +def _stream_log_waits_for_its_pipelines(data: Mapping[str, object]) -> bool: + return getattr(data.get("litellm_logging_obj"), "_on_deferred_stream_complete", None) is not None + + +def _drop_response_steps_that_did_not_run( + data: dict[str, object], # mutable-ok: the request's applied guardrails are trimmed in place + steps: Sequence["PipelineStep"], + steps_that_ran: frozenset[str], +) -> None: """ - Post_call pipeline steps the executor names, as it runs them, before the response - headers go out and the success log is built: every step of a whole response whose - success log waits for the response hooks, and of a streamed one the steps of 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") - if _success_log_waits_for_response_hooks(data) - else frozenset() - ) - 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, 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, the LLM pass-through route runs - its response steps after its headers are built, and a route whose success log does not - wait for the response hooks builds that log first, so those steps are named up front - like any other guardrail. + Drop from the applied guardrails the steps of a response pipeline that no pipeline ran, + unless their guardrail also runs outside pipelines. Every step is named up front, so a + log built before the pipeline finished, such as a cache hit's, still lists the steps + that ran. The list is trimmed in place, since a route's logging object may hold a + shallow copy of the metadata. """ pre_call: Final = pipeline_managed_guardrail_names(data, "pre_call") - 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 - if not _ran_or_runs_outside_pipelines( - PipelineExecutor.find_guardrail_callback(name), - data, - frozenset(mode for mode, names in (("pre_call", pre_call), ("post_call", post_call)) if name in names), - ) + post_call: Final = frozenset(step.guardrail for step in steps) + not_run: Final = frozenset( + name for name in post_call - steps_that_ran if _runs_only_in_pipelines(name, data, pre_call, post_call) ) + if not not_run: + return + for bucket_key in ("metadata", "litellm_metadata"): + _drop_applied_guardrails_in_place(data.get(bucket_key), not_run) + + +def _drop_applied_guardrails_in_place(bucket: object, names: frozenset[str]) -> None: + if not isinstance(bucket, dict): + return + metadata: Final = cast("dict[str, object]", bucket) # cast-ok: metadata buckets are str-keyed + applied: Final = metadata.get("applied_guardrails") + if not isinstance(applied, list): + return + entries: Final = cast("list[object]", applied) # cast-ok: entries are only compared with names + entries[:] = [name for name in entries if name not in names] + if not entries: + metadata.pop("applied_guardrails") def _partition_post_call_callbacks() -> tuple[tuple[CustomGuardrail, ...], tuple[CustomLogger, ...]]: @@ -995,23 +1015,19 @@ def _streaming_pipeline_translation(user_api_key_dict: UserAPIKeyAuth) -> "BaseT return None if resolved is None else resolved[1] -def _stream_gated_pipelines( - request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth -) -> tuple[tuple[str, "GuardrailPipeline"], ...]: - translation: Final = _streaming_pipeline_translation(user_api_key_dict) - if translation is None: - return () - return tuple( - (policy_name, pipeline) - for policy_name, pipeline in _post_call_pipelines(request_data) - if not _pipeline_unsupported_streaming_guardrails(pipeline, translation) - ) - - def stream_gated_guardrail_names( request_data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth ) -> frozenset[str]: - return _pipeline_step_guardrail_names(_stream_gated_pipelines(request_data, user_api_key_dict)) + translation: Final = _streaming_pipeline_translation(user_api_key_dict) + if translation is None: + return frozenset() + return _pipeline_step_guardrail_names( + tuple( + (policy_name, pipeline) + for policy_name, pipeline in _post_call_pipelines(request_data) + if not _pipeline_unsupported_streaming_guardrails(pipeline, translation) + ) + ) def _streamable_post_call_pipelines( @@ -2711,7 +2727,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, call_type: str) -> None: + def _process_guardrail_metadata(self, data: dict) -> None: """Process guardrails from metadata and add to applied_guardrails.""" from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, @@ -2736,10 +2752,8 @@ class ProxyLogging: if not isinstance(applied_guardrails, list): applied_guardrails = [] - not_run: Final = _pipeline_steps_that_have_not_run( - cast(dict[str, object], data), # cast-ok: request payload - user_api_key_dict, - call_type, + not_run: Final = _request_steps_that_have_not_run( + cast(dict[str, object], data) # cast-ok: request payload ) for guardrail_name in guardrails_in_metadata: if ( @@ -2778,6 +2792,9 @@ class ProxyLogging: return data, None current_response = response # rebind-ok: chains each pipeline's replacement response into the next + steps_that_ran: frozenset[str] = frozenset() # rebind-ok: a later pipeline keeps a step an earlier one ran + request: Final = cast(dict[str, object], data) # cast-ok: request payload, updated in place by every pipeline + log_waits: Final = event_hook == "post_call" and _whole_response_log_waits_for_its_pipelines(request, call_type) for policy_name, pipeline in pipelines: if pipeline.mode != event_hook: continue @@ -2794,6 +2811,10 @@ class ProxyLogging: raw_request_snapshot=raw_request_snapshot, ) + steps_that_ran = steps_that_ran | _steps_that_ran(result) + if log_waits: + _drop_response_steps_that_did_not_run(request, pipeline.steps, steps_that_ran) + data = self._handle_pipeline_result( result=result, data=data, @@ -3033,7 +3054,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, call_type) + self._process_guardrail_metadata(data) return data parallel_guardrails: Final[tuple[CustomGuardrail, ...]] = ( @@ -3137,17 +3158,17 @@ class ProxyLogging: ) if data is not None: - self._process_guardrail_metadata(data, user_api_key_dict, call_type) + self._process_guardrail_metadata(data) 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, call_type) + self._process_guardrail_metadata(data) return data except Exception: if data is not None: - self._process_guardrail_metadata(data, user_api_key_dict, call_type) + self._process_guardrail_metadata(data) raise async def _run_parallel_pre_call_guardrails( @@ -4711,6 +4732,7 @@ class ProxyLogging: call_type, endpoint_translation = translation + steps_that_ran: frozenset[str] = frozenset() # rebind-ok: a later pipeline keeps a step an earlier one ran for policy_name, pipeline in pipelines: result: PipelineExecutionResult = await PipelineExecutor.execute_steps( steps=pipeline.steps, @@ -4722,6 +4744,9 @@ class ProxyLogging: streaming_chunks=buffered, endpoint_translation=endpoint_translation, ) + steps_that_ran = steps_that_ran | _steps_that_ran(result) + if _stream_log_waits_for_its_pipelines(request_data): + _drop_response_steps_that_did_not_run(request_data, pipeline.steps, steps_that_ran) try: ProxyLogging._handle_pipeline_result( result, data=request_data, policy_name=policy_name, original_response=buffered diff --git a/tests/integration/observability/test_guardrail_pipeline_applied_header.py b/tests/integration/observability/test_guardrail_pipeline_applied_header.py index 55f6f0764c3..b62e38a1a37 100644 --- a/tests/integration/observability/test_guardrail_pipeline_applied_header.py +++ b/tests/integration/observability/test_guardrail_pipeline_applied_header.py @@ -7,6 +7,7 @@ import json import uuid from collections.abc import Callable, Iterator, Mapping from dataclasses import dataclass +from functools import partial from pathlib import Path from types import MappingProxyType from typing import Final @@ -17,7 +18,7 @@ import openai import pytest import yaml -from tests.integration._support.client import Gateway, eventually, gateway_from_environment, object_value +from tests.integration._support.client import Gateway, JsonValue, 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 @@ -43,6 +44,7 @@ _STOPS_AFTER_FIRST: Final = "stops-after-first" _RUNS_EVERY_STEP: Final = "runs-every-step" _BLOCKS_AT_FIRST: Final = "blocks-at-first" _RESPONSE_STOPS_AFTER_FIRST: Final = "response-stops-after-first" +_RESPONSE_RUNS_EVERY_STEP: Final = "response-runs-every-step" _SHIELDS_THEN_CHECKS: Final = "shields-then-checks" _MOVES_PAST_FIRST: Final = "moves-past-first" _RESPONSE_MOVES_PAST_FIRST: Final = "response-moves-past-first" @@ -62,6 +64,58 @@ class _Rig: upstream: Wire model: str traces: str + sync_success_log: Path + + +_GEMINI_RESPONSE: Final = { + "candidates": [{"content": {"role": "model", "parts": [{"text": "guarded"}]}, "finishReason": "STOP", "index": 0}], + "usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 1, "totalTokenCount": 4}, + "modelVersion": "gemini-2.0-flash", +} +_FINE_TUNING_JOB: Final[dict[str, JsonValue]] = { + "id": "ftjob-pipeline", + "object": "fine_tuning.job", + "created_at": 1700000000, + "error": None, + "fine_tuned_model": None, + "finished_at": None, + "hyperparameters": {"n_epochs": 3}, + "model": "gpt-4o-mini-2024-07-18", + "organization_id": "org-pipeline", + "result_files": [], + "seed": 1, + "status": "queued", + "trained_tokens": None, + "training_file": "file-pipeline", + "validation_file": None, +} +_VIDEO: Final = { + "id": "video_pipeline", + "object": "video", + "model": "sora-2", + "status": "queued", + "created_at": 1, + "progress": 0, + "seconds": "4", + "size": "720x1280", +} +_SYNC_SUCCESS_RECORDER: Final = """ +import json +from pathlib import Path + +from litellm.integrations.custom_logger import CustomLogger + + +class SyncSuccessRecorder(CustomLogger): + def log_success_event(self, kwargs, response_obj, start_time, end_time): + metadata = (kwargs.get("standard_logging_object") or {{}}).get("metadata") or {{}} + row = {{"team": metadata.get("user_api_key_team_id"), "applied": metadata.get("applied_guardrails")}} + with Path({path!r}).open("a") as sink: + sink.write(json.dumps(row) + "\\n") + + +recorder = SyncSuccessRecorder() +""" def _completion_chunk(identity: str, delta: dict[str, str], finish_reason: str | None) -> bytes: @@ -183,8 +237,14 @@ def _completion(request: Request) -> Reply: return _bedrock_reply(request) if request.target == "/v1/responses": return _responses_reply(request) - if request.target == "/v1/images/generations": + if request.target in ("/v1/images/generations", "/v1/images/edits"): return Reply(body=json.dumps({"created": 1, "data": [{"b64_json": "aW1n", "revised_prompt": None}]}).encode()) + if request.target == "/v1/videos": + return Reply(body=json.dumps(_VIDEO).encode()) + if request.target == "/v1/fine_tuning/jobs": + return Reply(body=json.dumps(_FINE_TUNING_JOB).encode()) + if request.target.split("?", 1)[0].endswith(":generateContent"): + return Reply(body=json.dumps(_GEMINI_RESPONSE).encode()) assert request.target == "/v1/chat/completions", request.target identity: Final = f"chatcmpl-{uuid.uuid4().hex}" if json.loads(request.body).get("stream") is True: @@ -265,8 +325,17 @@ def _policy(mode: str, steps: tuple[tuple[str, str, str], ...], *, on_error: str def _config(directory: Path, vendor: Wire, upstream: Wire) -> Path: + (directory / "sync_success_recorder.py").write_text( + _SYNC_SUCCESS_RECORDER.format(path=str(directory / "sync-success-log.jsonl")) + ) configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) - configuration["litellm_settings"] = {**configuration.get("litellm_settings", {}), "callbacks": ["otel"]} + configuration["litellm_settings"] = { + **configuration.get("litellm_settings", {}), + "callbacks": ["otel", "sync_success_recorder.recorder"], + } + configuration["finetune_settings"] = [ + {"custom_llm_provider": "openai", "api_key": "sk-fixture", "api_base": f"{upstream.url}/v1"} + ] configuration["model_list"] = [ *configuration.get("model_list", []), { @@ -347,6 +416,9 @@ def _config(directory: Path, vendor: Wire, upstream: Wire) -> Path: _RESPONSE_STOPS_AFTER_FIRST: _policy( "post_call", ((_RESPONSE_FIRST, "allow", "block"), (_RESPONSE_SECOND, "next", "block")) ), + _RESPONSE_RUNS_EVERY_STEP: _policy( + "post_call", ((_RESPONSE_FIRST, "next", "block"), (_RESPONSE_SECOND, "next", "block")) + ), _SHIELDS_THEN_CHECKS: _policy( "pre_call", ((_PROMPT_SHIELD, "next", "block"), (_REQUEST_FIRST, "next", "block")) ), @@ -388,7 +460,7 @@ def rig(tmp_path_factory: pytest.TempPathFactory, audit_sinks: SpanSinks) -> Ite owned.scenario() as scenario, ): model: Final = scenario.model(api_base=f"{upstream.url}/v1", api_key="sk-fixture") - yield _Rig(owned, vendor, upstream, model, audit_sinks.operator) + yield _Rig(owned, vendor, upstream, model, audit_sinks.operator, directory / "sync-success-log.jsonl") def _body(path: str, model: str, stream: bool) -> dict[str, object]: @@ -708,6 +780,190 @@ def test_an_image_generation_names_its_response_pipeline_up_front_because_its_sp ], rows +def _applied_in_header_spend_log_and_trace( + rig: _Rig, since: int, team: str, response: httpx.Response, *, rows_expected: int = 1 +) -> tuple[list[str], list[list[str]], list[list[str]]]: + rows: Final = eventually( + lambda: rig.gateway.rows('SELECT metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s', (team,)), + lambda found: len(found) == rows_expected, + seconds=70, + ) + spans: Final = eventually( + lambda: _spans_naming_guardrails(rig, since, team), lambda found: len(found) >= rows_expected, seconds=30 + ) + return ( + sorted(name for name in response.headers.get(_HEADER, "").split(",") if name), + [sorted(object_value(row["metadata"]).get("applied_guardrails") or []) for row in rows], + [sorted(ast.literal_eval(str(span["attributes"]["metadata.applied_guardrails"]))) for span in spans], + ) + + +def _cache_hits_and_metadata_logged(rig: _Rig, team: str, count: int) -> list[dict[str, JsonValue]]: + return eventually( + lambda: rig.gateway.rows( + 'SELECT cache_hit, metadata FROM "LiteLLM_SpendLogs" WHERE team_id = %s ORDER BY "startTime"', (team,) + ), + lambda found: len(found) == count, + seconds=70, + ) + + +def test_a_cache_hit_names_every_response_step_that_ran_in_the_spend_log_and_the_trace(rig: _Rig) -> None: + since: Final = recorded_spans(rig.traces)[0] + rig.vendor.drain() + body: Final = {"model": rig.model, "messages": [{"role": "user", "content": f"cache me {uuid.uuid4().hex}"}]} + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"tags": [_RESPONSE_RUNS_EVERY_STEP]}) + key: Final = scenario.key(team_id=team) + first: Final = rig.gateway.request("POST", "/v1/chat/completions", body, key=key) + _cache_hits_and_metadata_logged(rig, team, 1) + cached: Final = rig.gateway.request("POST", "/v1/chat/completions", body, key=key) + rows: Final = _cache_hits_and_metadata_logged(rig, team, 2) + spans: Final = eventually(lambda: _spans_naming_guardrails(rig, since, team), lambda found: len(found) >= 2) + expected: Final = [_RESPONSE_FIRST, _RESPONSE_SECOND] + assert [first.status_code, cached.status_code] == [200, 200], (first.text, cached.text) + assert _steps_that_ran(rig) == (_RESPONSE_FIRST, _RESPONSE_SECOND) * 2 + assert [row["cache_hit"] for row in rows] == ["False", "True"], rows + assert [sorted(response.headers.get(_HEADER, "").split(",")) for response in (first, cached)] == [expected] * 2 + assert [sorted(object_value(row["metadata"]).get("applied_guardrails") or []) for row in rows] == [expected] * 2 + assert [sorted(ast.literal_eval(str(span["attributes"]["metadata.applied_guardrails"]))) for span in spans] == [ + expected + ] * len(spans), spans + + +def test_a_sync_success_callback_sees_every_response_step_that_ran(rig: _Rig) -> None: + team, response = _ask_as_new_team(rig, _RESPONSE_RUNS_EVERY_STEP) + assert response.status_code == 200, response.text + assert _steps_that_ran(rig) == (_RESPONSE_FIRST, _RESPONSE_SECOND) + logged: Final = eventually( + lambda: ( + [row for row in map(json.loads, rig.sync_success_log.read_text().splitlines()) if row["team"] == team] + if rig.sync_success_log.exists() + else [] + ), + lambda found: len(found) == 1, + seconds=30, + ) + assert sorted(logged[0]["applied"] or []) == [_RESPONSE_FIRST, _RESPONSE_SECOND], logged + + +def _image_edit(rig: _Rig, scenario_model: Callable[..., str], key: str) -> httpx.Response: + return rig.gateway.request_multipart( + "/v1/images/edits", + {"model": scenario_model(model="openai/gpt-image-1"), "prompt": f"a cat {uuid.uuid4().hex}"}, + {"image": ("cat.png", _PNG, "image/png")}, + key=key, + ) + + +def _video(rig: _Rig, scenario_model: Callable[..., str], key: str) -> httpx.Response: + return rig.gateway.request( + "POST", + "/v1/videos", + {"model": scenario_model(model="openai/sora-2"), "prompt": f"a cat {uuid.uuid4().hex}"}, + key=key, + ) + + +def _generate_content(rig: _Rig, scenario_model: Callable[..., str], key: str) -> httpx.Response: + return rig.gateway.request( + "POST", + f"/v1beta/models/{rig.model}:generateContent", + {"contents": [{"role": "user", "parts": [{"text": f"hello {uuid.uuid4().hex}"}]}]}, + key=key, + ) + + +def _stream_generate_content(rig: _Rig, scenario_model: Callable[..., str], key: str) -> httpx.Response: + return rig.gateway.request( + "POST", + f"/v1beta/models/{rig.model}:streamGenerateContent", + {"contents": [{"role": "user", "parts": [{"text": f"hello {uuid.uuid4().hex}"}]}]}, + key=key, + params={"alt": "sse"}, + ) + + +def _native_generate_content(rig: _Rig, scenario_model: Callable[..., str], key: str) -> httpx.Response: + model: Final = scenario_model(model="gemini/gemini-2.0-flash", api_base=rig.upstream.url) + return rig.gateway.request( + "POST", + f"/v1beta/models/{model}:generateContent", + {"contents": [{"role": "user", "parts": [{"text": f"hello {uuid.uuid4().hex}"}]}]}, + key=key, + ) + + +_PNG: Final = bytes.fromhex( + "89504e470d0a1a0a0000000d4948445200000001000000010806000000" + "1f15c4890000000d49444154789c6360000002000154a24f5d0000000049454e44ae426082" +) +_ROUTES: Final = ( + pytest.param(_image_edit, id="image-edit"), + pytest.param(_video, id="video"), + pytest.param(_generate_content, id="generate-content"), + pytest.param(_stream_generate_content, id="stream-generate-content"), + pytest.param(_native_generate_content, id="native-generate-content"), +) + + +@pytest.mark.parametrize("send", _ROUTES) +def test_a_response_pipeline_that_runs_every_step_names_both_steps_on_every_route( + rig: _Rig, send: Callable[[_Rig, Callable[..., str], str], httpx.Response] +) -> None: + since: Final = recorded_spans(rig.traces)[0] + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"tags": [_RESPONSE_RUNS_EVERY_STEP]}) + response: Final = send( + rig, + partial(scenario.model, api_base=f"{rig.upstream.url}/v1", api_key="sk-fixture"), + scenario.key(team_id=team), + ) + assert response.status_code == 200, response.text + expected: Final = [_RESPONSE_FIRST, _RESPONSE_SECOND] + header, rows, spans = _applied_in_header_spend_log_and_trace(rig, since, team, response) + assert (header, rows, spans) == (expected, [expected], [expected] * len(spans)) + + +@pytest.mark.parametrize("send", _ROUTES) +def test_a_response_pipeline_that_stops_early_names_the_same_steps_in_the_header_and_the_spend_log_on_every_route( + rig: _Rig, send: Callable[[_Rig, Callable[..., str], str], httpx.Response] +) -> None: + since: Final = recorded_spans(rig.traces)[0] + with rig.gateway.scenario() as scenario: + team: Final = scenario.team(metadata={"tags": [_RESPONSE_STOPS_AFTER_FIRST]}) + response: Final = send( + rig, + partial(scenario.model, api_base=f"{rig.upstream.url}/v1", api_key="sk-fixture"), + scenario.key(team_id=team), + ) + assert response.status_code == 200, response.text + header, rows, spans = _applied_in_header_spend_log_and_trace(rig, since, team, response) + assert _RESPONSE_FIRST in header, header + assert (rows, spans) == ([header], [header] * len(spans)) + + +def test_a_fine_tuning_job_with_a_method_object_and_a_guardrail_is_created(rig: _Rig) -> None: + rig.upstream.drain() + method: Final = {"type": "dpo", "dpo": {"hyperparameters": {"beta": 0.1}}} + with rig.gateway.scenario() as scenario: + key: Final = scenario.key(team_id=scenario.team(metadata={"tags": [_STOPS_AFTER_FIRST]})) + response: Final = rig.gateway.request( + "POST", + "/v1/fine_tuning/jobs", + { + "model": "gpt-4o-mini-2024-07-18", + "training_file": "file-pipeline", + "custom_llm_provider": "openai", + "method": method, + "guardrails": [_REQUEST_FIRST], + }, + key=key, + ) + assert response.status_code == 200, response.text + assert [request.target for request in rig.upstream.drain()] == ["/v1/fine_tuning/jobs"] + + def _agent_card(name: str, url: str) -> dict[str, object]: return { "protocolVersion": "0.3", diff --git a/tests/unit/proxy/test_common_request_processing.py b/tests/unit/proxy/test_common_request_processing.py index 8d01edb0096..82dd3479321 100644 --- a/tests/unit/proxy/test_common_request_processing.py +++ b/tests/unit/proxy/test_common_request_processing.py @@ -25,7 +25,7 @@ from litellm.litellm_core_utils.bug_report import ( build_bug_report, ) from litellm.types.router import RouterRateLimitError, RouterRateLimitErrorBasic -from litellm.types.utils import CachingDetails, CallTypesLiteral, Delta, ModelResponseStream, StreamingChoices +from litellm.types.utils import CachingDetails, Delta, ModelResponseStream, StreamingChoices from litellm.constants import ( CLIENT_REQUESTED_MODEL_SCOPE_KEY, MAX_LITELLM_CALL_ID_LENGTH, @@ -207,24 +207,6 @@ def test_include_guardrail_response_requested_is_false_without_exact_true(): assert include_guardrail_response_requested({}) is False -class _RecordsDeferral(CustomLogger): - def __init__(self) -> None: - super().__init__() - self.deferred: list[bool] = [] - - async def async_pre_call_hook( - self, - user_api_key_dict: ProxyUserAPIKeyAuth, - cache: DualCache, - data: dict[str, object], - call_type: CallTypesLiteral, - ) -> dict[str, object]: - logging_obj: Final = data["litellm_logging_obj"] - assert isinstance(logging_obj, LiteLLMLoggingObj) - self.deferred.append(logging_obj.defer_async_logging) - return data - - class TestProxyBaseLLMRequestProcessing: @pytest.mark.asyncio async def test_base_passthrough_process_llm_request_preserves_litellm_headers_for_non_streaming_response( @@ -501,30 +483,6 @@ class TestProxyBaseLLMRequestProcessing: pytest.fail("litellm_call_id is not a valid UUID") assert data_passed["litellm_call_id"] == returned_data["litellm_call_id"] - @pytest.mark.asyncio - @pytest.mark.parametrize("defer_async_logging", [False, True]) - async def test_common_processing_pre_call_logic_defers_the_success_log_before_pre_call_hook_runs( - self, monkeypatch: pytest.MonkeyPatch, defer_async_logging: bool - ) -> None: - from litellm.proxy.proxy_server import ProxyConfig as RealProxyConfig - - recorder: Final = _RecordsDeferral() - monkeypatch.setattr(litellm, "callbacks", [recorder]) - - await ProxyBaseLLMRequestProcessing( - data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]} - ).common_processing_pre_call_logic( - request=Request(scope={"type": "http", "method": "POST", "path": "/v1/chat/completions", "headers": []}), - general_settings={}, - user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"), - proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), - proxy_config=RealProxyConfig(), - route_type="acompletion", - defer_async_logging=defer_async_logging, - ) - - assert recorder.deferred == [defer_async_logging] - @pytest.mark.asyncio @pytest.mark.parametrize("requested_model", [{"bad": "value"}, ["gpt-5.2"], 1]) async def test_common_processing_pre_call_logic_rejects_a_non_string_model_with_400( @@ -7543,12 +7501,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit: processor = ProxyBaseLLMRequestProcessing(data={"model": primary_model}) call_count = 0 - deferred_per_call: Final[list[bool]] = [] async def mock_pre_call_logic(**kwargs): nonlocal call_count call_count += 1 - deferred_per_call.append(kwargs["defer_async_logging"]) model_in_data = processor.data.get("model") if model_in_data == primary_model: raise ProxyRateLimitError( @@ -7581,12 +7537,10 @@ class TestPreCallWithFallbacksOnLocalRateLimit: model=primary_model, route_type="acompletion", llm_router=mock_router, - defer_async_logging=True, ) assert processor.data["model"] == fallback_model assert call_count == 2 - assert deferred_per_call == [True, True] @pytest.mark.asyncio async def test_raises_when_no_fallbacks_configured(self): 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 5839bdcda4d..a0ba2723443 100644 --- a/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py +++ b/tests/unit/proxy/utils/proxy_logging/test_guardrail_pipeline.py @@ -281,7 +281,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(), "completion") + proxy_logging._process_guardrail_metadata(data) snapshot = { "call_count": len(calls), "first_name": calls[0]["name"], @@ -306,7 +306,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(), "completion") + proxy_logging._process_guardrail_metadata(data) assert calls == ["g2"] @@ -318,12 +318,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(), "completion") + proxy_logging._process_guardrail_metadata({}) def test_process_guardrail_metadata_invalid_data_raises(proxy_logging): with pytest.raises(AttributeError): - proxy_logging._process_guardrail_metadata(None, UserAPIKeyAuth(), "completion") # type: ignore[arg-type] + proxy_logging._process_guardrail_metadata(None) # type: ignore[arg-type] # --------------------------------------------------------------------------- @@ -1204,24 +1204,21 @@ def _two_step_pipeline_data( } -def _logging_obj(*, defer: bool) -> LiteLLMLoggingObj: +def _logging_obj(*, log_parked: bool, call_type: str = "acompletion") -> LiteLLMLoggingObj: logging_obj: Final = LiteLLMLoggingObj( model="m", messages=[], stream=False, - call_type="acompletion", + call_type=call_type, start_time=None, litellm_call_id="call-1", function_id="fn-1", ) - logging_obj.defer_async_logging = defer + logging_obj.defer_async_logging = True + logging_obj.enqueue_deferred_logging = (lambda: None) if log_parked else None return logging_obj -def _logged_after_response_hooks(data: dict[str, object]) -> dict[str, object]: - return {**data, "litellm_logging_obj": _logging_obj(defer=True)} - - @pytest.mark.asyncio @pytest.mark.parametrize( ("first_on_pass", "expected_calls", "expected_header"), @@ -1360,27 +1357,205 @@ async def test_pre_call_hook_names_a_step_that_ran_in_the_litellm_metadata_a_lat @pytest.mark.asyncio -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, +@pytest.mark.parametrize( + ("log_parked", "call_type", "bucket", "expected_after_response"), + [ + pytest.param(True, "acompletion", "metadata", ["first"], id="success-log-waits-for-the-response-hooks"), + pytest.param(True, "acompletion", "litellm_metadata", ["first"], id="success-log-waits-litellm-metadata"), + pytest.param(False, "acompletion", "metadata", ["first", "second"], id="success-log-already-built"), + pytest.param( + True, "allm_passthrough_route", "metadata", ["first", "second"], id="pass-through-headers-already-out" + ), + ], +) +async def test_post_call_pipeline_drops_the_steps_that_did_not_run_only_while_nothing_has_read_them( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + log_parked: bool, + call_type: str, + bucket: str, + expected_after_response: list[str], ) -> None: 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) + request: Final = _two_step_pipeline_data("post_call", "allow") data: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), - data=_logged_after_response_hooks(_two_step_pipeline_data("post_call", "allow")), + data={ + "model": request["model"], + "messages": request["messages"], + bucket: request["metadata"], + "litellm_logging_obj": _logging_obj(log_parked=log_parked, call_type=call_type), + }, call_type="completion", ) - assert "applied_guardrails" not in data["metadata"] + assert data[bucket]["applied_guardrails"] == ["first", "second"] + logged_metadata: Final = data[bucket].copy() await proxy_logging.post_call_success_hook( data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() ) assert (first.calls, second.calls) == (1, 0) - assert data["metadata"]["applied_guardrails"] == ["first"] + assert data[bucket]["applied_guardrails"] == expected_after_response + assert logged_metadata["applied_guardrails"] == expected_after_response + + +def _three_step_two_pipeline_data() -> dict[str, object]: + runs_second: Final = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="first", on_pass="next", on_fail="block"), + PipelineStep(guardrail="second", on_pass="next", on_fail="block"), + ], + ) + stops_before_second: Final = GuardrailPipeline( + mode="post_call", + steps=[ + PipelineStep(guardrail="third", on_pass="allow", on_fail="block"), + PipelineStep(guardrail="second", on_pass="next", on_fail="block"), + ], + ) + return { + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": ["first", "second", "third"], + "_guardrail_pipelines": [("runs-second", runs_second), ("stops-before-second", stops_before_second)], + "_pipeline_managed_guardrails": {"first", "second", "third"}, + }, + } + + +@pytest.mark.asyncio +async def test_post_call_pipeline_keeps_a_step_an_earlier_pipeline_ran_that_a_later_one_never_reached( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second", "third")] + monkeypatch.setattr(litellm, "callbacks", guardrails) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={ + **_three_step_two_pipeline_data(), + "litellm_logging_obj": _logging_obj(log_parked=True), + }, + call_type="completion", + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert [guardrail.calls for guardrail in guardrails] == [1, 1, 1] + assert data["metadata"]["applied_guardrails"] == ["first", "second", "third"] + + +@pytest.mark.asyncio +async def test_held_stream_pipeline_keeps_a_step_an_earlier_pipeline_ran_that_a_later_one_never_reached( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, +) -> None: + guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second", "third")] + monkeypatch.setattr(litellm, "callbacks", guardrails) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + logging_obj: Final = _logging_obj(log_parked=False) + + async def stream_log(*args: object) -> None: + return None + + logging_obj._on_deferred_stream_complete = stream_log + auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions") + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=auth, + data={ + **_three_step_two_pipeline_data(), + "stream": True, + "litellm_logging_obj": logging_obj, + }, + call_type="completion", + ) + + delivered: Final[list[object]] = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data + ) + ] + + assert len(delivered) == 2 + assert [guardrail.calls for guardrail in guardrails] == [1, 1, 1] + assert data["metadata"]["applied_guardrails"] == ["first", "second", "third"] + + +@pytest.mark.asyncio +async def test_post_call_pipeline_keeps_a_step_that_never_ran_whose_guardrail_ran_outside_the_pipeline( + proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, +) -> None: + first: Final = _PipelineStepGuardrail("first", "post_call") + second: Final = _PipelineStepGuardrail("second", ["pre_call", "post_call"]) + monkeypatch.setattr(litellm, "callbacks", [first, second]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + + 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"), "litellm_logging_obj": _logging_obj(log_parked=True)}, + call_type="completion", + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert (first.calls, second.calls) == (1, 1) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("steps", "expected"), + [(("first", "second"), ["second"]), (("first",), None)], + ids=["skipped-then-ran", "only-step-skipped"], +) +async def test_post_call_pipeline_drops_a_step_its_guardrail_skips_on_a_whole_response( + proxy_logging: ProxyLogging, + make_user_api_key_auth: Callable[..., UserAPIKeyAuth], + monkeypatch: pytest.MonkeyPatch, + steps: tuple[str, ...], + expected: list[str] | None, +) -> None: + guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in steps] + guardrails[0].apply_stream_scope("streaming") + monkeypatch.setattr(litellm, "callbacks", guardrails) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False) + pipeline: Final = GuardrailPipeline( + mode="post_call", steps=[PipelineStep(guardrail=name, on_pass="next", on_fail="block") for name in steps] + ) + + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=make_user_api_key_auth(), + data={ + "model": "m", + "messages": [{"role": "user", "content": "hi"}], + "metadata": { + "guardrails": list(steps), + "_guardrail_pipelines": [("governance", pipeline)], + "_pipeline_managed_guardrails": set(steps), + }, + "litellm_logging_obj": _logging_obj(log_parked=True), + }, + call_type="completion", + ) + await proxy_logging.post_call_success_hook( + data=data, response=litellm.ModelResponse(), user_api_key_dict=make_user_api_key_auth() + ) + + assert [guardrail.calls for guardrail in guardrails] == [0, 1][: len(steps)] + assert data["metadata"].get("applied_guardrails") == expected class _OpensLitellmMetadataOnResponseGuardrail(_PipelineStepGuardrail): @@ -1402,7 +1577,7 @@ async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_ data: Final = await proxy_logging.pre_call_hook( user_api_key_dict=make_user_api_key_auth(), - data=_logged_after_response_hooks(_two_step_pipeline_data("post_call", "next")), + data={**_two_step_pipeline_data("post_call", "next"), "litellm_logging_obj": _logging_obj(log_parked=True)}, call_type="completion", ) await proxy_logging.post_call_success_hook( @@ -1415,78 +1590,47 @@ async def test_post_call_pipeline_names_every_step_that_ran_in_metadata_after_a_ @pytest.mark.asyncio @pytest.mark.parametrize( - ("asks_for", "request_route", "streams_live", "expected_header"), - [ - pytest.param({"stream": False}, "/custom/stream", False, None, id="whole-response"), - pytest.param( - {"litellm_logging_obj": None}, - "/v1/images/generations", - False, - ["first", "second"], - id="whole-response-logged-before-it-runs", - ), - pytest.param( - {"litellm_logging_obj": _logging_obj(defer=False)}, - "/v1/chat/completions", - False, - ["first", "second"], - id="whole-response-log-not-deferred", - ), - pytest.param({"stream": True}, "/v1/chat/completions", False, None, id="stream-held-for-the-pipeline"), - pytest.param({"stream": True}, "/v1/chat/completions", True, ["first", "second"], id="stream-released-live"), - pytest.param( - {"stream": True}, "/custom/stream", False, ["first", "second"], id="stream-runs-steps-on-their-own" - ), - pytest.param({"method": "message/send"}, "/a2a/agent", False, None, id="a2a-whole-response"), - pytest.param({"method": "message/stream"}, "/a2a/agent", False, ["first", "second"], id="a2a-stream"), - pytest.param({"method": "SubscribeToTask"}, "/a2a/agent", False, ["first", "second"], id="a2a-v1-resubscribe"), - ], + ("log_waits", "expected_after_stream"), + [(True, ["first"]), (False, ["first", "second"])], + ids=["stream-log-waits-for-the-pipeline", "stream-log-not-deferred"], ) -async def test_pre_call_hook_names_up_front_the_post_call_steps_a_stream_runs_after_its_headers( - proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, - asks_for: dict[str, object], - request_route: str, - streams_live: bool, - expected_header: list[str] | None, -) -> None: - 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: Final = _logged_after_response_hooks( - _two_step_pipeline_data("post_call", "next", first_on_fail="next", second_on_fail="next") - ) - data.update(asks_for) - - 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", - ) - - 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( +async def test_held_stream_pipeline_drops_the_steps_that_did_not_run_only_while_its_log_waits( proxy_logging: ProxyLogging, make_user_api_key_auth: Callable[..., UserAPIKeyAuth], monkeypatch: pytest.MonkeyPatch, - stream: bool, + log_waits: bool, + expected_after_stream: list[str], ) -> None: - guardrails: Final = [_PipelineStepGuardrail(name, "post_call") for name in ("first", "second")] - monkeypatch.setattr(litellm, "callbacks", guardrails) - data: Final = _logged_after_response_hooks(_two_step_pipeline_data("post_call", "allow")) - data["stream"] = stream + 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) + logging_obj: Final = _logging_obj(log_parked=False) - 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", + async def stream_log(*args: object) -> None: + return None + + logging_obj._on_deferred_stream_complete = stream_log if log_waits else None + auth: Final = make_user_api_key_auth(request_route="/v1/chat/completions") + data: Final = await proxy_logging.pre_call_hook( + user_api_key_dict=auth, + data={**_two_step_pipeline_data("post_call", "allow"), "stream": True, "litellm_logging_obj": logging_obj}, + call_type="completion", ) + assert data["metadata"]["applied_guardrails"] == ["first", "second"] + logged_metadata: Final = data["metadata"].copy() - assert result["metadata"].get("applied_guardrails") == ["first", "second"] + delivered: Final[list[object]] = [ + item + async for item in proxy_logging.async_post_call_streaming_iterator_hook( + user_api_key_dict=auth, response=_async_chunk_iter(_stream_chunks()), request_data=data + ) + ] + + assert len(delivered) == 2 + assert (first.calls, second.calls) == (1, 0) + assert data["metadata"]["applied_guardrails"] == expected_after_stream + assert logged_metadata["applied_guardrails"] == expected_after_stream @pytest.mark.asyncio 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 89a59d44380..b7fc8c6dbeb 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, call_type: str) -> None: + def fake_process(d): invoked["data"] = d proxy_logging._process_guardrail_metadata = fake_process # type: ignore[assignment]