mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
Merge pull request #40274 from BerriAI/litellm_post_call_pipeline_background_responses
feat(guardrails): run post_call policy pipelines on background Responses retrieval
This commit is contained in:
Mateo Wang
GitHub
commit
11b31c19be
6 changed files with 1034 additions and 113 deletions
|
|
@ -190,6 +190,7 @@ from litellm.proxy.litellm_pre_call_utils import (
|
|||
refresh_proxy_server_request_body_snapshot,
|
||||
reject_url_valued_destination,
|
||||
)
|
||||
from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval
|
||||
from litellm.types.utils import (
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
|
|
@ -1849,7 +1850,6 @@ class ProxyBaseLLMRequestProcessing:
|
|||
data=self.data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
# Calculate request queue time after add_litellm_data_to_request
|
||||
# which sets arrival_time in proxy_server_request. Ends at start_time
|
||||
# (not a freshly captured time.time() here) so this window is exactly
|
||||
|
|
@ -1997,6 +1997,12 @@ class ProxyBaseLLMRequestProcessing:
|
|||
data=self.data,
|
||||
call_type=route_type,
|
||||
)
|
||||
if route_type == "aget_responses":
|
||||
attach_post_call_pipelines_to_retrieval(
|
||||
data=self.data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
|
||||
# Refresh AFTER pre_call_hook: guardrails (e.g. Presidio PII masking) may
|
||||
# have mutated `self.data` in place, and the audit-trail snapshot taken in
|
||||
|
|
|
|||
152
litellm/proxy/policy_engine/response_retrieval.py
Normal file
152
litellm/proxy/policy_engine/response_retrieval.py
Normal file
|
|
@ -0,0 +1,152 @@
|
|||
from collections.abc import Mapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
|
||||
from litellm.proxy.common_utils.callback_utils import (
|
||||
add_guardrail_to_applied_guardrails_header,
|
||||
add_policy_sources_to_metadata,
|
||||
add_policy_to_applied_policies_header,
|
||||
)
|
||||
from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
|
||||
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
|
||||
from litellm.proxy.policy_engine.policy_matcher import PolicyMatcher
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.router_utils.common_utils import resolve_model_group_alias
|
||||
from litellm.types.proxy.policy_engine import PolicyMatchContext
|
||||
from litellm.types.proxy.policy_engine.pipeline_types import GuardrailPipeline
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.router import Router
|
||||
|
||||
PolicyPipelines: TypeAlias = tuple[tuple[str, GuardrailPipeline], ...]
|
||||
|
||||
_POLICY_PIPELINES_ADAPTER: Final = TypeAdapter(PolicyPipelines)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class UngovernedRetrieval:
|
||||
reason: Literal["no router", "response id names no deployment", "deployment no longer in the router"]
|
||||
|
||||
|
||||
def _model_group_for_response_id(response_id: object, llm_router: "Router | None") -> str | UngovernedRetrieval:
|
||||
if llm_router is None:
|
||||
return UngovernedRetrieval("no router")
|
||||
model_id: Final = (
|
||||
ResponsesAPIRequestUtils.get_model_id_from_response_id(response_id) if isinstance(response_id, str) else None
|
||||
)
|
||||
if model_id is None:
|
||||
return UngovernedRetrieval("response id names no deployment")
|
||||
deployment: Final = llm_router.get_deployment(model_id)
|
||||
if deployment is None:
|
||||
return UngovernedRetrieval("deployment no longer in the router")
|
||||
hidden_by: Final = _submit_model_hidden_by(deployment.model_name, llm_router.model_group_alias)
|
||||
if hidden_by is not None:
|
||||
verbose_proxy_logger.warning(
|
||||
"Policy engine: background response %s re-matches policies on retrieval as model group %s (%s), "
|
||||
"so a policy attached to the model name it was submitted as does not run on it",
|
||||
response_id,
|
||||
deployment.model_name,
|
||||
hidden_by,
|
||||
)
|
||||
return deployment.model_name
|
||||
|
||||
|
||||
def _submit_model_hidden_by(model_group: str, model_group_alias: Mapping[str, object]) -> str | None:
|
||||
if "*" in model_group:
|
||||
return "a wildcard deployment"
|
||||
aliases: Final = tuple(
|
||||
alias for alias in model_group_alias if resolve_model_group_alias(model_group_alias, alias) == model_group
|
||||
)
|
||||
if not aliases:
|
||||
return None
|
||||
return f"the target of model_group_alias {', '.join(aliases)}"
|
||||
|
||||
|
||||
def _retrieval_context(
|
||||
data: Mapping[str, object], user_api_key_dict: "UserAPIKeyAuth", model_group: str
|
||||
) -> PolicyMatchContext:
|
||||
team_alias: Final = user_api_key_dict.team_alias
|
||||
key_alias: Final = user_api_key_dict.key_alias
|
||||
return PolicyMatchContext(
|
||||
team_alias=team_alias if isinstance(team_alias, str) else None,
|
||||
key_alias=key_alias if isinstance(key_alias, str) else None,
|
||||
model=model_group,
|
||||
tags=get_tags_from_request_body(data) or None,
|
||||
)
|
||||
|
||||
|
||||
def _post_call_pipelines_for_context(context: PolicyMatchContext) -> tuple[PolicyPipelines, Mapping[str, str]]:
|
||||
matches: Final = get_attachment_registry().get_attached_policies_with_reasons(context)
|
||||
if not matches:
|
||||
return (), MappingProxyType({})
|
||||
applied_policy_names: Final = PolicyMatcher.get_policies_with_matching_conditions(
|
||||
policy_names=[match["policy_name"] for match in matches], # mutable-ok: the matcher takes a list
|
||||
context=context,
|
||||
)
|
||||
post_call_pipelines: Final = tuple(
|
||||
(policy_name, pipeline)
|
||||
for policy_name, pipeline in PolicyResolver.resolve_pipelines_for_context(
|
||||
context=context, policy_names=applied_policy_names
|
||||
)
|
||||
if pipeline.mode == "post_call"
|
||||
)
|
||||
return post_call_pipelines, MappingProxyType({match["policy_name"]: match["matched_via"] for match in matches})
|
||||
|
||||
|
||||
def attach_post_call_pipelines_to_retrieval(
|
||||
data: dict[str, object], # mutable-ok: request-state dict the policy engine hooks all write in place
|
||||
user_api_key_dict: "UserAPIKeyAuth",
|
||||
llm_router: "Router | None",
|
||||
) -> None:
|
||||
if not get_policy_registry().is_initialized():
|
||||
return
|
||||
model_group: Final = _model_group_for_response_id(data.get("response_id"), llm_router)
|
||||
if isinstance(model_group, UngovernedRetrieval):
|
||||
verbose_proxy_logger.warning(
|
||||
"Policy engine: background response %s is retrieved without its post_call policy pipelines (%s)",
|
||||
data.get("response_id"),
|
||||
model_group.reason,
|
||||
)
|
||||
return
|
||||
context: Final = _retrieval_context(data, user_api_key_dict, model_group)
|
||||
post_call_pipelines, policy_sources = _post_call_pipelines_for_context(context)
|
||||
_, bucket = get_or_create_metadata_bucket(data)
|
||||
already_attached: Final = _POLICY_PIPELINES_ADAPTER.validate_python(bucket.get("_guardrail_pipelines") or ())
|
||||
attached_policy_names: Final = frozenset(policy_name for policy_name, _pipeline in already_attached)
|
||||
added: Final = tuple(
|
||||
(policy_name, pipeline)
|
||||
for policy_name, pipeline in post_call_pipelines
|
||||
if policy_name not in attached_policy_names
|
||||
)
|
||||
if not added:
|
||||
return
|
||||
pipelines: Final = (*already_attached, *added)
|
||||
bucket["_guardrail_pipelines"] = pipelines
|
||||
bucket["_pipeline_managed_guardrails"] = frozenset(
|
||||
step.guardrail for _policy_name, pipeline in pipelines for step in pipeline.steps
|
||||
)
|
||||
for policy_name, _pipeline in added:
|
||||
add_policy_to_applied_policies_header(request_data=data, policy_name=policy_name)
|
||||
for _policy_name, pipeline in added:
|
||||
for step in pipeline.steps:
|
||||
add_guardrail_to_applied_guardrails_header(request_data=data, guardrail_name=step.guardrail)
|
||||
add_policy_sources_to_metadata(
|
||||
request_data=data,
|
||||
policy_sources={ # mutable-ok: add_policy_sources_to_metadata takes a dict
|
||||
policy_name: policy_sources[policy_name] for policy_name, _pipeline in added
|
||||
},
|
||||
)
|
||||
verbose_proxy_logger.debug(
|
||||
"Policy engine: attached post_call pipelines to the retrieval of background response %s (model group %s): %s",
|
||||
data.get("response_id"),
|
||||
model_group,
|
||||
", ".join(policy_name for policy_name, _pipeline in added),
|
||||
)
|
||||
|
|
@ -94,6 +94,7 @@ from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
|
|||
from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
coerce_token_limit,
|
||||
get_or_create_metadata_bucket,
|
||||
independent_snapshot,
|
||||
is_expected_client_error,
|
||||
)
|
||||
|
|
@ -157,6 +158,8 @@ from litellm.proxy.hooks.sensitive_data_routing import (
|
|||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup, add_guardrails_from_auth_metadata
|
||||
from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at
|
||||
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
from litellm.proxy.policy_engine.policy_resolver import PolicyResolver
|
||||
from litellm.repositories.budget_repository import BudgetRepository
|
||||
from litellm.repositories.config_repository import ConfigRepository
|
||||
from litellm.repositories.table_repositories import (
|
||||
|
|
@ -172,6 +175,7 @@ from litellm.repositories.verification_token_repository import (
|
|||
)
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.mcp import (
|
||||
MCPDuringCallResponseObject,
|
||||
MCPPreCallRequestObject,
|
||||
|
|
@ -529,17 +533,127 @@ def _post_call_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, "Guardr
|
|||
)
|
||||
|
||||
|
||||
def _warn_background_skips_post_call_pipelines(data: Mapping[str, object]) -> None:
|
||||
if data.get("background") is not True:
|
||||
return
|
||||
policy_names: Final = tuple(policy_name for policy_name, _pipeline in _post_call_pipelines(data))
|
||||
if not policy_names:
|
||||
return
|
||||
verbose_proxy_logger.warning(
|
||||
"Policies with post_call guardrail pipelines do not run on background responses yet; "
|
||||
"the response is released ungoverned by them: %s",
|
||||
", ".join(policy_names),
|
||||
_PENDING_BACKGROUND_RESPONSE_STATUSES: Final = frozenset(("queued", "in_progress"))
|
||||
|
||||
|
||||
def _is_pending_background_response(response: LLMResponseTypes) -> bool:
|
||||
return isinstance(response, ResponsesAPIResponse) and response.status in _PENDING_BACKGROUND_RESPONSE_STATUSES
|
||||
|
||||
|
||||
def _guardrails_outside_pipeline(policy_name: str, pipeline: "GuardrailPipeline") -> frozenset[str]:
|
||||
resolved: Final = PolicyResolver.resolve_policy_guardrails(
|
||||
policy_name=policy_name, policies=get_policy_registry().get_all_policies()
|
||||
)
|
||||
return frozenset(resolved.guardrails) - frozenset(step.guardrail for step in pipeline.steps)
|
||||
|
||||
|
||||
def _guardrails_run_standalone_pre_call(data: Mapping[str, object]) -> frozenset[str]:
|
||||
return frozenset(
|
||||
callback.guardrail_name
|
||||
for callback in litellm.callbacks
|
||||
if isinstance(callback, CustomGuardrail)
|
||||
and callback.guardrail_name is not None
|
||||
and callback.should_run_guardrail(data=data, event_type=GuardrailEventHooks.pre_call)
|
||||
)
|
||||
|
||||
|
||||
def _without_names(
|
||||
bucket: dict[str, object], # mutable-ok: the applied_* header slots live in the request-state dict hooks write
|
||||
slot: str,
|
||||
names: frozenset[str],
|
||||
) -> None:
|
||||
claimed: Final = bucket.get(slot)
|
||||
if not isinstance(claimed, list):
|
||||
return
|
||||
remaining: Final = [ # mutable-ok: the slot stays a list, the shape every applied_* header writer appends to
|
||||
name for name in claimed if name not in names
|
||||
]
|
||||
if remaining:
|
||||
bucket[slot] = remaining # rebind-ok: the slot lives in the shared request-state dict, rewritten in place
|
||||
else:
|
||||
bucket.pop(slot)
|
||||
|
||||
|
||||
def _withdraw_deferred_claims(
|
||||
data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data
|
||||
deferred: Sequence[tuple[str, "GuardrailPipeline"]],
|
||||
) -> None:
|
||||
outside_by_policy: Final = MappingProxyType(
|
||||
{policy_name: _guardrails_outside_pipeline(policy_name, pipeline) for policy_name, pipeline in deferred}
|
||||
)
|
||||
running_elsewhere: Final = _pipeline_managed_guardrail_names(data, "pre_call").union(
|
||||
_guardrails_run_standalone_pre_call(data), *outside_by_policy.values()
|
||||
)
|
||||
withdrawn_policies: Final = frozenset(name for name, outside in outside_by_policy.items() if not outside)
|
||||
withdrawn_guardrails: Final = _pipeline_step_guardrail_names(deferred) - running_elsewhere
|
||||
_, bucket = get_or_create_metadata_bucket(data)
|
||||
_without_names(bucket, "applied_policies", withdrawn_policies)
|
||||
_without_names(bucket, "applied_guardrails", withdrawn_guardrails)
|
||||
sources: Final = bucket.get("policy_sources")
|
||||
if not isinstance(sources, dict):
|
||||
return
|
||||
remaining_sources: Final = { # mutable-ok: policy_sources stays a dict, the shape its writer updates in place
|
||||
name: reason for name, reason in sources.items() if name not in withdrawn_policies
|
||||
}
|
||||
if remaining_sources:
|
||||
bucket["policy_sources"] = remaining_sources
|
||||
else:
|
||||
bucket.pop("policy_sources")
|
||||
|
||||
|
||||
def _defer_post_call_pipelines(
|
||||
data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data
|
||||
response: ResponsesAPIResponse,
|
||||
) -> None:
|
||||
deferred: Final = _post_call_pipelines(data)
|
||||
if not deferred:
|
||||
return
|
||||
verbose_proxy_logger.debug(
|
||||
"Post_call guardrail pipelines wait for background response %s (status=%s) to be retrieved complete: %s",
|
||||
response.id,
|
||||
response.status,
|
||||
", ".join(policy_name for policy_name, _pipeline in deferred),
|
||||
)
|
||||
tag_matched: Final = _tag_matched_deferrals(data, deferred)
|
||||
if tag_matched:
|
||||
verbose_proxy_logger.warning(
|
||||
"Policy engine: background response %s matched post_call policies through a request tag at submit; "
|
||||
"retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body "
|
||||
"does not govern the completed response: %s",
|
||||
response.id,
|
||||
", ".join(tag_matched),
|
||||
)
|
||||
body_selected: Final = _body_selected_deferrals(data, deferred)
|
||||
if body_selected:
|
||||
verbose_proxy_logger.warning(
|
||||
"Policy engine: background response %s matched post_call policies through the request body's policies "
|
||||
"list at submit; retrieval carries no request body, so those policies do not govern the completed "
|
||||
"response: %s",
|
||||
response.id,
|
||||
", ".join(body_selected),
|
||||
)
|
||||
_withdraw_deferred_claims(data, deferred)
|
||||
|
||||
|
||||
def _tag_matched_deferrals(
|
||||
data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]]
|
||||
) -> tuple[str, ...]:
|
||||
sources: Final = _policy_state_metadata(data).get("policy_sources")
|
||||
if not isinstance(sources, dict):
|
||||
return ()
|
||||
return tuple(
|
||||
policy_name
|
||||
for policy_name, _pipeline in deferred
|
||||
if policy_name in sources and "tag:" in str(sources[policy_name])
|
||||
)
|
||||
|
||||
|
||||
def _body_selected_deferrals(
|
||||
data: Mapping[str, object], deferred: Sequence[tuple[str, "GuardrailPipeline"]]
|
||||
) -> tuple[str, ...]:
|
||||
sources: Final = _policy_state_metadata(data).get("policy_sources")
|
||||
attributed: Final = frozenset(sources) if isinstance(sources, dict) else frozenset()
|
||||
return tuple(policy_name for policy_name, _pipeline in deferred if policy_name not in attributed)
|
||||
|
||||
|
||||
def _pipeline_is_streamable(policy_name: str, pipeline: "GuardrailPipeline") -> bool:
|
||||
|
|
@ -1985,8 +2099,6 @@ class ProxyLogging:
|
|||
)
|
||||
|
||||
try:
|
||||
_warn_background_skips_post_call_pipelines(data)
|
||||
|
||||
# Execute guardrail pipelines before the normal callback loop
|
||||
data, _ = await self._maybe_execute_pipelines( # rebind-ok: pipeline edits feed the callback loop below
|
||||
data=data,
|
||||
|
|
@ -2956,6 +3068,24 @@ class ProxyLogging:
|
|||
daemon=True,
|
||||
).start()
|
||||
|
||||
async def _run_post_call_pipelines(
|
||||
self,
|
||||
data: dict[str, object], # mutable-ok: same request-payload shape as post_call_success_hook's data
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
response: LLMResponseTypes,
|
||||
) -> LLMResponseTypes | None:
|
||||
if _is_pending_background_response(response):
|
||||
_defer_post_call_pipelines(data, response)
|
||||
return None
|
||||
_, pipeline_response = await self._maybe_execute_pipelines(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion",
|
||||
event_hook="post_call",
|
||||
response=response,
|
||||
)
|
||||
return pipeline_response
|
||||
|
||||
async def post_call_success_hook(
|
||||
self,
|
||||
data: dict,
|
||||
|
|
@ -2975,11 +3105,9 @@ class ProxyLogging:
|
|||
from litellm.proxy.proxy_server import llm_router
|
||||
from litellm.types.guardrails import GuardrailEventHooks
|
||||
|
||||
_, pipeline_response = await self._maybe_execute_pipelines(
|
||||
pipeline_response: Final = await self._run_post_call_pipelines(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_type=getattr(data.get("litellm_logging_obj"), "call_type", None) or "acompletion",
|
||||
event_hook="post_call",
|
||||
response=response,
|
||||
)
|
||||
if pipeline_response is not None:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,262 @@
|
|||
import logging
|
||||
from collections.abc import Iterator, Mapping
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
from litellm.proxy.policy_engine.response_retrieval import attach_post_call_pipelines_to_retrieval
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
from litellm.types.router import Deployment, LiteLLM_Params
|
||||
|
||||
GOVERNED_MODEL_GROUP = "gpt-5.4-mini"
|
||||
GOVERNED_MODEL_ID = "deployment-governed"
|
||||
UNGOVERNED_MODEL_GROUP = "gpt-4.1-mini"
|
||||
UNGOVERNED_MODEL_ID = "deployment-ungoverned"
|
||||
WILDCARD_MODEL_GROUP = "openai/*"
|
||||
WILDCARD_MODEL_ID = "deployment-wildcard"
|
||||
|
||||
|
||||
class FakeRouter:
|
||||
def __init__(self, deployments: dict[str, Deployment], model_group_alias: dict[str, object] | None = None):
|
||||
self._deployments = deployments
|
||||
self.model_group_alias = model_group_alias or {}
|
||||
|
||||
def get_deployment(self, model_id: str) -> Deployment | None:
|
||||
return self._deployments.get(model_id)
|
||||
|
||||
|
||||
def _deployment(model_group: str, model_id: str) -> Deployment:
|
||||
return Deployment(
|
||||
model_name=model_group,
|
||||
litellm_params=LiteLLM_Params(model=f"openai/{model_group}"),
|
||||
model_info={"id": model_id},
|
||||
)
|
||||
|
||||
|
||||
def _router(model_group_alias: dict[str, object] | None = None) -> FakeRouter:
|
||||
return FakeRouter(
|
||||
{
|
||||
GOVERNED_MODEL_ID: _deployment(GOVERNED_MODEL_GROUP, GOVERNED_MODEL_ID),
|
||||
UNGOVERNED_MODEL_ID: _deployment(UNGOVERNED_MODEL_GROUP, UNGOVERNED_MODEL_ID),
|
||||
WILDCARD_MODEL_ID: _deployment(WILDCARD_MODEL_GROUP, WILDCARD_MODEL_ID),
|
||||
},
|
||||
model_group_alias,
|
||||
)
|
||||
|
||||
|
||||
def _encoded_response_id(model_id: str) -> str:
|
||||
return ResponsesAPIRequestUtils._build_responses_api_response_id(
|
||||
custom_llm_provider="openai", model_id=model_id, response_id="resp_upstream"
|
||||
)
|
||||
|
||||
|
||||
def _pipeline_policy(guardrail: str, mode: str = "post_call") -> dict[str, object]:
|
||||
return {
|
||||
"guardrails": {"add": [guardrail]},
|
||||
"pipeline": {"mode": mode, "steps": [{"guardrail": guardrail, "on_pass": "allow", "on_fail": "block"}]},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def policy_engine() -> Iterator[None]:
|
||||
policy_registry = get_policy_registry()
|
||||
attachment_registry = get_attachment_registry()
|
||||
policy_registry.load_policies(
|
||||
{
|
||||
"response-governance": _pipeline_policy("output-word-filter"),
|
||||
"input-governance": _pipeline_policy("input-word-filter", mode="pre_call"),
|
||||
"team-governance": _pipeline_policy("team-word-filter"),
|
||||
"tag-governance": _pipeline_policy("tag-word-filter"),
|
||||
}
|
||||
)
|
||||
attachment_registry.load_attachments(
|
||||
[
|
||||
{"policy": "response-governance", "models": [GOVERNED_MODEL_GROUP]},
|
||||
{"policy": "input-governance", "models": [GOVERNED_MODEL_GROUP]},
|
||||
{"policy": "team-governance", "teams": ["governed-team"]},
|
||||
{"policy": "tag-governance", "tags": ["governed"]},
|
||||
]
|
||||
)
|
||||
yield
|
||||
policy_registry.clear()
|
||||
attachment_registry.clear()
|
||||
|
||||
|
||||
def _retrieval_data(model_id: str) -> dict[str, object]:
|
||||
return {"response_id": _encoded_response_id(model_id), "litellm_metadata": {}}
|
||||
|
||||
|
||||
def _attached_pipelines(data: Mapping[str, object]) -> tuple[tuple[str, str], ...]:
|
||||
bucket = data["litellm_metadata"]
|
||||
assert isinstance(bucket, dict)
|
||||
return tuple(
|
||||
(policy_name, ",".join(step.guardrail for step in pipeline.steps))
|
||||
for policy_name, pipeline in bucket["_guardrail_pipelines"]
|
||||
)
|
||||
|
||||
|
||||
def test_attaches_model_scoped_post_call_pipeline_to_retrieval(policy_engine: None) -> None:
|
||||
data = _retrieval_data(GOVERNED_MODEL_ID)
|
||||
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
|
||||
|
||||
assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),)
|
||||
assert data["litellm_metadata"]["_pipeline_managed_guardrails"] == frozenset({"output-word-filter"})
|
||||
assert data["litellm_metadata"]["applied_policies"] == ["response-governance"]
|
||||
assert data["litellm_metadata"]["applied_guardrails"] == ["output-word-filter"]
|
||||
assert data["litellm_metadata"]["policy_sources"] == {"response-governance": "model:gpt-5.4-mini"}
|
||||
assert "model" not in data
|
||||
assert "guardrails" not in data["litellm_metadata"]
|
||||
|
||||
|
||||
def test_key_and_team_context_also_governs_retrieval(policy_engine: None) -> None:
|
||||
data = _retrieval_data(UNGOVERNED_MODEL_ID)
|
||||
|
||||
attach_post_call_pipelines_to_retrieval(
|
||||
data=data, user_api_key_dict=UserAPIKeyAuth(team_alias="governed-team"), llm_router=_router()
|
||||
)
|
||||
|
||||
assert _attached_pipelines(data) == (("team-governance", "team-word-filter"),)
|
||||
|
||||
|
||||
def test_tag_attached_policy_is_not_re_matched_when_the_retrieval_carries_no_tag(policy_engine: None) -> None:
|
||||
data = _retrieval_data(GOVERNED_MODEL_ID)
|
||||
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
|
||||
|
||||
assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),)
|
||||
|
||||
|
||||
def test_tag_attached_policy_governs_a_retrieval_whose_metadata_carries_the_tag(policy_engine: None) -> None:
|
||||
data: dict[str, object] = {
|
||||
"response_id": _encoded_response_id(UNGOVERNED_MODEL_ID),
|
||||
"litellm_metadata": {"tags": ["governed"]},
|
||||
}
|
||||
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
|
||||
|
||||
assert _attached_pipelines(data) == (("tag-governance", "tag-word-filter"),)
|
||||
assert data["litellm_metadata"]["policy_sources"] == {"tag-governance": "tag:governed"}
|
||||
|
||||
|
||||
def test_retrieval_of_an_ungoverned_model_attaches_nothing(policy_engine: None) -> None:
|
||||
data = _retrieval_data(UNGOVERNED_MODEL_ID)
|
||||
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
|
||||
|
||||
assert data == _retrieval_data(UNGOVERNED_MODEL_ID)
|
||||
|
||||
|
||||
def test_already_attached_policy_is_not_attached_twice(policy_engine: None) -> None:
|
||||
data = _retrieval_data(GOVERNED_MODEL_ID)
|
||||
router = _router()
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router)
|
||||
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router)
|
||||
|
||||
assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),)
|
||||
assert data["litellm_metadata"]["applied_policies"] == ["response-governance"]
|
||||
|
||||
|
||||
def _hidden_submit_model_warnings(caplog: pytest.LogCaptureFixture) -> list[str]:
|
||||
return [
|
||||
record.getMessage()
|
||||
for record in caplog.records
|
||||
if record.levelno == logging.WARNING and "the model name it was submitted as" in record.getMessage()
|
||||
]
|
||||
|
||||
|
||||
def test_wildcard_deployment_attaches_nothing_for_the_submitted_model_and_warns(
|
||||
policy_engine: None, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
data = _retrieval_data(WILDCARD_MODEL_ID)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
|
||||
|
||||
assert data == _retrieval_data(WILDCARD_MODEL_ID)
|
||||
assert [
|
||||
"as model group openai/* (a wildcard deployment)" in message
|
||||
for message in _hidden_submit_model_warnings(caplog)
|
||||
] == [True]
|
||||
|
||||
|
||||
def test_aliased_model_group_still_attaches_its_own_policies_and_warns(
|
||||
policy_engine: None, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
data = _retrieval_data(GOVERNED_MODEL_ID)
|
||||
router = _router({"gpt-mini": GOVERNED_MODEL_GROUP, "gpt-hidden": {"model": GOVERNED_MODEL_GROUP, "hidden": True}})
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=router)
|
||||
|
||||
assert _attached_pipelines(data) == (("response-governance", "output-word-filter"),)
|
||||
assert [
|
||||
"(the target of model_group_alias gpt-mini, gpt-hidden)" in message
|
||||
for message in _hidden_submit_model_warnings(caplog)
|
||||
] == [True]
|
||||
|
||||
|
||||
def test_plain_model_group_retrieval_does_not_warn_about_the_submitted_model(
|
||||
policy_engine: None, caplog: pytest.LogCaptureFixture
|
||||
) -> None:
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
attach_post_call_pipelines_to_retrieval(
|
||||
data=_retrieval_data(GOVERNED_MODEL_ID),
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
llm_router=_router({"other-alias": UNGOVERNED_MODEL_GROUP}),
|
||||
)
|
||||
|
||||
assert _hidden_submit_model_warnings(caplog) == []
|
||||
|
||||
|
||||
def _ungoverned_retrieval_warnings(caplog: pytest.LogCaptureFixture) -> list[str]:
|
||||
return [
|
||||
record.getMessage()
|
||||
for record in caplog.records
|
||||
if record.levelno == logging.WARNING
|
||||
and "retrieved without its post_call policy pipelines" in record.getMessage()
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("response_id", "reason"),
|
||||
[
|
||||
("resp_plain_upstream_id", "response id names no deployment"),
|
||||
(_encoded_response_id("deployment-missing-from-router"), "deployment no longer in the router"),
|
||||
(None, "response id names no deployment"),
|
||||
],
|
||||
)
|
||||
def test_unresolvable_response_id_attaches_nothing_and_warns(
|
||||
policy_engine: None, caplog: pytest.LogCaptureFixture, response_id: str, reason: str
|
||||
) -> None:
|
||||
data = {"response_id": response_id, "litellm_metadata": {}}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=_router())
|
||||
|
||||
assert data == {"response_id": response_id, "litellm_metadata": {}}
|
||||
assert [message.endswith(f"({reason})") for message in _ungoverned_retrieval_warnings(caplog)] == [True]
|
||||
|
||||
|
||||
def test_without_a_router_attaches_nothing_and_warns(policy_engine: None, caplog: pytest.LogCaptureFixture) -> None:
|
||||
data = _retrieval_data(GOVERNED_MODEL_ID)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None)
|
||||
|
||||
assert data == _retrieval_data(GOVERNED_MODEL_ID)
|
||||
assert [message.endswith("(no router)") for message in _ungoverned_retrieval_warnings(caplog)] == [True]
|
||||
|
||||
|
||||
def test_without_policy_engine_attaches_nothing_quietly(caplog: pytest.LogCaptureFixture) -> None:
|
||||
get_policy_registry().clear()
|
||||
data = _retrieval_data(GOVERNED_MODEL_ID)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
attach_post_call_pipelines_to_retrieval(data=data, user_api_key_dict=UserAPIKeyAuth(), llm_router=None)
|
||||
|
||||
assert data == _retrieval_data(GOVERNED_MODEL_ID)
|
||||
assert _ungoverned_retrieval_warnings(caplog) == []
|
||||
|
|
@ -3,7 +3,7 @@ import copy
|
|||
import datetime
|
||||
import json
|
||||
from types import MappingProxyType, SimpleNamespace
|
||||
from typing import AsyncGenerator, Callable, Final, Optional
|
||||
from typing import AsyncGenerator, Callable, Final, Iterator, Optional
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import httpx
|
||||
|
|
@ -8307,3 +8307,116 @@ async def test_handle_llm_api_exception_forwards_provider_headers_on_http_status
|
|||
|
||||
assert exc_info.value.headers is not None
|
||||
assert exc_info.value.headers["llm_provider-x-amzn-requestid"] == "req-passthrough-500"
|
||||
|
||||
|
||||
class TestBackgroundResponseRetrievalGovernance:
|
||||
"""LIT-7175: retrieving a background Response attaches the model's post_call policy pipelines."""
|
||||
|
||||
GOVERNED_MODEL_GROUP = "gpt-5.4-mini"
|
||||
GOVERNED_MODEL_ID = "deployment-governed"
|
||||
|
||||
@pytest.fixture
|
||||
def policy_engine(self) -> Iterator[None]:
|
||||
from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
|
||||
get_policy_registry().load_policies(
|
||||
{
|
||||
"response-governance": {
|
||||
"guardrails": {"add": ["output-word-filter"]},
|
||||
"pipeline": {
|
||||
"mode": "post_call",
|
||||
"steps": [{"guardrail": "output-word-filter", "on_pass": "allow", "on_fail": "block"}],
|
||||
},
|
||||
}
|
||||
}
|
||||
)
|
||||
get_attachment_registry().load_attachments(
|
||||
[{"policy": "response-governance", "models": [self.GOVERNED_MODEL_GROUP]}]
|
||||
)
|
||||
yield
|
||||
get_policy_registry().clear()
|
||||
get_attachment_registry().clear()
|
||||
|
||||
def _router(self) -> MagicMock:
|
||||
from litellm.types.router import Deployment, LiteLLM_Params
|
||||
|
||||
router = MagicMock()
|
||||
router.get_deployment.side_effect = lambda model_id: (
|
||||
Deployment(
|
||||
model_name=self.GOVERNED_MODEL_GROUP,
|
||||
litellm_params=LiteLLM_Params(model=f"openai/{self.GOVERNED_MODEL_GROUP}"),
|
||||
model_info={"id": model_id},
|
||||
)
|
||||
if model_id == self.GOVERNED_MODEL_ID
|
||||
else None
|
||||
)
|
||||
return router
|
||||
|
||||
async def _pre_call(self, route_type: str, monkeypatch: pytest.MonkeyPatch) -> dict[str, object]:
|
||||
from litellm.responses.utils import ResponsesAPIRequestUtils
|
||||
|
||||
client_facing_response_id = "resp_opaque-client-facing-id"
|
||||
encoded_response_id = ResponsesAPIRequestUtils._build_responses_api_response_id(
|
||||
custom_llm_provider="openai", model_id=self.GOVERNED_MODEL_ID, response_id="resp_upstream"
|
||||
)
|
||||
processing_obj = ProxyBaseLLMRequestProcessing(
|
||||
data={"response_id": client_facing_response_id, "litellm_metadata": {}}
|
||||
)
|
||||
mock_request = MagicMock(spec=Request)
|
||||
mock_request.headers = {}
|
||||
|
||||
async def passthrough_add_litellm_data_to_request(
|
||||
data: dict[str, object], **kwargs: object
|
||||
) -> dict[str, object]:
|
||||
return data
|
||||
|
||||
async def decrypting_pre_call_hook(
|
||||
user_api_key_dict: ProxyUserAPIKeyAuth, data: dict[str, object], call_type: str
|
||||
) -> dict[str, object]:
|
||||
if data.get("response_id") == client_facing_response_id:
|
||||
data["response_id"] = encoded_response_id
|
||||
return data
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm.proxy.common_request_processing,
|
||||
"add_litellm_data_to_request",
|
||||
passthrough_add_litellm_data_to_request,
|
||||
)
|
||||
proxy_logging_obj = MagicMock(spec=ProxyLogging)
|
||||
proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=decrypting_pre_call_hook)
|
||||
proxy_config = MagicMock(spec=ProxyConfig)
|
||||
proxy_config._get_hierarchical_router_settings = AsyncMock(return_value=None)
|
||||
returned_data, _ = await processing_obj.common_processing_pre_call_logic(
|
||||
request=mock_request,
|
||||
general_settings={},
|
||||
user_api_key_dict=ProxyUserAPIKeyAuth(),
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
proxy_config=proxy_config,
|
||||
route_type=route_type,
|
||||
llm_router=self._router(),
|
||||
)
|
||||
return returned_data
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieving_a_background_response_attaches_its_model_post_call_pipeline(
|
||||
self, policy_engine: None, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
data = await self._pre_call("aget_responses", monkeypatch)
|
||||
|
||||
assert data["response_id"].startswith("resp_bGl0ZWxsbTpjdXN0b21f")
|
||||
pipelines = data["litellm_metadata"]["_guardrail_pipelines"]
|
||||
assert [(policy_name, [step.guardrail for step in pipeline.steps]) for policy_name, pipeline in pipelines] == [
|
||||
("response-governance", ["output-word-filter"])
|
||||
]
|
||||
assert data["litellm_metadata"]["applied_policies"] == ["response-governance"]
|
||||
assert data["model"] is None
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_submitting_a_response_does_not_attach_pipelines_from_its_response_id(
|
||||
self, policy_engine: None, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
data = await self._pre_call("aresponses", monkeypatch)
|
||||
|
||||
assert "_guardrail_pipelines" not in data["litellm_metadata"]
|
||||
assert "applied_policies" not in data["litellm_metadata"]
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ from __future__ import annotations
|
|||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
from collections.abc import Iterator
|
||||
from typing import Any, Callable, Dict, List
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -30,6 +31,7 @@ from litellm.proxy.common_utils.callback_utils import add_guardrail_to_applied_g
|
|||
from litellm.proxy.utils import ProxyLogging, _streamable_post_call_pipelines
|
||||
from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail
|
||||
from litellm.types.guardrails import BlockedWord, ContentFilterAction, GuardrailEventHooks
|
||||
from litellm.types.llms.openai import ResponsesAPIResponse
|
||||
from litellm.types.proxy.policy_engine.pipeline_types import (
|
||||
GuardrailPipeline,
|
||||
PipelineStep,
|
||||
|
|
@ -154,9 +156,7 @@ async def test_execute_guardrail_hook_unknown_hook_type_raises(proxy_logging, ma
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_guardrail_with_load_balancing_routes_through_router(
|
||||
proxy_logging, make_user_api_key_auth
|
||||
):
|
||||
async def test_execute_guardrail_with_load_balancing_routes_through_router(proxy_logging, make_user_api_key_auth):
|
||||
cb = _make_guardrail()
|
||||
router = MagicMock()
|
||||
router.get_available_guardrail = MagicMock(return_value={"callback": cb})
|
||||
|
|
@ -172,9 +172,7 @@ async def test_execute_guardrail_with_load_balancing_routes_through_router(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_guardrail_with_load_balancing_router_none_raises(
|
||||
proxy_logging, make_user_api_key_auth
|
||||
):
|
||||
async def test_execute_guardrail_with_load_balancing_router_none_raises(proxy_logging, make_user_api_key_auth):
|
||||
with patch("litellm.proxy.proxy_server.llm_router", None):
|
||||
with pytest.raises(ValueError, match="Router not initialized"):
|
||||
await proxy_logging._execute_guardrail_with_load_balancing(
|
||||
|
|
@ -187,9 +185,7 @@ async def test_execute_guardrail_with_load_balancing_router_none_raises(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_guardrail_with_load_balancing_no_callback_raises(
|
||||
proxy_logging, make_user_api_key_auth
|
||||
):
|
||||
async def test_execute_guardrail_with_load_balancing_no_callback_raises(proxy_logging, make_user_api_key_auth):
|
||||
router = MagicMock()
|
||||
router.get_available_guardrail = MagicMock(return_value={"callback": None})
|
||||
with patch("litellm.proxy.proxy_server.llm_router", router):
|
||||
|
|
@ -209,9 +205,7 @@ async def test_execute_guardrail_with_load_balancing_no_callback_raises(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_guardrail_callback_skipped_when_should_run_false(
|
||||
proxy_logging, make_user_api_key_auth
|
||||
):
|
||||
async def test_process_guardrail_callback_skipped_when_should_run_false(proxy_logging, make_user_api_key_auth):
|
||||
cb = _make_guardrail()
|
||||
cb.should_run_guardrail = MagicMock(return_value=False)
|
||||
out = await proxy_logging._process_guardrail_callback(
|
||||
|
|
@ -225,9 +219,7 @@ async def test_process_guardrail_callback_skipped_when_should_run_false(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_process_guardrail_callback_returns_data_on_success(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
async def test_process_guardrail_callback_returns_data_on_success(proxy_logging, make_user_api_key_auth, monkeypatch):
|
||||
cb = _make_guardrail()
|
||||
cb.should_run_guardrail = MagicMock(return_value=True)
|
||||
proxy_logging._should_use_guardrail_load_balancing = MagicMock(return_value=False)
|
||||
|
|
@ -342,14 +334,14 @@ async def test_maybe_execute_pipelines_no_pipelines_returns_data(proxy_logging,
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode(proxy_logging, make_user_api_key_auth, monkeypatch):
|
||||
async def test_maybe_execute_pipelines_skips_pipelines_with_other_mode(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
pipeline = MagicMock()
|
||||
pipeline.mode = "post_call" # not pre_call
|
||||
data = {"metadata": {"_guardrail_pipelines": [("p1", pipeline)]}, "model": "m", "messages": []}
|
||||
executed = MagicMock()
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps", executed
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.policy_engine.pipeline_executor.PipelineExecutor.execute_steps", executed)
|
||||
out, replacement = await proxy_logging._maybe_execute_pipelines(
|
||||
data=data,
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
|
|
@ -537,9 +529,7 @@ def test_handle_pipeline_result_block_enriches_with_guardrail_name_and_mode():
|
|||
litellm.callbacks = [cb]
|
||||
try:
|
||||
with pytest.raises(HTTPException) as info:
|
||||
ProxyLogging._handle_pipeline_result(
|
||||
result=result, data={"model": "m"}, policy_name="p"
|
||||
)
|
||||
ProxyLogging._handle_pipeline_result(result=result, data={"model": "m"}, policy_name="p")
|
||||
finally:
|
||||
litellm.callbacks = saved
|
||||
|
||||
|
|
@ -651,9 +641,7 @@ async def test_run_guardrail_with_metrics_records_error_and_enriches(monkeypatch
|
|||
monkeypatch.setattr(litellm, "callbacks", [prom])
|
||||
|
||||
with pytest.raises(HTTPException):
|
||||
await ProxyLogging._run_guardrail_with_metrics(
|
||||
callback=cb, coro=task(), hook_type="post_call"
|
||||
)
|
||||
await ProxyLogging._run_guardrail_with_metrics(callback=cb, coro=task(), hook_type="post_call")
|
||||
|
||||
assert detail["guardrail_name"] == "presidio"
|
||||
recorded = prom._record_guardrail_metrics.call_args.kwargs
|
||||
|
|
@ -681,9 +669,7 @@ def _moderation_guardrail() -> MagicMock:
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_during_call_hook_records_latency_metric(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
async def test_during_call_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch):
|
||||
cb = _moderation_guardrail()
|
||||
prom = _prometheus_callback()
|
||||
monkeypatch.setattr(litellm, "callbacks", [prom, cb])
|
||||
|
|
@ -702,9 +688,7 @@ async def test_during_call_hook_records_latency_metric(
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_success_hook_records_latency_metric(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch
|
||||
):
|
||||
async def test_post_call_success_hook_records_latency_metric(proxy_logging, make_user_api_key_auth, monkeypatch):
|
||||
cb = _moderation_guardrail()
|
||||
prom = _prometheus_callback()
|
||||
monkeypatch.setattr(litellm, "callbacks", [prom, cb])
|
||||
|
|
@ -732,9 +716,7 @@ async def test_post_call_success_hook_records_latency_metric(
|
|||
async def test_process_prompt_template_no_op_when_no_prompt_spec(proxy_logging, monkeypatch):
|
||||
from litellm.proxy.prompts import prompt_registry
|
||||
|
||||
monkeypatch.setattr(
|
||||
prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: None
|
||||
)
|
||||
monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: None)
|
||||
data: Dict[str, Any] = {"messages": [{"role": "user"}], "model": "m", "temperature": 0.1}
|
||||
await proxy_logging._process_prompt_template(
|
||||
data=data,
|
||||
|
|
@ -759,9 +741,7 @@ async def test_process_prompt_template_applies_when_spec_resolves(proxy_logging,
|
|||
"get_prompt_callback_for_prompt",
|
||||
lambda *a, **kw: custom_logger,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec
|
||||
)
|
||||
monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec)
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.async_get_chat_completion_prompt = AsyncMock(
|
||||
|
|
@ -809,9 +789,7 @@ async def test_process_prompt_template_async_get_prompt_error_raises(proxy_loggi
|
|||
"get_prompt_callback_for_prompt",
|
||||
lambda *a, **kw: custom_logger,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec
|
||||
)
|
||||
monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec)
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.async_get_chat_completion_prompt = AsyncMock(side_effect=RuntimeError("bad prompt"))
|
||||
with pytest.raises(RuntimeError):
|
||||
|
|
@ -912,9 +890,7 @@ async def test_process_prompt_template_aresponses_swaps_model_and_merges_input(p
|
|||
"get_prompt_callback_for_prompt",
|
||||
lambda *a, **kw: custom_logger,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec
|
||||
)
|
||||
monkeypatch.setattr(prompt_registry.IN_MEMORY_PROMPT_REGISTRY, "resolve_prompt_spec", lambda *a, **kw: prompt_spec)
|
||||
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.async_get_chat_completion_prompt = AsyncMock(
|
||||
|
|
@ -965,6 +941,7 @@ def _post_call_pipeline_data(
|
|||
"metadata": {
|
||||
"_guardrail_pipelines": [("response-governance", pipeline)],
|
||||
"_pipeline_managed_guardrails": {guardrail},
|
||||
"policy_sources": {"response-governance": "model:m"},
|
||||
},
|
||||
**extra,
|
||||
}
|
||||
|
|
@ -1123,9 +1100,7 @@ async def test_pre_call_hook_still_runs_guardrail_managed_only_by_post_call_pipe
|
|||
},
|
||||
}
|
||||
|
||||
await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion"
|
||||
)
|
||||
await proxy_logging.pre_call_hook(user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion")
|
||||
|
||||
assert seen["count"] == 1
|
||||
|
||||
|
|
@ -1288,11 +1263,7 @@ async def test_post_call_pipeline_block_keeps_guardrail_metadata_writes(
|
|||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"callbacks",
|
||||
[
|
||||
BlockingWriterGuardrail(
|
||||
guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False
|
||||
)
|
||||
],
|
||||
[BlockingWriterGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)],
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data()
|
||||
|
|
@ -1378,9 +1349,7 @@ async def test_pre_call_pipeline_managed_parallel_guardrail_runs_exactly_once(
|
|||
},
|
||||
}
|
||||
|
||||
await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion"
|
||||
)
|
||||
await proxy_logging.pre_call_hook(user_api_key_dict=make_user_api_key_auth(), data=data, call_type="completion")
|
||||
|
||||
assert seen["count"] == 1
|
||||
|
||||
|
|
@ -1419,10 +1388,295 @@ async def test_streaming_request_whose_pipeline_guardrail_is_missing_streams_ver
|
|||
assert any("response-governance" in message and "gr-post" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
def _background_response(status: str, text: str = "") -> ResponsesAPIResponse:
|
||||
output = (
|
||||
[{"type": "message", "role": "assistant", "content": [{"type": "output_text", "text": text}]}] if text else []
|
||||
)
|
||||
return ResponsesAPIResponse(id="resp_bg", created_at=0, output=output, status=status)
|
||||
|
||||
|
||||
def _output_blocking_callbacks(seen: dict[str, object]) -> list[CustomGuardrail]:
|
||||
class OutputBlockingGuardrail(CustomGuardrail):
|
||||
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
||||
seen["response"] = response
|
||||
raise HTTPException(status_code=400, detail={"error": "output blocked"})
|
||||
|
||||
return [
|
||||
OutputBlockingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_accepts_background_request_with_post_call_pipeline(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
|
||||
@pytest.mark.parametrize("pending_status", ["queued", "in_progress"])
|
||||
async def test_post_call_success_hook_waits_for_pending_background_response_before_running_pipeline(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
pending_status: str,
|
||||
) -> None:
|
||||
seen: dict[str, object] = {}
|
||||
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data(background=True)
|
||||
response = _background_response(pending_status)
|
||||
|
||||
with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"):
|
||||
out = await proxy_logging.post_call_success_hook(
|
||||
data=data, response=response, user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert out is response
|
||||
assert "response" not in seen
|
||||
assert not _warnings(caplog)
|
||||
assert any(
|
||||
"response-governance" in record.getMessage() and pending_status in record.getMessage()
|
||||
for record in caplog.records
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("final_status", ["completed", "incomplete"])
|
||||
async def test_post_call_success_hook_runs_pipeline_on_retrieved_background_response(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
final_status: str,
|
||||
) -> None:
|
||||
seen: dict[str, object] = {}
|
||||
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks(seen))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data()
|
||||
response = _background_response(final_status, text="kumquat")
|
||||
|
||||
with pytest.raises(HTTPException) as info:
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=response, user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert info.value.detail["error"] == "output blocked"
|
||||
assert seen["response"] is response
|
||||
|
||||
|
||||
def _output_passing_callbacks() -> list[CustomGuardrail]:
|
||||
class OutputPassingGuardrail(CustomGuardrail):
|
||||
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
||||
return response
|
||||
|
||||
return [
|
||||
OutputPassingGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)
|
||||
]
|
||||
|
||||
|
||||
def _claimed_post_call_pipeline_data(
|
||||
*policy_names: str, extra_guardrails: dict[str, list[str]] | None = None, policy_source: str | None = "model:m"
|
||||
):
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
|
||||
step = {"guardrail": "gr-post", "on_pass": "allow", "on_fail": "block"}
|
||||
get_policy_registry().load_policies(
|
||||
{
|
||||
policy_name: {
|
||||
"guardrails": {"add": ["gr-post", *(extra_guardrails or {}).get(policy_name, [])]},
|
||||
"pipeline": {"mode": "post_call", "steps": [step]},
|
||||
}
|
||||
for policy_name in policy_names
|
||||
}
|
||||
)
|
||||
pipeline = GuardrailPipeline(mode="post_call", steps=[PipelineStep(**step)])
|
||||
return {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"metadata": {
|
||||
"_guardrail_pipelines": [(policy_name, pipeline) for policy_name in policy_names],
|
||||
"_pipeline_managed_guardrails": {"gr-post"},
|
||||
"applied_policies": list(policy_names),
|
||||
"applied_guardrails": ["gr-post", *(g for gs in (extra_guardrails or {}).values() for g in gs)],
|
||||
"policy_sources": {policy_name: policy_source for policy_name in policy_names if policy_source is not None},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def clear_policy_registry() -> Iterator[None]:
|
||||
from litellm.proxy.policy_engine.policy_registry import get_policy_registry
|
||||
|
||||
yield
|
||||
get_policy_registry().clear()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_background_response_withdraws_the_deferred_policy_claims(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
clear_policy_registry: None,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({}))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _claimed_post_call_pipeline_data("response-governance")
|
||||
|
||||
out = await proxy_logging.post_call_success_hook(
|
||||
data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert out.status == "queued"
|
||||
assert "applied_policies" not in data["metadata"]
|
||||
assert "policy_sources" not in data["metadata"]
|
||||
assert "applied_guardrails" not in data["metadata"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_background_response_warns_when_the_deferred_policy_was_matched_through_a_tag(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
clear_policy_registry: None,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({}))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _claimed_post_call_pipeline_data("response-governance", policy_source="tag:governed+model:m")
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
out = await proxy_logging.post_call_success_hook(
|
||||
data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert out.status == "queued"
|
||||
assert "policy_sources" not in data["metadata"]
|
||||
assert [message for message in _warnings(caplog) if "through a request tag" in message] == [
|
||||
"Policy engine: background response resp_bg matched post_call policies through a request tag at submit; "
|
||||
"retrieval re-matches only the key, team, and model scopes, so a tag carried in the request body "
|
||||
"does not govern the completed response: response-governance"
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_background_response_warns_when_the_deferred_policy_came_from_the_request_body(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
clear_policy_registry: None,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({}))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _claimed_post_call_pipeline_data("body-governance", policy_source=None)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
out = await proxy_logging.post_call_success_hook(
|
||||
data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert out.status == "queued"
|
||||
assert "policy_sources" not in data["metadata"]
|
||||
assert _warnings(caplog) == [
|
||||
"Policy engine: background response resp_bg matched post_call policies through the request body's policies "
|
||||
"list at submit; retrieval carries no request body, so those policies do not govern the completed "
|
||||
"response: body-governance"
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_background_response_matched_through_its_model_does_not_warn(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
clear_policy_registry: None,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({}))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=_claimed_post_call_pipeline_data("response-governance"),
|
||||
response=_background_response("queued"),
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
)
|
||||
|
||||
assert _warnings(caplog) == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_background_response_keeps_the_claim_of_a_policy_that_runs_outside_its_pipeline(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
clear_policy_registry: None,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", _output_blocking_callbacks({}))
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _claimed_post_call_pipeline_data(
|
||||
"input-and-output-governance",
|
||||
"response-governance",
|
||||
extra_guardrails={"input-and-output-governance": ["gr-pre"]},
|
||||
)
|
||||
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=_background_response("in_progress"), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert data["metadata"]["applied_policies"] == ["input-and-output-governance"]
|
||||
assert data["metadata"]["applied_guardrails"] == ["gr-pre"]
|
||||
assert data["metadata"]["policy_sources"] == {"input-and-output-governance": "model:m"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pending_background_response_keeps_the_claim_of_a_default_on_guardrail_that_ran_pre_call(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
clear_policy_registry: None,
|
||||
) -> None:
|
||||
class DualStageGuardrail(CustomGuardrail):
|
||||
async def async_post_call_success_hook(self, data, user_api_key_dict, response):
|
||||
return response
|
||||
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"callbacks",
|
||||
[DualStageGuardrail(guardrail_name="gr-post", event_hook=["pre_call", "post_call"], default_on=True)],
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _claimed_post_call_pipeline_data("response-governance")
|
||||
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=_background_response("queued"), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert "applied_policies" not in data["metadata"]
|
||||
assert data["metadata"]["applied_guardrails"] == ["gr-post"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_retrieved_background_response_keeps_the_policy_claim_once_its_pipeline_ran(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
clear_policy_registry: None,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", _output_passing_callbacks())
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _claimed_post_call_pipeline_data("response-governance")
|
||||
|
||||
await proxy_logging.post_call_success_hook(
|
||||
data=data, response=_background_response("completed", text="fine"), user_api_key_dict=make_user_api_key_auth()
|
||||
)
|
||||
|
||||
assert data["metadata"]["applied_policies"] == ["response-governance"]
|
||||
assert data["metadata"]["policy_sources"] == {"response-governance": "model:m"}
|
||||
assert data["metadata"]["applied_guardrails"] == ["gr-post"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_stays_quiet_on_background_request_with_post_call_pipeline(
|
||||
proxy_logging: ProxyLogging,
|
||||
make_user_api_key_auth: Callable[..., UserAPIKeyAuth],
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
caplog: pytest.LogCaptureFixture,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "callbacks", [])
|
||||
data = _post_call_pipeline_data(background=True)
|
||||
|
||||
|
|
@ -1436,36 +1690,7 @@ async def test_pre_call_hook_accepts_background_request_with_post_call_pipeline(
|
|||
|
||||
assert out is not None
|
||||
assert out.get("background") is True
|
||||
assert any("response-governance" in message and "background" in message for message in _warnings(caplog))
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_pre_call_hook_stays_quiet_on_background_request_without_post_call_pipeline(
|
||||
proxy_logging, make_user_api_key_auth, monkeypatch, caplog
|
||||
):
|
||||
seen: Dict[str, Any] = {}
|
||||
monkeypatch.setattr(litellm, "callbacks", [_unified_stream_guardrail(seen)])
|
||||
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-post", on_fail="block")])
|
||||
data = {
|
||||
"model": "m",
|
||||
"messages": [{"role": "user", "content": "hi"}],
|
||||
"background": True,
|
||||
"metadata": {
|
||||
"_guardrail_pipelines": [("request-governance", pre_call)],
|
||||
"_pipeline_managed_guardrails": {"gr-post"},
|
||||
},
|
||||
}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
out = await proxy_logging.pre_call_hook(
|
||||
user_api_key_dict=make_user_api_key_auth(),
|
||||
data=data,
|
||||
call_type="aresponses",
|
||||
guardrails_only=True,
|
||||
)
|
||||
|
||||
assert out is not None
|
||||
assert not any("background" in message for message in _warnings(caplog))
|
||||
assert not _warnings(caplog)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
@ -1512,7 +1737,9 @@ def test_streamable_post_call_pipelines_keeps_supported_and_drops_unsupported(
|
|||
steps=[PipelineStep(guardrail="gr-post", on_fail="next"), PipelineStep(guardrail="gr-native", on_fail="block")],
|
||||
)
|
||||
pre_call = GuardrailPipeline(mode="pre_call", steps=[PipelineStep(guardrail="gr-native", on_fail="block")])
|
||||
data = {"metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]}}
|
||||
data = {
|
||||
"metadata": {"_guardrail_pipelines": [("governed", governed), ("ungoverned", ungoverned), ("req", pre_call)]}
|
||||
}
|
||||
|
||||
with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
|
||||
streamable = _streamable_post_call_pipelines(data, make_user_api_key_auth(request_route="/v1/chat/completions"))
|
||||
|
|
@ -1812,7 +2039,9 @@ def _rewriting_stream_guardrail(transform: Callable[[Dict[str, Any]], Dict[str,
|
|||
async def apply_guardrail(self, inputs, request_data, input_type, logging_obj=None):
|
||||
return {**inputs, **transform(inputs)}
|
||||
|
||||
return RewritingStreamGuardrail(guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False)
|
||||
return RewritingStreamGuardrail(
|
||||
guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=False
|
||||
)
|
||||
|
||||
|
||||
def _tool_call_stream_chunks() -> List[Any]:
|
||||
|
|
@ -1991,11 +2220,38 @@ async def test_streaming_iterator_hook_pipeline_releases_originals_on_unresolvab
|
|||
|
||||
def _anthropic_sse_chunks() -> List[bytes]:
|
||||
events = [
|
||||
("message_start", {"type": "message_start", "message": {"id": "msg_1", "type": "message", "role": "assistant", "model": "m", "content": [], "stop_reason": None, "usage": {"input_tokens": 1, "output_tokens": 0}}}),
|
||||
("content_block_start", {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}),
|
||||
("content_block_delta", {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello world"}}),
|
||||
(
|
||||
"message_start",
|
||||
{
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_1",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "m",
|
||||
"content": [],
|
||||
"stop_reason": None,
|
||||
"usage": {"input_tokens": 1, "output_tokens": 0},
|
||||
},
|
||||
},
|
||||
),
|
||||
(
|
||||
"content_block_start",
|
||||
{"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}},
|
||||
),
|
||||
(
|
||||
"content_block_delta",
|
||||
{"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "hello world"}},
|
||||
),
|
||||
("content_block_stop", {"type": "content_block_stop", "index": 0}),
|
||||
("message_delta", {"type": "message_delta", "delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 2}}),
|
||||
(
|
||||
"message_delta",
|
||||
{
|
||||
"type": "message_delta",
|
||||
"delta": {"stop_reason": "end_turn", "stop_sequence": None},
|
||||
"usage": {"output_tokens": 2},
|
||||
},
|
||||
),
|
||||
("message_stop", {"type": "message_stop"}),
|
||||
]
|
||||
return [f"event: {name}\ndata: {json.dumps(payload)}\n\n".encode() for name, payload in events]
|
||||
|
|
@ -2062,7 +2318,13 @@ async def test_streaming_iterator_hook_pipeline_delivers_text_rewrite_on_anthrop
|
|||
assert "hello [MASKED]" in raw
|
||||
assert "hello world" not in raw
|
||||
assert raw.count("event: content_block_delta") == 1
|
||||
for expected_event in ("message_start", "content_block_start", "content_block_stop", "message_delta", "message_stop"):
|
||||
for expected_event in (
|
||||
"message_start",
|
||||
"content_block_start",
|
||||
"content_block_stop",
|
||||
"message_delta",
|
||||
"message_stop",
|
||||
):
|
||||
assert f"event: {expected_event}" in raw
|
||||
|
||||
|
||||
|
|
@ -2135,9 +2397,7 @@ async def test_per_chunk_streaming_hook_skips_pipeline_managed_guardrail(
|
|||
managed = UnifiedRecordingGuardrail(
|
||||
guardrail_name="gr-post", event_hook=GuardrailEventHooks.post_call, default_on=True
|
||||
)
|
||||
free = RecordingGuardrail(
|
||||
guardrail_name="gr-free", event_hook=GuardrailEventHooks.post_call, default_on=True
|
||||
)
|
||||
free = RecordingGuardrail(guardrail_name="gr-free", event_hook=GuardrailEventHooks.post_call, default_on=True)
|
||||
monkeypatch.setattr(litellm, "callbacks", [managed, free])
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None, raising=False)
|
||||
data = _post_call_pipeline_data(stream=True)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue