diff --git a/litellm/proxy/policy_engine/policy_endpoints.py b/litellm/proxy/policy_engine/policy_endpoints.py index 3a25e249a4e..35e7f09e806 100644 --- a/litellm/proxy/policy_engine/policy_endpoints.py +++ b/litellm/proxy/policy_engine/policy_endpoints.py @@ -4,7 +4,7 @@ CRUD ENDPOINTS FOR POLICIES Provides REST API endpoints for managing policies and policy attachments. """ -from typing import Optional +from typing import Any, Dict, Optional from fastapi import APIRouter, Depends, HTTPException @@ -16,6 +16,7 @@ from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor from litellm.proxy.policy_engine.policy_registry import get_policy_registry from litellm.types.proxy.policy_engine import ( GuardrailPipeline, + PipelineExecutionResult, PipelineTestRequest, PolicyAttachmentCreateRequest, PolicyAttachmentDBResponse, @@ -32,6 +33,43 @@ from litellm.types.proxy.policy_engine import ( router = APIRouter() +_PIPELINE_RESPONSE_INTERNAL_KEYS = { + "litellm_parent_otel_span", + "parent_otel_span", + "user_api_key_auth", +} + + +def _sanitize_pipeline_response_value(value: Any) -> Any: + if isinstance(value, dict): + return { + str(key): _sanitize_pipeline_response_value(item) + for key, item in value.items() + if str(key) not in _PIPELINE_RESPONSE_INTERNAL_KEYS + } + + if isinstance(value, (list, tuple, set)): + return [_sanitize_pipeline_response_value(item) for item in value] + + if value is None or isinstance(value, (str, int, float, bool)): + return value + + if hasattr(value, "model_dump"): + try: + return _sanitize_pipeline_response_value( + value.model_dump(mode="json", fallback=str) + ) + except Exception: + pass + + return str(value) + + +def _dump_pipeline_result_for_response( + result: PipelineExecutionResult, +) -> Dict[str, Any]: + return _sanitize_pipeline_response_value(result.model_dump()) + # ───────────────────────────────────────────────────────────────────────────── # Policy CRUD Endpoints @@ -601,7 +639,7 @@ async def test_pipeline( call_type="completion", policy_name="test-pipeline", ) - return result.model_dump() + return _dump_pipeline_result_for_response(result) except Exception as e: verbose_proxy_logger.exception(f"Error testing pipeline: {e}") raise HTTPException(status_code=500, detail=str(e)) diff --git a/tests/test_litellm/proxy/policy_engine/test_policy_endpoints.py b/tests/test_litellm/proxy/policy_engine/test_policy_endpoints.py new file mode 100644 index 00000000000..ee4051e866f --- /dev/null +++ b/tests/test_litellm/proxy/policy_engine/test_policy_endpoints.py @@ -0,0 +1,60 @@ +import json + +import pytest +from fastapi.encoders import jsonable_encoder + +from litellm.proxy.policy_engine.policy_endpoints import ( + _dump_pipeline_result_for_response, +) +from litellm.types.proxy.policy_engine import ( + PipelineExecutionResult, + PipelineStepResult, +) + + +class OpaqueSpan: + __slots__ = () + + def __str__(self) -> str: + return "opaque-span" + + +def test_pipeline_test_response_removes_internal_non_json_values(): + opaque_span = OpaqueSpan() + result = PipelineExecutionResult( + terminal_action="allow", + step_results=[ + PipelineStepResult( + guardrail_name="pii", + outcome="pass", + action_taken="allow", + modified_data={ + "metadata": { + "parent_otel_span": opaque_span, + "litellm_parent_otel_span": opaque_span, + "safe": "kept", + }, + "non_internal_object": opaque_span, + }, + ) + ], + modified_data={ + "metadata": { + "parent_otel_span": opaque_span, + "safe": "kept", + }, + }, + ) + + with pytest.raises(ValueError): + jsonable_encoder(result.model_dump()) + + dumped = _dump_pipeline_result_for_response(result) + + json.dumps(dumped) + assert dumped["modified_data"]["metadata"] == {"safe": "kept"} + assert dumped["step_results"][0]["modified_data"]["metadata"] == {"safe": "kept"} + assert ( + dumped["step_results"][0]["modified_data"]["non_internal_object"] + == "opaque-span" + )