mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-07 02:59:05 +00:00
Integrate pipeline execution into proxy pre_call_hook
This commit is contained in:
parent
ba392754a2
commit
c94074558e
1 changed files with 114 additions and 1 deletions
|
|
@ -77,7 +77,10 @@ from litellm._logging import verbose_proxy_logger
|
|||
from litellm._service_logger import ServiceLogging, ServiceTypes
|
||||
from litellm.caching.caching import DualCache, RedisCache
|
||||
from litellm.exceptions import RejectedRequestError
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_guardrail import (
|
||||
CustomGuardrail,
|
||||
ModifyResponseException,
|
||||
)
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
|
||||
from litellm.integrations.SlackAlerting.utils import _add_langfuse_trace_id_to_alert
|
||||
|
|
@ -110,6 +113,7 @@ from litellm.proxy.hooks.parallel_request_limiter import (
|
|||
_PROXY_MaxParallelRequestsHandler,
|
||||
)
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor
|
||||
from litellm.secret_managers.main import str_to_bool
|
||||
from litellm.types.integrations.slack_alerting import DEFAULT_ALERT_TYPES
|
||||
from litellm.types.mcp import (
|
||||
|
|
@ -117,6 +121,7 @@ from litellm.types.mcp import (
|
|||
MCPPreCallRequestObject,
|
||||
MCPPreCallResponseObject,
|
||||
)
|
||||
from litellm.types.proxy.policy_engine.pipeline_types import PipelineExecutionResult
|
||||
from litellm.types.utils import LLMResponseTypes, LoggedLiteLLMParams
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -1141,6 +1146,98 @@ class ProxyLogging:
|
|||
request_data=data, guardrail_name=guardrail_name
|
||||
)
|
||||
|
||||
async def _maybe_execute_pipelines(
|
||||
self,
|
||||
data: dict,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
call_type: str,
|
||||
event_hook: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Execute guardrail pipelines if any are configured for this request.
|
||||
|
||||
Checks metadata for pipelines resolved by the policy engine
|
||||
and executes them. Handles the result (allow/block/modify_response).
|
||||
|
||||
Returns the (possibly modified) data dict.
|
||||
"""
|
||||
metadata = data.get("metadata", data.get("litellm_metadata", {})) or {}
|
||||
pipelines = metadata.get("_guardrail_pipelines")
|
||||
if not pipelines:
|
||||
return data
|
||||
|
||||
for policy_name, pipeline in pipelines:
|
||||
if pipeline.mode != event_hook:
|
||||
continue
|
||||
|
||||
result: PipelineExecutionResult = await PipelineExecutor.execute_steps(
|
||||
steps=pipeline.steps,
|
||||
mode=pipeline.mode,
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_type=call_type,
|
||||
policy_name=policy_name,
|
||||
)
|
||||
|
||||
data = self._handle_pipeline_result(
|
||||
result=result,
|
||||
data=data,
|
||||
policy_name=policy_name,
|
||||
)
|
||||
|
||||
return data
|
||||
|
||||
@staticmethod
|
||||
def _handle_pipeline_result(
|
||||
result: Any,
|
||||
data: dict,
|
||||
policy_name: str,
|
||||
) -> dict:
|
||||
"""
|
||||
Handle a PipelineExecutionResult — allow, block, or modify_response.
|
||||
|
||||
Returns data dict if allowed, raises on block/modify_response.
|
||||
"""
|
||||
if result.terminal_action == "allow":
|
||||
if result.modified_data is not None:
|
||||
data.update(result.modified_data)
|
||||
return data
|
||||
|
||||
if result.terminal_action == "block":
|
||||
step_results_serializable = [
|
||||
{
|
||||
"guardrail": sr.guardrail_name,
|
||||
"outcome": sr.outcome,
|
||||
"action": sr.action_taken,
|
||||
}
|
||||
for sr in result.step_results
|
||||
]
|
||||
error_detail = {
|
||||
"error": {
|
||||
"message": f"Content blocked by guardrail pipeline '{policy_name}'",
|
||||
"type": "guardrail_pipeline_error",
|
||||
"pipeline_context": {
|
||||
"policy": policy_name,
|
||||
"step_results": step_results_serializable,
|
||||
},
|
||||
}
|
||||
}
|
||||
if HTTPException is not None:
|
||||
raise HTTPException(status_code=400, detail=error_detail)
|
||||
else:
|
||||
raise Exception(str(error_detail))
|
||||
|
||||
if result.terminal_action == "modify_response":
|
||||
raise ModifyResponseException(
|
||||
message=result.modify_response_message or "Response modified by pipeline",
|
||||
model=data.get("model", "unknown"),
|
||||
request_data=data,
|
||||
guardrail_name=f"pipeline:{policy_name}",
|
||||
detection_info=None,
|
||||
)
|
||||
|
||||
return data
|
||||
|
||||
# The actual implementation of the function
|
||||
@overload
|
||||
async def pre_call_hook(
|
||||
|
|
@ -1203,6 +1300,18 @@ class ProxyLogging:
|
|||
)
|
||||
|
||||
try:
|
||||
# Execute guardrail pipelines before the normal callback loop
|
||||
data = await self._maybe_execute_pipelines(
|
||||
data=data,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
call_type=call_type,
|
||||
event_hook="pre_call",
|
||||
)
|
||||
|
||||
# Get pipeline-managed guardrails to skip in normal loop
|
||||
metadata = data.get("metadata", data.get("litellm_metadata", {})) or {}
|
||||
pipeline_managed: set = metadata.get("_pipeline_managed_guardrails", set())
|
||||
|
||||
for callback in litellm.callbacks:
|
||||
start_time = time.time()
|
||||
_callback = None
|
||||
|
|
@ -1217,6 +1326,10 @@ class ProxyLogging:
|
|||
and isinstance(_callback, CustomGuardrail)
|
||||
and data is not None
|
||||
):
|
||||
# Skip guardrails managed by a pipeline
|
||||
if _callback.guardrail_name and _callback.guardrail_name in pipeline_managed:
|
||||
continue
|
||||
|
||||
result = await self._process_guardrail_callback(
|
||||
callback=_callback,
|
||||
data=data, # type: ignore
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue