mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(guardrails): name every response step up front and drop the ones that did not run while nothing has read them
The success log of a cache hit, a sync callback, image edits, videos, generateContent and other routes is built before or beside the response pipeline, so leaving response steps out of the up-front fill lost steps that ran. Every response step is named up front again, as on main, and the steps that never ran are removed after the pipeline only while the success log is still parked (whole responses outside the LLM pass-through route) or the held stream's log is still waiting. The removal is in place so the logging object's metadata copy sees it. This drops the early deferral flag and the A2A method check, which also fixes the 500 on fine-tuning jobs that send a method object.
This commit is contained in:
yucheng
parent
076e3e374c
commit
05d2a0c63f
8 changed files with 602 additions and 241 deletions
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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).
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue