From 022846baae7702a176594183467d8567627bf8b7 Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Thu, 12 Feb 2026 22:05:55 -0600 Subject: [PATCH 1/9] fix(router): remove repeated provider parsing in budget limiter hot path (#21043) * fix(router): remove budget limiter provider hot-path overhead - avoid LiteLLM_Params instantiation from dict deployments in provider lookup\n- resolve provider once per deployment and reuse in budget filtering\n- add router unit tests for hot-path behavior\n\nFixes #21042 * fix(router): handle None provider cache entries in budget limiter - avoid recomputing provider when cached value is None\n- clarify deployment_provider_map uses id(deployment) keys\n- add regression test covering None-provider cache path * refactor(router): avoid id()-based provider cache coupling - switch provider cache handoff to index-aligned list between budget-limiter loops\n- remove implicit dependency on object identity stability\n- move hot-path tests to tests/test_litellm/router_strategy per template guidance * chore(router): make use_litellm_proxy default explicit Use deployment_litellm_params.get('use_litellm_proxy', False) for clarity and parity with LiteLLM_Params default behavior. * test(router): add provider-resolution parity guard - wrap dict litellm_params with lightweight attribute view for get_llm_provider\n- reduce drift risk from manual field extraction vs LiteLLM_Params defaults\n- add parity test matrix comparing optimized path to legacy LiteLLM_Params behavior for dict deployments * test(router): harden dict view compatibility for provider resolution - extend _LiteLLMParamsDictView with mapping-like and dump methods\n- add regression test simulating future get_llm_provider method-based access\n- keep hot-path optimization while reducing duck-typing break risk --------- Co-authored-by: Codex --- litellm/router_strategy/budget_limiter.py | 97 +++++++- .../test_budget_limiter_hotpath.py | 232 ++++++++++++++++++ 2 files changed, 319 insertions(+), 10 deletions(-) create mode 100644 tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py diff --git a/litellm/router_strategy/budget_limiter.py b/litellm/router_strategy/budget_limiter.py index 9e4001b67b9..64dc5fe4741 100644 --- a/litellm/router_strategy/budget_limiter.py +++ b/litellm/router_strategy/budget_limiter.py @@ -41,6 +41,53 @@ from litellm.types.utils import GenericBudgetConfigType, StandardLoggingPayload DEFAULT_REDIS_SYNC_INTERVAL = 1 +class _LiteLLMParamsDictView: + """ + Lightweight attribute view over `litellm_params` dict. + + This avoids pydantic construction in request hot-path while preserving + attribute-style access used by `litellm.get_llm_provider(...)`. + """ + + __slots__ = ("_params",) + + def __init__(self, params: Dict[str, Any]): + self._params = params + + def __getattr__(self, key: str) -> Any: + return self._params.get(key) + + def __getitem__(self, key: str) -> Any: + return self._params.get(key) + + def __contains__(self, key: str) -> bool: + return key in self._params + + def get(self, key: str, default: Any = None) -> Any: + return self._params.get(key, default) + + def keys(self): + return self._params.keys() + + def values(self): + return self._params.values() + + def items(self): + return self._params.items() + + def __iter__(self): + return iter(self._params) + + def __len__(self) -> int: + return len(self._params) + + def dict(self) -> Dict[str, Any]: + return dict(self._params) + + def model_dump(self) -> Dict[str, Any]: + return dict(self._params) + + class RouterBudgetLimiting(CustomLogger): def __init__( self, @@ -98,6 +145,7 @@ class RouterBudgetLimiting(CustomLogger): cache_keys, provider_configs, deployment_configs, + deployment_providers, ) = await self._async_get_cache_keys_for_router_budget_limiting( healthy_deployments=healthy_deployments, request_kwargs=request_kwargs, @@ -123,6 +171,7 @@ class RouterBudgetLimiting(CustomLogger): healthy_deployments=healthy_deployments, provider_configs=provider_configs, deployment_configs=deployment_configs, + deployment_providers=deployment_providers, spend_map=spend_map, potential_deployments=potential_deployments, request_tags=_get_tags_from_request_kwargs( @@ -145,6 +194,7 @@ class RouterBudgetLimiting(CustomLogger): healthy_deployments: List[Dict[str, Any]], provider_configs: Dict[str, GenericBudgetInfo], deployment_configs: Dict[str, GenericBudgetInfo], + deployment_providers: List[Optional[str]], spend_map: Dict[str, float], request_tags: List[str], ) -> Tuple[List[Dict[str, Any]], str]: @@ -161,12 +211,15 @@ class RouterBudgetLimiting(CustomLogger): """ # Filter deployments based on both provider and deployment budgets deployment_above_budget_info: str = "" - for deployment in healthy_deployments: + for idx, deployment in enumerate(healthy_deployments): is_within_budget = True # Check provider budget if self.provider_budget_config: - provider = self._get_llm_provider_for_deployment(deployment) + if idx < len(deployment_providers): + provider = deployment_providers[idx] + else: + provider = self._get_llm_provider_for_deployment(deployment) if provider in provider_configs: config = provider_configs[provider] if config.max_budget is None: @@ -230,24 +283,32 @@ class RouterBudgetLimiting(CustomLogger): self, healthy_deployments: List[Dict[str, Any]], request_kwargs: Optional[Dict] = None, - ) -> Tuple[List[str], Dict[str, GenericBudgetInfo], Dict[str, GenericBudgetInfo]]: + ) -> Tuple[ + List[str], + Dict[str, GenericBudgetInfo], + Dict[str, GenericBudgetInfo], + List[Optional[str]], + ]: """ Returns list of cache keys to fetch from router cache for budget limiting and provider and deployment configs Returns: - Tuple[List[str], Dict[str, GenericBudgetInfo], Dict[str, GenericBudgetInfo]]: + Tuple[List[str], Dict[str, GenericBudgetInfo], Dict[str, GenericBudgetInfo], List[Optional[str]]]: - List of cache keys to fetch from router cache for budget limiting - Dict of provider budget configs `provider_configs` - Dict of deployment budget configs `deployment_configs` + - List of resolved providers aligned by deployment index `deployment_providers` """ cache_keys: List[str] = [] provider_configs: Dict[str, GenericBudgetInfo] = {} deployment_configs: Dict[str, GenericBudgetInfo] = {} + deployment_providers: List[Optional[str]] = [] for deployment in healthy_deployments: # Check provider budgets if self.provider_budget_config: provider = self._get_llm_provider_for_deployment(deployment) + deployment_providers.append(provider) if provider is not None: budget_config = self._get_budget_config_for_provider(provider) if ( @@ -280,7 +341,12 @@ class RouterBudgetLimiting(CustomLogger): cache_keys.append( f"tag_spend:{_tag}:{_tag_budget_config.budget_duration}" ) - return cache_keys, provider_configs, deployment_configs + return ( + cache_keys, + provider_configs, + deployment_configs, + deployment_providers, + ) async def _get_or_set_budget_start_time( self, start_time_key: str, current_time: float, ttl_seconds: int @@ -597,12 +663,23 @@ class RouterBudgetLimiting(CustomLogger): def _get_llm_provider_for_deployment(self, deployment: Dict) -> Optional[str]: try: - _litellm_params: LiteLLM_Params = LiteLLM_Params( - **deployment.get("litellm_params", {"model": ""}) - ) + deployment_litellm_params = deployment.get("litellm_params") or {} + + if isinstance(deployment_litellm_params, LiteLLM_Params): + model = deployment_litellm_params.model or "" + provider_resolution_params: Any = deployment_litellm_params + elif isinstance(deployment_litellm_params, dict): + model = deployment_litellm_params.get("model") or "" + provider_resolution_params = _LiteLLMParamsDictView( + deployment_litellm_params + ) + else: + model = "" + provider_resolution_params = _LiteLLMParamsDictView({}) + _, custom_llm_provider, _, _ = litellm.get_llm_provider( - model=_litellm_params.model, - litellm_params=_litellm_params, + model=str(model), + litellm_params=provider_resolution_params, ) except Exception: verbose_router_logger.error( diff --git a/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py new file mode 100644 index 00000000000..82b7fc4d42c --- /dev/null +++ b/tests/test_litellm/router_strategy/test_budget_limiter_hotpath.py @@ -0,0 +1,232 @@ +import pytest + +import litellm +from litellm.caching.caching import DualCache +from litellm.router_strategy.budget_limiter import RouterBudgetLimiting +from litellm.types.router import LiteLLM_Params +from litellm.types.utils import BudgetConfig + + +@pytest.fixture +def disable_budget_sync(monkeypatch): + async def noop(*args, **kwargs): + return None + + monkeypatch.setattr( + "litellm.router_strategy.budget_limiter.RouterBudgetLimiting.periodic_sync_in_memory_spend_with_redis", + noop, + ) + + +@pytest.mark.asyncio +async def test_get_llm_provider_for_deployment_dict_does_not_require_litellm_params_instantiation( + disable_budget_sync, monkeypatch +): + class RaiseOnInit: + def __init__(self, *args, **kwargs): + raise AssertionError("LiteLLM_Params should not be instantiated in hot path") + + monkeypatch.setattr( + "litellm.router_strategy.budget_limiter.LiteLLM_Params", + RaiseOnInit, + ) + + provider_budget = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={}, + ) + + deployment = {"litellm_params": {"model": "openai/gpt-4o-mini"}} + provider = provider_budget._get_llm_provider_for_deployment(deployment) + + assert provider == "openai" + + +@pytest.mark.asyncio +async def test_get_llm_provider_for_deployment_dict_view_supports_mapping_and_attr_access( + disable_budget_sync, monkeypatch +): + observed = {} + + def _future_style_get_llm_provider( + model, + custom_llm_provider=None, + api_base=None, + api_key=None, + litellm_params=None, + ): + assert litellm_params is not None + observed["model_attr"] = litellm_params.model + observed["provider_get"] = litellm_params.get("custom_llm_provider") + observed["api_base_item"] = litellm_params["api_base"] + observed["has_api_key"] = "api_key" in litellm_params + observed["model_dump"] = litellm_params.model_dump() + return model, "openai", None, None + + monkeypatch.setattr( + "litellm.router_strategy.budget_limiter.litellm.get_llm_provider", + _future_style_get_llm_provider, + ) + + provider_budget = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={}, + ) + + deployment = { + "litellm_params": { + "model": "openai/gpt-4o-mini", + "custom_llm_provider": "openai", + "api_base": "https://api.openai.com/v1", + } + } + provider = provider_budget._get_llm_provider_for_deployment(deployment) + + assert provider == "openai" + assert observed["model_attr"] == "openai/gpt-4o-mini" + assert observed["provider_get"] == "openai" + assert observed["api_base_item"] == "https://api.openai.com/v1" + assert observed["has_api_key"] is False + assert observed["model_dump"]["model"] == "openai/gpt-4o-mini" + + +@pytest.mark.asyncio +async def test_async_filter_deployments_resolves_provider_once_per_deployment( + disable_budget_sync, monkeypatch +): + provider_budget = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={ + "openai": BudgetConfig(budget_duration="1d", max_budget=100.0), + }, + ) + + healthy_deployments = [ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": {"id": "deployment-1"}, + }, + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "openai/gpt-4o-mini"}, + "model_info": {"id": "deployment-2"}, + }, + ] + + provider_resolution_calls = 0 + + def _count_provider_calls(deployment): + nonlocal provider_resolution_calls + provider_resolution_calls += 1 + return "openai" + + monkeypatch.setattr( + provider_budget, + "_get_llm_provider_for_deployment", + _count_provider_calls, + ) + + filtered_deployments = await provider_budget.async_filter_deployments( + model="gpt-4o-mini", + healthy_deployments=healthy_deployments, + messages=[], + request_kwargs={}, + parent_otel_span=None, + ) + + assert len(filtered_deployments) == len(healthy_deployments) + assert provider_resolution_calls == len(healthy_deployments) + + +@pytest.mark.asyncio +async def test_async_filter_deployments_does_not_recompute_provider_when_resolved_none( + disable_budget_sync, monkeypatch +): + provider_budget = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={ + "openai": BudgetConfig(budget_duration="1d", max_budget=100.0), + }, + model_list=[ + { + "model_name": "gpt-4o-mini", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "max_budget": 100.0, + "budget_duration": "1d", + }, + "model_info": {"id": "deployment-1"}, + } + ], + ) + + healthy_deployments = [ + { + "model_name": "gpt-4o-mini", + "litellm_params": {"model": "unknown-provider/model"}, + "model_info": {"id": "deployment-1"}, + } + ] + + provider_resolution_calls = 0 + + def _provider_returns_none(deployment): + nonlocal provider_resolution_calls + provider_resolution_calls += 1 + return None + + monkeypatch.setattr( + provider_budget, + "_get_llm_provider_for_deployment", + _provider_returns_none, + ) + + filtered_deployments = await provider_budget.async_filter_deployments( + model="gpt-4o-mini", + healthy_deployments=healthy_deployments, + messages=[], + request_kwargs={}, + parent_otel_span=None, + ) + + assert len(filtered_deployments) == len(healthy_deployments) + assert provider_resolution_calls == len(healthy_deployments) + + +def _legacy_provider_resolution(deployment): + """ + Reference implementation used before hot-path optimization. + """ + try: + _litellm_params = LiteLLM_Params(**deployment.get("litellm_params", {"model": ""})) + _, custom_llm_provider, _, _ = litellm.get_llm_provider( + model=_litellm_params.model, + litellm_params=_litellm_params, + ) + except Exception: + return None + return custom_llm_provider + + +@pytest.mark.parametrize( + "deployment", + [ + {"litellm_params": {"model": "openai/gpt-4o-mini"}}, + {"litellm_params": {"model": "gpt-4o-mini", "custom_llm_provider": "openai"}}, + {"litellm_params": {"model": "unknown-provider/model"}}, + ], +) +@pytest.mark.asyncio +async def test_get_llm_provider_for_deployment_matches_legacy_behavior( + disable_budget_sync, deployment +): + provider_budget = RouterBudgetLimiting( + dual_cache=DualCache(), + provider_budget_config={}, + ) + + current_provider = provider_budget._get_llm_provider_for_deployment(deployment) + legacy_provider = _legacy_provider_resolution(deployment) + + assert current_provider == legacy_provider From da31dd19da19a81284d7ffc01ede17d924d0962b Mon Sep 17 00:00:00 2001 From: joaokopernico <111400483+joaokopernico@users.noreply.github.com> Date: Fri, 13 Feb 2026 01:10:08 -0300 Subject: [PATCH 2/9] fix(anthropic): use Authorization Bearer for OAuth tokens instead of x-api-key (#21039) OAuth tokens (sk-ant-oat*) require Authorization: Bearer header per Anthropic's OAuth specification, but were being sent via x-api-key which Anthropic rejects with 'invalid x-api-key'. - optionally_handle_anthropic_oauth: detect OAuth tokens in api_key param (standard chat flow), not just Authorization header - get_anthropic_headers: use Authorization: Bearer + required OAuth headers for OAuth tokens, x-api-key for regular API keys - Passthrough messages: skip x-api-key when Authorization is set - Add oauth-2025-04-20 to beta headers whitelist config --- litellm/anthropic_beta_headers_config.json | 1 + litellm/llms/anthropic/common_utils.py | 137 +++++--- .../messages/transformation.py | 34 +- .../anthropic/test_anthropic_common_utils.py | 313 ++++++++++++++---- 4 files changed, 369 insertions(+), 116 deletions(-) diff --git a/litellm/anthropic_beta_headers_config.json b/litellm/anthropic_beta_headers_config.json index 5edb8067a08..a066a5e95bc 100644 --- a/litellm/anthropic_beta_headers_config.json +++ b/litellm/anthropic_beta_headers_config.json @@ -19,6 +19,7 @@ "mcp-client-2025-11-20": "mcp-client-2025-11-20", "mcp-client-2025-04-04": "mcp-client-2025-04-04", "mcp-servers-2025-12-04": "mcp-servers-2025-12-04", + "oauth-2025-04-20": "oauth-2025-04-20", "output-128k-2025-02-19": "output-128k-2025-02-19", "prompt-caching-scope-2026-01-05": "prompt-caching-scope-2026-01-05", "skills-2025-10-02": "skills-2025-10-02", diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index cb23d21fbc9..c665e084261 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -38,9 +38,18 @@ def optionally_handle_anthropic_oauth( Returns: Tuple of (updated headers, api_key) """ + # Check Authorization header (passthrough / forwarded requests) auth_header = headers.get("authorization", "") if auth_header and auth_header.startswith(f"Bearer {ANTHROPIC_OAUTH_TOKEN_PREFIX}"): api_key = auth_header.replace("Bearer ", "") + headers.pop("x-api-key", None) + headers["anthropic-beta"] = ANTHROPIC_OAUTH_BETA_HEADER + headers["anthropic-dangerous-direct-browser-access"] = "true" + return headers, api_key + # Check api_key directly (standard chat/completion flow) + if api_key and api_key.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX): + headers.pop("x-api-key", None) + headers["authorization"] = f"Bearer {api_key}" headers["anthropic-beta"] = ANTHROPIC_OAUTH_BETA_HEADER headers["anthropic-dangerous-direct-browser-access"] = "true" return headers, api_key @@ -108,7 +117,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): if tools is None: return False for tool in tools: - if "type" in tool and tool["type"].startswith(ANTHROPIC_HOSTED_TOOLS.WEB_SEARCH.value): + if "type" in tool and tool["type"].startswith( + ANTHROPIC_HOSTED_TOOLS.WEB_SEARCH.value + ): return True return False @@ -134,111 +145,126 @@ class AnthropicModelInfo(BaseLLMModelInfo): """ if not tools: return False - + for tool in tools: tool_type = tool.get("type", "") - if tool_type in ["tool_search_tool_regex_20251119", "tool_search_tool_bm25_20251119"]: + if tool_type in [ + "tool_search_tool_regex_20251119", + "tool_search_tool_bm25_20251119", + ]: return True return False - + def is_programmatic_tool_calling_used(self, tools: Optional[List]) -> bool: """ Check if programmatic tool calling is being used (tools with allowed_callers field). - + Returns True if any tool has allowed_callers containing 'code_execution_20250825'. """ if not tools: return False - + for tool in tools: # Check top-level allowed_callers allowed_callers = tool.get("allowed_callers", None) if allowed_callers and isinstance(allowed_callers, list): if "code_execution_20250825" in allowed_callers: return True - + # Check function.allowed_callers for OpenAI format tools function = tool.get("function", {}) if isinstance(function, dict): function_allowed_callers = function.get("allowed_callers", None) - if function_allowed_callers and isinstance(function_allowed_callers, list): + if function_allowed_callers and isinstance( + function_allowed_callers, list + ): if "code_execution_20250825" in function_allowed_callers: return True - + return False - + def is_input_examples_used(self, tools: Optional[List]) -> bool: """ Check if input_examples is being used in any tools. - + Returns True if any tool has input_examples field. """ if not tools: return False - + for tool in tools: # Check top-level input_examples input_examples = tool.get("input_examples", None) - if input_examples and isinstance(input_examples, list) and len(input_examples) > 0: + if ( + input_examples + and isinstance(input_examples, list) + and len(input_examples) > 0 + ): return True - + # Check function.input_examples for OpenAI format tools function = tool.get("function", {}) if isinstance(function, dict): function_input_examples = function.get("input_examples", None) - if function_input_examples and isinstance(function_input_examples, list) and len(function_input_examples) > 0: + if ( + function_input_examples + and isinstance(function_input_examples, list) + and len(function_input_examples) > 0 + ): return True - + return False - - def is_effort_used(self, optional_params: Optional[dict], model: Optional[str] = None) -> bool: + + def is_effort_used( + self, optional_params: Optional[dict], model: Optional[str] = None + ) -> bool: """ Check if effort parameter is being used. - + Returns True if effort-related parameters are present. """ if not optional_params: return False - + # Check if reasoning_effort is provided for Claude Opus 4.5 if model and ("opus-4-5" in model.lower() or "opus_4_5" in model.lower()): reasoning_effort = optional_params.get("reasoning_effort") if reasoning_effort and isinstance(reasoning_effort, str): return True - + # Check if output_config is directly provided output_config = optional_params.get("output_config") if output_config and isinstance(output_config, dict): effort = output_config.get("effort") if effort and isinstance(effort, str): return True - + return False def is_code_execution_tool_used(self, tools: Optional[List]) -> bool: """ Check if code execution tool is being used. - + Returns True if any tool has type "code_execution_20250825". """ if not tools: return False - + for tool in tools: tool_type = tool.get("type", "") if tool_type == "code_execution_20250825": return True return False - + def is_container_with_skills_used(self, optional_params: Optional[dict]) -> bool: """ Check if container with skills is being used. - + Returns True if optional_params contains container with skills. """ if not optional_params: return False - + container = optional_params.get("container") if container and isinstance(container, dict): skills = container.get("skills") @@ -256,10 +282,10 @@ class AnthropicModelInfo(BaseLLMModelInfo): def get_computer_tool_beta_header(self, computer_tool_version: str) -> str: """ Get the appropriate beta header for a given computer tool version. - + Args: computer_tool_version: The computer tool version (e.g., 'computer_20250124', 'computer_20241022') - + Returns: The corresponding beta header string """ @@ -282,37 +308,37 @@ class AnthropicModelInfo(BaseLLMModelInfo): ) -> List[str]: """ Get list of common beta headers based on the features that are active. - + Returns: List of beta header strings """ from litellm.types.llms.anthropic import ( ANTHROPIC_EFFORT_BETA_HEADER, ) - + betas = [] - + # Detect features effort_used = self.is_effort_used(optional_params, model) - + if effort_used: betas.append(ANTHROPIC_EFFORT_BETA_HEADER) # effort-2025-11-24 - + if computer_tool_used: beta_header = self.get_computer_tool_beta_header(computer_tool_used) betas.append(beta_header) - + # Anthropic no longer requires the prompt-caching beta header # Prompt caching now works automatically when cache_control is used in messages # Reference: https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching - + if file_id_used: betas.append("files-api-2025-04-14") betas.append("code-execution-2025-05-22") - + if mcp_server_used: betas.append("mcp-client-2025-04-04") - + return list(set(betas)) def get_anthropic_headers( @@ -351,27 +377,35 @@ class AnthropicModelInfo(BaseLLMModelInfo): # Tool search, programmatic tool calling, and input_examples all use the same beta header if tool_search_used or programmatic_tool_calling_used or input_examples_used: from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER + betas.add(ANTHROPIC_TOOL_SEARCH_BETA_HEADER) - + # Effort parameter uses a separate beta header if effort_used: from litellm.types.llms.anthropic import ANTHROPIC_EFFORT_BETA_HEADER + betas.add(ANTHROPIC_EFFORT_BETA_HEADER) - + # Code execution tool uses a separate beta header if code_execution_tool_used: betas.add("code-execution-2025-08-25") - + # Container with skills uses a separate beta header if container_with_skills_used: betas.add("skills-2025-10-02") + _is_oauth = api_key and api_key.startswith(ANTHROPIC_OAUTH_TOKEN_PREFIX) headers = { "anthropic-version": anthropic_version or "2023-06-01", - "x-api-key": api_key, "accept": "application/json", "content-type": "application/json", } + if _is_oauth: + headers["authorization"] = f"Bearer {api_key}" + headers["anthropic-dangerous-direct-browser-access"] = "true" + betas.add(ANTHROPIC_OAUTH_BETA_HEADER) + else: + headers["x-api-key"] = api_key if user_anthropic_beta_headers is not None: betas.update(user_anthropic_beta_headers) @@ -381,7 +415,10 @@ class AnthropicModelInfo(BaseLLMModelInfo): # Vertex AI requires web search beta header for web search to work if web_search_tool_used: from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES - headers["anthropic-beta"] = ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value + + headers[ + "anthropic-beta" + ] = ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value elif len(betas) > 0: headers["anthropic-beta"] = ",".join(betas) @@ -398,7 +435,9 @@ class AnthropicModelInfo(BaseLLMModelInfo): api_base: Optional[str] = None, ) -> Dict: # Check for Anthropic OAuth token in headers - headers, api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key) + headers, api_key = optionally_handle_anthropic_oauth( + headers=headers, api_key=api_key + ) if api_key is None: raise litellm.AuthenticationError( message="Missing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via params. Please set `ANTHROPIC_API_KEY` in your environment vars", @@ -416,11 +455,15 @@ class AnthropicModelInfo(BaseLLMModelInfo): file_id_used = self.is_file_id_used(messages=messages) web_search_tool_used = self.is_web_search_tool_used(tools=tools) tool_search_used = self.is_tool_search_used(tools=tools) - programmatic_tool_calling_used = self.is_programmatic_tool_calling_used(tools=tools) + programmatic_tool_calling_used = self.is_programmatic_tool_calling_used( + tools=tools + ) input_examples_used = self.is_input_examples_used(tools=tools) effort_used = self.is_effort_used(optional_params=optional_params, model=model) code_execution_tool_used = self.is_code_execution_tool_used(tools=tools) - container_with_skills_used = self.is_container_with_skills_used(optional_params=optional_params) + container_with_skills_used = self.is_container_with_skills_used( + optional_params=optional_params + ) user_anthropic_beta_headers = self._get_user_anthropic_beta_headers( anthropic_beta_header=headers.get("anthropic-beta") ) @@ -499,7 +542,7 @@ class AnthropicModelInfo(BaseLLMModelInfo): def get_token_counter(self) -> Optional[BaseTokenCounter]: """ Factory method to create an Anthropic token counter. - + Returns: AnthropicTokenCounter instance for this provider. """ diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py index 8f2f3bf3545..8275ba2b3e1 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/transformation.py @@ -49,15 +49,15 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): # TODO: Add Anthropic `metadata` support # "metadata", ] - + @staticmethod def _filter_billing_headers_from_system(system_param): """ Filter out x-anthropic-billing-header metadata from system parameter. - + Args: system_param: Can be a string or a list of system message content blocks - + Returns: Filtered system parameter (string or list), or None if all content was filtered """ @@ -74,7 +74,9 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): text = content_block.get("text", "") content_type = content_block.get("type", "") # Skip text blocks that start with billing header - if content_type == "text" and text.startswith("x-anthropic-billing-header:"): + if content_type == "text" and text.startswith( + "x-anthropic-billing-header:" + ): continue filtered_list.append(content_block) else: @@ -111,11 +113,13 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): import os # Check for Anthropic OAuth token in Authorization header - headers, api_key = optionally_handle_anthropic_oauth(headers=headers, api_key=api_key) + headers, api_key = optionally_handle_anthropic_oauth( + headers=headers, api_key=api_key + ) if api_key is None: api_key = os.getenv("ANTHROPIC_API_KEY") - if "x-api-key" not in headers and api_key: + if "x-api-key" not in headers and "authorization" not in headers and api_key: headers["x-api-key"] = api_key if "anthropic-version" not in headers: headers["anthropic-version"] = DEFAULT_ANTHROPIC_API_VERSION @@ -149,7 +153,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): message="max_tokens is required for Anthropic /v1/messages API", status_code=400, ) - + # Filter out x-anthropic-billing-header from system messages system_param = anthropic_messages_optional_request_params.get("system") if system_param is not None: @@ -159,7 +163,7 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): else: # Remove system parameter if all content was filtered out anthropic_messages_optional_request_params.pop("system", None) - + ####### get required params for all anthropic messages requests ###### verbose_logger.debug(f"TRANSFORMATION DEBUG - Messages: {messages}") anthropic_messages_request: AnthropicMessagesRequest = AnthropicMessagesRequest( @@ -244,25 +248,29 @@ class AnthropicMessagesConfig(BaseAnthropicMessagesConfig): edits = context_management_param.get("edits", []) has_compact = False has_other = False - + for edit in edits: edit_type = edit.get("type", "") if edit_type == "compact_20260112": has_compact = True else: has_other = True - + # Add compact header if any compact edits exist if has_compact: beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value) - + # Add context management header if any other edits exist if has_other: - beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value) + beta_values.add( + ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value + ) # Check for structured outputs if optional_params.get("output_format") is not None: - beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.STRUCTURED_OUTPUT_2025_09_25.value) + beta_values.add( + ANTHROPIC_BETA_HEADER_VALUES.STRUCTURED_OUTPUT_2025_09_25.value + ) # Check for fast mode if optional_params.get("speed") == "fast": diff --git a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py index 0a397d116e7..a321a24540f 100644 --- a/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py +++ b/tests/test_litellm/llms/anthropic/test_anthropic_common_utils.py @@ -1,84 +1,285 @@ """ -Tests for Anthropic OAuth token handling for Claude Code Max integration. +Tests for Anthropic OAuth token handling in common_utils. + +Verifies that OAuth tokens (sk-ant-oat*) are sent via Authorization: Bearer +instead of x-api-key, per Anthropic's OAuth specification. """ import os import sys -# Add litellm to path -sys.path.insert(0, os.path.abspath("../../../../..")) +sys.path.insert( + 0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../..")) +) # Fake OAuth token for testing (not a real secret) FAKE_OAUTH_TOKEN = "sk-ant-oat01-fake-token-for-testing-123456789abcdef" +FAKE_REGULAR_KEY = "sk-ant-api03-regular-key-for-testing-123456789" -def test_oauth_detection_in_common_utils(): - """Test 1: OAuth token detection in common_utils""" - from litellm.llms.anthropic.common_utils import optionally_handle_anthropic_oauth +class TestOptionallyHandleAnthropicOAuth: + """Tests for optionally_handle_anthropic_oauth function.""" - headers = {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"} - updated_headers, extracted_api_key = optionally_handle_anthropic_oauth(headers, None) + def test_oauth_token_in_authorization_header(self): + """OAuth token in Authorization header should be detected and headers set correctly.""" + from litellm.llms.anthropic.common_utils import ( + optionally_handle_anthropic_oauth, + ) - assert extracted_api_key == FAKE_OAUTH_TOKEN - assert updated_headers["anthropic-beta"] == "oauth-2025-04-20" - assert updated_headers["anthropic-dangerous-direct-browser-access"] == "true" + headers = {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"} + updated_headers, extracted_api_key = optionally_handle_anthropic_oauth( + headers, None + ) + + assert extracted_api_key == FAKE_OAUTH_TOKEN + assert updated_headers["anthropic-beta"] == "oauth-2025-04-20" + assert updated_headers["anthropic-dangerous-direct-browser-access"] == "true" + assert "x-api-key" not in updated_headers + + def test_oauth_token_in_api_key_directly(self): + """OAuth token passed as api_key should set Authorization: Bearer header.""" + from litellm.llms.anthropic.common_utils import ( + optionally_handle_anthropic_oauth, + ) + + headers = {} + updated_headers, returned_api_key = optionally_handle_anthropic_oauth( + headers, FAKE_OAUTH_TOKEN + ) + + assert returned_api_key == FAKE_OAUTH_TOKEN + assert updated_headers["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + assert updated_headers["anthropic-beta"] == "oauth-2025-04-20" + assert updated_headers["anthropic-dangerous-direct-browser-access"] == "true" + assert "x-api-key" not in updated_headers + + def test_oauth_removes_existing_x_api_key(self): + """When OAuth is detected, any existing x-api-key should be removed.""" + from litellm.llms.anthropic.common_utils import ( + optionally_handle_anthropic_oauth, + ) + + headers = {"x-api-key": FAKE_OAUTH_TOKEN} + updated_headers, _ = optionally_handle_anthropic_oauth( + headers, FAKE_OAUTH_TOKEN + ) + + assert "x-api-key" not in updated_headers + assert updated_headers["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + + def test_regular_api_key_unchanged(self): + """Regular API keys (non-OAuth) should pass through unmodified.""" + from litellm.llms.anthropic.common_utils import ( + optionally_handle_anthropic_oauth, + ) + + headers = {} + updated_headers, returned_api_key = optionally_handle_anthropic_oauth( + headers, FAKE_REGULAR_KEY + ) + + assert returned_api_key == FAKE_REGULAR_KEY + assert "authorization" not in updated_headers + assert "anthropic-dangerous-direct-browser-access" not in updated_headers + assert "anthropic-beta" not in updated_headers + + def test_regular_key_in_authorization_header(self): + """Non-OAuth token in Authorization header should not trigger OAuth handling.""" + from litellm.llms.anthropic.common_utils import ( + optionally_handle_anthropic_oauth, + ) + + headers = {"authorization": f"Bearer {FAKE_REGULAR_KEY}"} + updated_headers, returned_api_key = optionally_handle_anthropic_oauth( + headers, FAKE_REGULAR_KEY + ) + + assert returned_api_key == FAKE_REGULAR_KEY + assert "anthropic-dangerous-direct-browser-access" not in updated_headers + + def test_none_api_key_no_error(self): + """None api_key with empty headers should not raise errors.""" + from litellm.llms.anthropic.common_utils import ( + optionally_handle_anthropic_oauth, + ) + + headers = {} + updated_headers, returned_api_key = optionally_handle_anthropic_oauth( + headers, None + ) + + assert returned_api_key is None + assert "authorization" not in updated_headers -def test_oauth_integration_in_validate_environment(): - """Test 2: OAuth integration in AnthropicConfig validate_environment""" - from litellm.llms.anthropic.common_utils import AnthropicModelInfo +class TestGetAnthropicHeaders: + """Tests for get_anthropic_headers method with OAuth support.""" - config = AnthropicModelInfo() - headers = {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"} + def test_oauth_token_uses_authorization_bearer(self): + """OAuth token should produce Authorization: Bearer header, not x-api-key.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo - updated_headers = config.validate_environment( - headers=headers, - model="claude-3-haiku-20240307", - messages=[{"role": "user", "content": "Hello"}], - optional_params={}, - litellm_params={}, - api_key=None, - api_base=None, - ) + config = AnthropicModelInfo() + headers = config.get_anthropic_headers( + api_key=FAKE_OAUTH_TOKEN, + computer_tool_used=False, + prompt_caching_set=False, + pdf_used=False, + is_vertex_request=False, + ) - assert updated_headers["x-api-key"] == FAKE_OAUTH_TOKEN - assert updated_headers["anthropic-dangerous-direct-browser-access"] == "true" + assert headers["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + assert headers["anthropic-dangerous-direct-browser-access"] == "true" + assert "oauth-2025-04-20" in headers.get("anthropic-beta", "") + assert "x-api-key" not in headers + + def test_regular_key_uses_x_api_key(self): + """Regular API key should produce x-api-key header, not Authorization.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = config.get_anthropic_headers( + api_key=FAKE_REGULAR_KEY, + computer_tool_used=False, + prompt_caching_set=False, + pdf_used=False, + is_vertex_request=False, + ) + + assert headers["x-api-key"] == FAKE_REGULAR_KEY + assert "authorization" not in headers + assert "anthropic-dangerous-direct-browser-access" not in headers + + def test_oauth_includes_standard_headers(self): + """OAuth path should still include standard Anthropic headers.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = config.get_anthropic_headers( + api_key=FAKE_OAUTH_TOKEN, + computer_tool_used=False, + prompt_caching_set=False, + pdf_used=False, + is_vertex_request=False, + ) + + assert headers["anthropic-version"] == "2023-06-01" + assert headers["accept"] == "application/json" + assert headers["content-type"] == "application/json" -def test_oauth_detection_in_messages_transformation(): - """Test 3: OAuth detection in messages transformation""" - from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( - AnthropicMessagesConfig, - ) +class TestValidateEnvironmentOAuth: + """Tests for validate_environment with OAuth tokens.""" - config = AnthropicMessagesConfig() - headers = {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"} + def test_oauth_via_authorization_header(self): + """validate_environment should produce correct headers for OAuth tokens.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo - updated_headers, _ = config.validate_anthropic_messages_environment( - headers=headers, - model="claude-3-haiku-20240307", - messages=[{"role": "user", "content": "Hello"}], - optional_params={}, - litellm_params={}, - api_key=None, - api_base=None, - ) + config = AnthropicModelInfo() + headers = {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"} - assert updated_headers["x-api-key"] == FAKE_OAUTH_TOKEN - assert "oauth-2025-04-20" in updated_headers["anthropic-beta"] - assert updated_headers["anthropic-dangerous-direct-browser-access"] == "true" + updated_headers = config.validate_environment( + headers=headers, + model="claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + + assert updated_headers["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + assert updated_headers["anthropic-dangerous-direct-browser-access"] == "true" + assert "oauth-2025-04-20" in updated_headers.get("anthropic-beta", "") + assert "x-api-key" not in updated_headers + + def test_oauth_via_api_key_param(self): + """validate_environment with OAuth token as api_key should use Bearer auth.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = {} + + updated_headers = config.validate_environment( + headers=headers, + model="claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + api_key=FAKE_OAUTH_TOKEN, + api_base=None, + ) + + assert updated_headers["authorization"] == f"Bearer {FAKE_OAUTH_TOKEN}" + assert updated_headers["anthropic-dangerous-direct-browser-access"] == "true" + assert "x-api-key" not in updated_headers + + def test_regular_key_via_api_key_param(self): + """validate_environment with regular API key should use x-api-key.""" + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + config = AnthropicModelInfo() + headers = {} + + updated_headers = config.validate_environment( + headers=headers, + model="claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + api_key=FAKE_REGULAR_KEY, + api_base=None, + ) + + assert updated_headers["x-api-key"] == FAKE_REGULAR_KEY + assert "authorization" not in updated_headers + assert "anthropic-dangerous-direct-browser-access" not in updated_headers -def test_regular_api_keys_still_work(): - """Test 4: Regular API keys still work (regression test)""" - from litellm.llms.anthropic.common_utils import optionally_handle_anthropic_oauth +class TestPassthroughOAuth: + """Tests for passthrough messages endpoint with OAuth tokens.""" - regular_key = "sk-ant-api03-regular-key-123" - headers = {"authorization": f"Bearer {regular_key}"} + def test_passthrough_oauth_no_x_api_key(self): + """Passthrough endpoint should not add x-api-key for OAuth tokens.""" + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) - updated_headers, extracted_api_key = optionally_handle_anthropic_oauth(headers, regular_key) + config = AnthropicMessagesConfig() + headers = {"authorization": f"Bearer {FAKE_OAUTH_TOKEN}"} - # Regular key should be unchanged - assert extracted_api_key == regular_key - # OAuth headers should NOT be added - assert "anthropic-dangerous-direct-browser-access" not in updated_headers \ No newline at end of file + updated_headers, _ = config.validate_anthropic_messages_environment( + headers=headers, + model="claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + api_key=None, + api_base=None, + ) + + assert "oauth-2025-04-20" in updated_headers.get("anthropic-beta", "") + assert updated_headers["anthropic-dangerous-direct-browser-access"] == "true" + assert "x-api-key" not in updated_headers + + def test_passthrough_regular_key_uses_x_api_key(self): + """Passthrough endpoint should still use x-api-key for regular API keys.""" + from litellm.llms.anthropic.experimental_pass_through.messages.transformation import ( + AnthropicMessagesConfig, + ) + + config = AnthropicMessagesConfig() + headers = {} + + updated_headers, _ = config.validate_anthropic_messages_environment( + headers=headers, + model="claude-sonnet-4-5-20250929", + messages=[{"role": "user", "content": "Hello"}], + optional_params={}, + litellm_params={}, + api_key=FAKE_REGULAR_KEY, + api_base=None, + ) + + assert updated_headers["x-api-key"] == FAKE_REGULAR_KEY + assert "authorization" not in updated_headers From 1bc90d2db795dbb37dfb5dc70154457a2885fdc9 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Thu, 12 Feb 2026 23:19:13 -0500 Subject: [PATCH 3/9] fix guardrail status error (#20972) * fix guardrail status error * fix function imports --- litellm/integrations/custom_guardrail.py | 25 +++++- .../guardrail_hooks/test_model_armor.py | 77 +++++++++++++++++++ 2 files changed, 101 insertions(+), 1 deletion(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 407bc581f71..967de21493a 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -29,6 +29,7 @@ from litellm.types.utils import ( LLMResponseTypes, StandardLoggingGuardrailInformation, ) +from fastapi.exceptions import HTTPException if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -648,6 +649,23 @@ class CustomGuardrail(CustomLogger): ) return response + @staticmethod + def _is_guardrail_intervention(e: Exception) -> bool: + """ + Returns True if the exception represents an intentional guardrail block + (this was logged previously as an API failure - guardrail_failed_to_respond). + + Guardrails signal intentional blocks by raising: + - HTTPException with status 400 (content policy violation) + - ModifyResponseException (passthrough mode violation) + """ + + if isinstance(e, ModifyResponseException): + return True + if isinstance(e, HTTPException) and e.status_code == 400: + return True + return False + def _process_error( self, e: Exception, @@ -662,6 +680,11 @@ class CustomGuardrail(CustomLogger): This gets logged on downsteam Langfuse, DataDog, etc. """ + guardrail_status: GuardrailStatus = ( + "guardrail_intervened" + if self._is_guardrail_intervention(e) + else "guardrail_failed_to_respond" + ) # For custom_code_guardrail scenario, log as "deny" instead of full exception # Check if this is from custom_code_guardrail by checking the class name guardrail_response: Union[Exception, str] = e @@ -671,7 +694,7 @@ class CustomGuardrail(CustomLogger): self.add_standard_logging_guardrail_information_to_request_data( guardrail_json_response=guardrail_response, request_data=request_data, - guardrail_status="guardrail_failed_to_respond", + guardrail_status=guardrail_status, duration=duration, start_time=start_time, end_time=end_time, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py index 987388a80c7..8080491f662 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_model_armor.py @@ -1122,6 +1122,83 @@ async def test_model_armor_non_model_response(): assert not guardrail.async_handler.post.called +@pytest.mark.asyncio +async def test_model_armor_guardrail_status_intervened_vs_failed(): + """ + regression test for bug where _process_error always set 'guardrail_failed_to_respond' + even for intentional blocks (error 400). + """ + mock_user_api_key_dict = UserAPIKeyAuth() + mock_cache = MagicMock(spec=DualCache) + + #1: Blocked content should raise exception and show guardrail status: guardrail_intervened" + guardrail = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test", + ) + + mock_response = AsyncMock() + mock_response.status_code = 200 + mock_response.json = AsyncMock(return_value={ + "sanitizationResult": { + "filterMatchState": "MATCH_FOUND", + "filterResults": { + "rai": { + "raiFilterResult": { + "matchState": "MATCH_FOUND", + } + } + } + } + }) + + guardrail._ensure_access_token_async = AsyncMock(return_value=("token", "test-project")) + with patch.object(guardrail.async_handler, "post", AsyncMock(return_value=mock_response)): + request_data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "bad content"}], + "metadata": {"guardrails": ["model-armor-test"]}, + } + with pytest.raises(HTTPException): + await guardrail.async_pre_call_hook( + user_api_key_dict=mock_user_api_key_dict, + cache=mock_cache, + data=request_data, + call_type="completion", + ) + + info = request_data["metadata"]["standard_logging_guardrail_information"] + assert info[0]["guardrail_status"] == "guardrail_intervened" + + #2: if an API error - guardrail status should be guardrail_failed_to_respond" + guardrail2 = ModelArmorGuardrail( + template_id="test-template", + project_id="test-project", + location="us-central1", + guardrail_name="model-armor-test2", + fail_on_error=True, + ) + + guardrail2._ensure_access_token_async = AsyncMock(side_effect=ConnectionError("timeout")) + request_data2 = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "hello"}], + "metadata": {"guardrails": ["model-armor-test2"]}, + } + with pytest.raises(ConnectionError): + await guardrail2.async_pre_call_hook( + user_api_key_dict=mock_user_api_key_dict, + cache=mock_cache, + data=request_data2, + call_type="completion", + ) + + info2 = request_data2["metadata"]["standard_logging_guardrail_information"] + assert info2[0]["guardrail_status"] == "guardrail_failed_to_respond" + + def mock_open(read_data=''): """Helper to create a mock file object""" import io From f74fdfbb61ff22436bda3cfe8db19422153d4301 Mon Sep 17 00:00:00 2001 From: datzscaler Date: Thu, 12 Feb 2026 20:27:44 -0800 Subject: [PATCH 4/9] feat(ui): added UI for Zscaler AI Guard (#21077) * fix: allow Management keys to access user/daily/activity and team/daily/activity * feat(ui): added UI for Zscaler AI Guard * feat(ui): addressed UI comment * Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --------- Co-authored-by: naaa760 Co-authored-by: yuneng-jiang Co-authored-by: Krish Dholakia Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/_types.py | 3 + .../zscaler_ai_guard/zscaler_ai_guard.py | 9 ++ .../guardrail_hooks/zscaler_ai_guard.py | 132 ++++++++++++++++++ .../public/assets/logos/zscaler.svg | 5 + .../guardrails/guardrail_info_helpers.tsx | 1 + 5 files changed, 150 insertions(+) create mode 100644 litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py create mode 100644 ui/litellm-dashboard/public/assets/logos/zscaler.svg diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0c85e9ba4f2..de785b221bc 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -514,6 +514,8 @@ class LiteLLMRoutes(enum.Enum): "/user/delete", "/user/info", "/user/list", + "/user/daily/activity", + "/user/daily/activity/aggregated", # team "/team/new", "/team/update", @@ -526,6 +528,7 @@ class LiteLLMRoutes(enum.Enum): "/team/available", "/team/permissions_list", "/team/permissions_update", + "/team/daily/activity", # model "/model/new", "/model/update", diff --git a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py index ff00cd73ca5..b0b18166d2d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py +++ b/litellm/proxy/guardrails/guardrail_hooks/zscaler_ai_guard/zscaler_ai_guard.py @@ -21,6 +21,7 @@ from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel GUARDRAIL_TIMEOUT = 5 @@ -334,3 +335,11 @@ class ZscalerAIGuard(CustomGuardrail): user_facing_error = self._create_user_facing_error(f"{str(e)})") # This exception will be caught by the proxy and returned to the user raise HTTPException(status_code=500, detail=user_facing_error) + + @staticmethod + def get_config_model() -> Optional[type["GuardrailConfigModel"]]: + from litellm.types.proxy.guardrails.guardrail_hooks.zscaler_ai_guard import ( + ZscalerAIGuardConfigModel, + ) + + return ZscalerAIGuardConfigModel diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py b/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py new file mode 100644 index 00000000000..7cbdf751e1b --- /dev/null +++ b/litellm/types/proxy/guardrails/guardrail_hooks/zscaler_ai_guard.py @@ -0,0 +1,132 @@ +from typing import Optional + +from pydantic import Field, model_validator + +from litellm._logging import verbose_proxy_logger +from litellm.types.guardrails import GuardrailParamUITypes + +from .base import GuardrailConfigModel + + +class ZscalerAIGuardConfigModel(GuardrailConfigModel): + api_key: Optional[str] = Field( + default=None, + description=( + "API key for Zscaler AI Guard authentication. " + "If not provided, falls back to ZSCALER_AI_GUARD_API_KEY environment variable." + ), + ) + + api_base: Optional[str] = Field( + default=None, + description=( + "Zscaler AI Guard API endpoint. Determines policy resolution behavior:\n" + "• /execute-policy (default) - Requires explicit policy_id in configuration\n" + "• /resolve-and-execute-policy - Infers policy from user-api-key-alias header\n" + "Default: https://api.us1.zseclipse.net/v1/detection/execute-policy\n" + "Falls back to ZSCALER_AI_GUARD_URL environment variable." + ), + json_schema_extra={ + "examples": [ + "https://api.us1.zseclipse.net/v1/detection/execute-policy", + "https://api.us1.zseclipse.net/v1/detection/resolve-and-execute-policy", + ] + }, + ) + + policy_id: Optional[int] = Field( + default=None, + description=( + "Global policy ID for Zscaler AI Guard. Required when using /execute-policy endpoint.\n\n" + "Set to 0 or leave empty when using /resolve-and-execute-policy with dynamic policy resolution.\n" + "Falls back to ZSCALER_AI_GUARD_POLICY_ID environment variable." + ), + json_schema_extra={ + "ui_hint": "conditional_required", + "condition": "Required when api_base ends with /execute-policy", + }, + ) + + send_user_api_key_alias: Optional[bool] = Field( + default=False, + description=( + "Send user API key alias in request headers as 'user-api-key-alias'. " + "CRITICAL when using /resolve-and-execute-policy endpoint - the policy is inferred from this value. " + "Also useful for tracking/auditing with /execute-policy endpoint." + ), + json_schema_extra={ + "ui_type": GuardrailParamUITypes.BOOL, + "ui_hint": "recommended_when", + "condition": "Recommended when api_base ends with /resolve-and-execute-policy", + }, + ) + + send_user_api_key_user_id: Optional[bool] = Field( + default=False, + description=( + "Send user API key user_id in request headers as 'user-api-key-user-id'. " + "Enables user-level tracking and analytics in Zscaler AI Guard." + ), + json_schema_extra={"ui_type": GuardrailParamUITypes.BOOL}, + ) + + send_user_api_key_team_id: Optional[bool] = Field( + default=False, + description=( + "Send user API key team_id in request headers as 'user-api-key-team-id'. " + "Enables team-level tracking and analytics in Zscaler AI Guard." + ), + json_schema_extra={"ui_type": GuardrailParamUITypes.BOOL}, + ) + + @model_validator(mode="after") + def validate_endpoint_configuration(self) -> "ZscalerAIGuardConfigModel": + """ + Validate configuration consistency between api_base and other fields. + Provides warnings but doesn't block (since env vars might provide values). + """ + import os + + # Resolve actual api_base value (including env fallback) + api_base = self.api_base or os.getenv( + "ZSCALER_AI_GUARD_URL", + "https://api.us1.zseclipse.net/v1/detection/execute-policy", + ) + + # Resolve actual policy_id value + policy_id = self.policy_id + if policy_id is None: + env_policy = os.getenv("ZSCALER_AI_GUARD_POLICY_ID") + if env_policy: + try: + policy_id = int(env_policy) + except ValueError: + verbose_proxy_logger.warning( + f"ZSCALER_AI_GUARD_POLICY_ID env var is not a valid integer: {env_policy}" + ) + + # Check for configuration issues + is_resolve_policy = api_base.endswith("/resolve-and-execute-policy") + is_execute_policy = api_base.endswith("/execute-policy") and not is_resolve_policy + + # Scenario A: execute-policy without policy_id + if is_execute_policy and (policy_id is None or policy_id < 1): + verbose_proxy_logger.warning( + "Using /execute-policy endpoint without a valid policy_id. " + "Ensure ZSCALER_AI_GUARD_POLICY_ID environment variable is set, " + "or provide policy_id via request/key/team metadata." + ) + + # Scenario B: resolve-and-execute-policy without user_api_key_alias + if is_resolve_policy and not self.send_user_api_key_alias: + verbose_proxy_logger.warning( + "Using /resolve-and-execute-policy endpoint without send_user_api_key_alias=true. " + "The endpoint requires user-api-key-alias header to resolve the policy. " + "Set send_user_api_key_alias to true or ensure the header is sent via other means." + ) + + return self + + @staticmethod + def ui_friendly_name() -> str: + return "Zscaler AI Guard" diff --git a/ui/litellm-dashboard/public/assets/logos/zscaler.svg b/ui/litellm-dashboard/public/assets/logos/zscaler.svg new file mode 100644 index 00000000000..2a95cb02aed --- /dev/null +++ b/ui/litellm-dashboard/public/assets/logos/zscaler.svg @@ -0,0 +1,5 @@ + + + + + diff --git a/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx b/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx index c6314c95bef..bde21661d89 100644 --- a/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/guardrails/guardrail_info_helpers.tsx @@ -104,6 +104,7 @@ export const shouldRenderContentFilterConfigSettings = (provider: string | null) const asset_logos_folder = "../ui/assets/logos/"; export const guardrailLogoMap: Record = { + "Zscaler AI Guard": `${asset_logos_folder}zscaler.svg`, "Presidio PII": `${asset_logos_folder}presidio.png`, "Bedrock Guardrail": `${asset_logos_folder}bedrock.svg`, Lakera: `${asset_logos_folder}lakeraai.jpeg`, From 20ae67a3ba36b1dc7e8d67c638762d013596c7ee Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Fri, 13 Feb 2026 01:39:55 -0300 Subject: [PATCH 5/9] fix(vertex): map IMAGE_PROHIBITED_CONTENT finish reason to content_filter (#20524) Closes #20357 --- .../llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 04ae4b6beb8..ad046cf7109 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -1196,6 +1196,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "PROHIBITED_CONTENT": "The token generation was stopped as the response was flagged for the prohibited contents.", "SPII": "The token generation was stopped as the response was flagged for Sensitive Personally Identifiable Information (SPII) contents.", "IMAGE_SAFETY": "The token generation was stopped as the response was flagged for image safety reasons.", + "IMAGE_PROHIBITED_CONTENT": "The token generation was stopped as the response was flagged for prohibited image content.", } @staticmethod @@ -1218,6 +1219,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): "SPII": "content_filter", "MALFORMED_FUNCTION_CALL": "malformed_function_call", # openai doesn't have a way of representing this "IMAGE_SAFETY": "content_filter", + "IMAGE_PROHIBITED_CONTENT": "content_filter", } def translate_exception_str(self, exception_string: str): From 68d2306dd4a4b4217856c6af3a89e3dc69b0b870 Mon Sep 17 00:00:00 2001 From: Lei Nie Date: Thu, 12 Feb 2026 20:44:59 -0800 Subject: [PATCH 6/9] feat(vertex_ai): preserve usageMetadata in _hidden_params (#20559) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: allow Management keys to access user/daily/activity and team/daily/activity * feat(vertex): surface trafficType via generic provider_specific_fields in Responses API Extract Vertex AI's trafficType from usageMetadata in both streaming and non-streaming paths, storing it in _hidden_params["provider_specific_fields"]. The Responses API transformation layer generically passes any _hidden_params["provider_specific_fields"] dict to the ResponsesAPIResponse, avoiding provider-specific logic in the bridge. Also fix stream_chunk_builder to propagate _hidden_params from the last streaming chunk to the rebuilt ModelResponse, ensuring provider metadata survives the chunk→response rebuild. --------- Co-authored-by: naaa760 Co-authored-by: yuneng-jiang --- .../vertex_and_google_ai_studio_gemini.py | 22 ++++- litellm/main.py | 10 ++ .../transformation.py | 6 ++ ...test_vertex_and_google_ai_studio_gemini.py | 91 +++++++++++++++++++ 4 files changed, 127 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index ad046cf7109..bef83b6d35e 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -480,7 +480,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): tool = {VertexToolName.COMPUTER_USE.value: computer_use_config} # Handle OpenAI-style web_search and web_search_preview tools # Transform them to Gemini's googleSearch tool - elif "type" in tool and tool["type"] in ("web_search", "web_search_preview"): + elif "type" in tool and tool["type"] in ( + "web_search", + "web_search_preview", + ): verbose_logger.info( f"Gemini: Transforming OpenAI-style '{tool['type']}' tool to googleSearch" ) @@ -1632,7 +1635,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): completion_image_tokens = response_tokens_details.image_tokens or 0 completion_audio_tokens = response_tokens_details.audio_tokens or 0 calculated_text_tokens = ( - candidates_token_count - completion_image_tokens - completion_audio_tokens + candidates_token_count + - completion_image_tokens + - completion_audio_tokens ) response_tokens_details.text_tokens = calculated_text_tokens ######################################################### @@ -2250,6 +2255,13 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): citation_metadata # older approach - maintaining to prevent regressions ) + ## ADD TRAFFIC TYPE ## + traffic_type = completion_response.get("usageMetadata", {}).get( + "trafficType" + ) + if traffic_type: + model_response._hidden_params.setdefault("provider_specific_fields", {})["traffic_type"] = traffic_type + except Exception as e: raise VertexAIError( message="Received={}, Error converting to valid response block={}. File an issue if litellm error - https://github.com/BerriAI/litellm/issues".format( @@ -2908,6 +2920,12 @@ class ModelResponseIterator: PromptTokensDetailsWrapper, usage.prompt_tokens_details ).web_search_requests = web_search_requests + traffic_type = processed_chunk.get("usageMetadata", {}).get( + "trafficType" + ) + if traffic_type: + model_response._hidden_params.setdefault("provider_specific_fields", {})["traffic_type"] = traffic_type + setattr(model_response, "usage", usage) # type: ignore model_response._hidden_params["is_finished"] = False diff --git a/litellm/main.py b/litellm/main.py index bca023e65ec..80a2f74c571 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7383,6 +7383,16 @@ def stream_chunk_builder( # noqa: PLR0915 setattr(response, "usage", usage) + # Propagate provider_specific_fields from the last chunk (contains provider + # metadata like traffic_type set during streaming) + for chunk in reversed(chunks): + hidden = getattr(chunk, "_hidden_params", None) + if hidden and "provider_specific_fields" in hidden: + response._hidden_params.setdefault( + "provider_specific_fields", {} + ).update(hidden["provider_specific_fields"]) + break + # Add cost to usage object if include_cost_in_streaming_usage is True if litellm.include_cost_in_streaming_usage and logging_obj is not None: setattr( diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index 2bff4e23c78..08e31c59662 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1512,6 +1512,12 @@ class LiteLLMCompletionResponsesConfig: user=getattr(chat_completion_response, "user", None), ) responses_api_response._hidden_params = getattr(chat_completion_response, "_hidden_params", {}) + + # Surface provider-specific fields (generic passthrough from any provider) + provider_fields = responses_api_response._hidden_params.get("provider_specific_fields") + if provider_fields: + responses_api_response.provider_specific_fields = provider_fields + return responses_api_response @staticmethod diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 75fd597ffa1..581d1e603dd 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -3338,3 +3338,94 @@ def test_chunk_parser_handles_prompt_feedback_block_with_usage(): assert result.usage.completion_tokens == 0, f"completion_tokens should be 0, got {result.usage.completion_tokens}" assert result.usage.total_tokens == 8175, f"total_tokens should be 8175, got {result.usage.total_tokens}" + +def test_vertex_ai_traffic_type_preserved_in_hidden_params_streaming(): + """Test trafficType is preserved in _hidden_params for streaming.""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + chunk = { + "candidates": [{"content": {"parts": [{"text": "Hello"}]}}], + "usageMetadata": { + "promptTokenCount": 100, + "candidatesTokenCount": 200, + "totalTokenCount": 300, + "trafficType": "ON_DEMAND", + }, + } + + iterator = ModelResponseIterator( + streaming_response=[], sync_stream=True, logging_obj=MagicMock() + ) + result = iterator.chunk_parser(chunk) + + assert result._hidden_params["provider_specific_fields"]["traffic_type"] == "ON_DEMAND" + + +def test_vertex_ai_traffic_type_preserved_in_hidden_params_non_streaming(): + """Test trafficType is preserved in _hidden_params for non-streaming.""" + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + completion_response = { + "candidates": [ + { + "content": {"parts": [{"text": "Hello"}], "role": "model"}, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 50, + "candidatesTokenCount": 100, + "totalTokenCount": 150, + "trafficType": "PROVISIONED_THROUGHPUT", + }, + } + + raw_response = MagicMock() + raw_response.json.return_value = completion_response + + result = VertexGeminiConfig().transform_response( + model="gemini-pro", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert result._hidden_params["provider_specific_fields"]["traffic_type"] == "PROVISIONED_THROUGHPUT" + + +def test_vertex_ai_traffic_type_surfaced_in_responses_api(): + """Test trafficType is surfaced as provider_specific_fields in ResponsesAPIResponse.""" + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + # Create a ModelResponse with provider_specific_fields in _hidden_params + from litellm.types.utils import Choices, Message + + model_response = ModelResponse() + model_response._hidden_params["provider_specific_fields"] = {"traffic_type": "ON_DEMAND"} + model_response.choices = [ + Choices( + message=Message(content="Hello", role="assistant"), + finish_reason="stop", + index=0, + ) + ] + + responses_api_response = LiteLLMCompletionResponsesConfig.transform_chat_completion_response_to_responses_api_response( + request_input="test", + chat_completion_response=model_response, + responses_api_request={}, + ) + + assert responses_api_response.provider_specific_fields["traffic_type"] == "ON_DEMAND" + From 99b4d17ee851187363e6d2f39097d8767d91caac Mon Sep 17 00:00:00 2001 From: The Mavik <179817126+themavik@users.noreply.github.com> Date: Fri, 13 Feb 2026 10:15:30 +0530 Subject: [PATCH 7/9] fix: guard against None litellm_metadata in batch logging (#20832) When litellm_metadata is explicitly set to None in litellm_params, `dict.get("litellm_metadata", {})` returns None (not the default {}), because the key exists. The subsequent .get() call on None raises `AttributeError: 'NoneType' object has no attribute 'get'`. Use `or {}` instead, consistent with line 4924 in the same file. Fixes #15836 --- litellm/litellm_core_utils/litellm_logging.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 82a7af64f97..80a333f52b6 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2331,7 +2331,7 @@ class Logging(LiteLLMLoggingBaseClass): result, LiteLLMBatch ): litellm_params = self.litellm_params or {} - litellm_metadata = litellm_params.get("litellm_metadata", {}) + litellm_metadata = litellm_params.get("litellm_metadata") or {} if ( litellm_metadata.get("batch_ignore_default_logging", False) is True ): # polling job will query these frequently, don't spam db logs From b1a67666ea4ccd77684e0728acbe0a20a85cc6ff Mon Sep 17 00:00:00 2001 From: Otavio Brito <69211663+otaviofbrito@users.noreply.github.com> Date: Fri, 13 Feb 2026 01:53:47 -0300 Subject: [PATCH 8/9] refactor: reuse get_instance_fn in initialize_custom_guardrail - allow module level import (#20917) --- .../proxy/guardrails/guardrail_registry.py | 21 ++++--------------- 1 file changed, 4 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_registry.py b/litellm/proxy/guardrails/guardrail_registry.py index c3da6892209..c0903a35b6d 100644 --- a/litellm/proxy/guardrails/guardrail_registry.py +++ b/litellm/proxy/guardrails/guardrail_registry.py @@ -12,6 +12,7 @@ from litellm._uuid import uuid from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.safe_json_dumps import safe_dumps from litellm.proxy.utils import PrismaClient +from litellm.proxy.types_utils.utils import get_instance_fn from litellm.secret_managers.main import get_secret from litellm.types.guardrails import ( Guardrail, @@ -489,7 +490,7 @@ class InMemoryGuardrailHandler: config_file_path: Optional[str] = None, ) -> Optional[CustomGuardrail]: """ - Initialize a Custom Guardrail from a python file + Initialize a Custom Guardrail from a python file or module path This initializes it by adding it to the litellm callback manager """ @@ -498,26 +499,12 @@ class InMemoryGuardrailHandler: "GuardrailsAIException - Please pass the config_file_path to initialize_guardrails_v2" ) - _file_name, _class_name = guardrail_type.split(".") verbose_proxy_logger.debug( - "Initializing custom guardrail: %s, file_name: %s, class_name: %s", + "Initializing custom guardrail: %s", guardrail_type, - _file_name, - _class_name, ) - directory = os.path.dirname(config_file_path) - module_file_path = os.path.join(directory, _file_name) + ".py" - - spec = importlib.util.spec_from_file_location(_class_name, module_file_path) # type: ignore - if not spec: - raise ImportError( - f"Could not find a module specification for {module_file_path}" - ) - - module = importlib.util.module_from_spec(spec) # type: ignore - spec.loader.exec_module(module) # type: ignore - _guardrail_class = getattr(module, _class_name) + _guardrail_class = get_instance_fn(guardrail_type, config_file_path=config_file_path) mode = litellm_params.mode if mode is None: From 7f6563f1a6735e3aa60d39a2ab6d84704e6a30a3 Mon Sep 17 00:00:00 2001 From: Harshit Jain <48647625+Harshit28j@users.noreply.github.com> Date: Fri, 13 Feb 2026 10:33:11 +0530 Subject: [PATCH 9/9] fix: openai moderation guardrails (#20718) * fix: openai moderation guardrails * adds missing import * mv: test file to right place --- .../guardrail_hooks/openai/moderations.py | 111 +------ .../openai/test_moderations.py | 271 +++++++++++------- .../test_openai_moderation_streaming.py | 172 +++++++++++ 3 files changed, 336 insertions(+), 218 deletions(-) create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py index a196937ef6c..6160fb41439 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py +++ b/litellm/proxy/guardrails/guardrail_hooks/openai/moderations.py @@ -5,14 +5,9 @@ OpenAI Moderation Guardrail Integration for LiteLLM from typing import ( TYPE_CHECKING, - Any, - AsyncGenerator, - Dict, - List, Literal, Optional, Type, - Union, ) from fastapi import HTTPException @@ -20,7 +15,7 @@ from fastapi import HTTPException from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ( CustomGuardrail, - log_guardrail_information, + log_guardrail_information ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( @@ -32,10 +27,8 @@ from litellm.types.utils import GenericGuardrailAPIInputs from .base import OpenAIGuardrailBase if TYPE_CHECKING: - from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import OpenAIModerationResponse from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel - from litellm.types.utils import ModelResponse, ModelResponseStream class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): @@ -236,108 +229,6 @@ class OpenAIModerationGuardrail(OpenAIGuardrailBase, CustomGuardrail): # Moderation doesn't modify content, just blocks - return inputs unchanged return inputs - @log_guardrail_information - async def async_post_call_streaming_iterator_hook( - self, - user_api_key_dict: "UserAPIKeyAuth", - response: Any, - request_data: Dict[str, Any], - ) -> AsyncGenerator["ModelResponseStream", None]: - """ - Process streaming response chunks for OpenAI moderation. - - Collects all chunks from the stream, assembles them into a complete response, - and applies moderation check. If content violates moderation policy, raises HTTPException. - """ - # Import here to avoid circular imports - from litellm.llms.base_llm.base_model_iterator import MockResponseIterator - from litellm.main import stream_chunk_builder - from litellm.types.utils import TextCompletionResponse - - verbose_proxy_logger.debug("OpenAI Moderation: Running streaming response scan") - - # Collect all chunks to process them together - all_chunks: List["ModelResponseStream"] = [] - async for chunk in response: - all_chunks.append(chunk) - - # Assemble the complete response from chunks - assembled_model_response: Optional[ - Union["ModelResponse", TextCompletionResponse] - ] = stream_chunk_builder( - chunks=all_chunks, - ) - - if isinstance(assembled_model_response, (type(None), TextCompletionResponse)): - # If we can't assemble a ModelResponse or it's a text completion, - # just yield the original chunks without moderation - verbose_proxy_logger.warning( - "OpenAI Moderation: Could not assemble ModelResponse from chunks, skipping moderation" - ) - for chunk in all_chunks: - yield chunk - return - - # Extract response text for moderation - response_text = self._extract_response_text(assembled_model_response) - if response_text: - verbose_proxy_logger.debug( - f"OpenAI Moderation: Streaming response text: {response_text[:100]}..." # Log first 100 chars - ) - - # Make moderation request - this will raise HTTPException if content is flagged - moderation_response = await self.async_make_request( - input_text=response_text, - ) - - # Check if content is flagged and raise exception if needed - self._check_moderation_result(moderation_response) - - # If we reach here, content passed moderation - yield the original chunks - mock_response = MockResponseIterator(model_response=assembled_model_response) - - # Return the reconstructed stream - async for chunk in mock_response: - yield chunk - - def _extract_response_text(self, response: "ModelResponse") -> Optional[str]: - """ - Extract text content from the model response for moderation. - """ - if not hasattr(response, "choices") or not response.choices: - return None - - response_texts = [] - for choice in response.choices: - try: - # Try to get content from message (chat completion) - message = getattr(choice, "message", None) - if message: - content = getattr(message, "content", None) - if content and isinstance(content, str): - response_texts.append(content) - continue - - # Try to get text (text completion) - text = getattr(choice, "text", None) - if text and isinstance(text, str): - response_texts.append(text) - continue - - # Try to get content from delta (streaming) - delta = getattr(choice, "delta", None) - if delta: - content = getattr(delta, "content", None) - if content and isinstance(content, str): - response_texts.append(content) - continue - - except (AttributeError, TypeError): - # Skip choices that don't have expected attributes - continue - - return "\n".join(response_texts) if response_texts else None - @staticmethod def get_config_model() -> Optional[Type["GuardrailConfigModel"]]: """ diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index cebba2ff5e1..3a17bbd0025 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -7,7 +7,6 @@ import sys sys.path.insert(0, os.path.abspath("../../../../../..")) -import asyncio from unittest.mock import MagicMock, patch import pytest @@ -26,7 +25,7 @@ async def test_openai_moderation_guardrail_init(): guardrail = OpenAIModerationGuardrail( guardrail_name="test-openai-moderation", ) - + assert guardrail.guardrail_name == "test-openai-moderation" assert guardrail.api_key == "test-key" assert guardrail.model == "omni-moderation-latest" @@ -49,27 +48,27 @@ async def test_openai_moderation_guardrail_adds_to_litellm_callbacks(): # Clear existing callbacks for clean test original_callbacks = litellm.callbacks.copy() litellm.logging_callback_manager._reset_all_callbacks() - + try: with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): guardrail_litellm_params = LitellmParams( guardrail=SupportedGuardrailIntegrations.OPENAI_MODERATION, api_key="test-key", model="omni-moderation-latest", - mode="pre_call" + mode="pre_call", ) guardrail = openai_initialize_guardrail( litellm_params=guardrail_litellm_params, guardrail=Guardrail( guardrail_name="test-openai-moderation", - litellm_params=guardrail_litellm_params - ) + litellm_params=guardrail_litellm_params, + ), ) - + # Check that the guardrail was added to litellm callbacks assert guardrail in litellm.callbacks assert len(litellm.callbacks) == 1 - + # Verify it's the correct guardrail callback = litellm.callbacks[0] assert isinstance(callback, OpenAIModerationGuardrail) @@ -85,12 +84,12 @@ async def test_openai_moderation_guardrail_adds_to_litellm_callbacks(): async def test_openai_moderation_guardrail_safe_content(): """Test OpenAI moderation guardrail with safe content via apply_guardrail""" from litellm.types.utils import GenericGuardrailAPIInputs - + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): guardrail = OpenAIModerationGuardrail( guardrail_name="test-openai-moderation", ) - + # Mock safe moderation response mock_response = OpenAIModerationResponse( id="modr-123", @@ -118,25 +117,29 @@ async def test_openai_moderation_guardrail_safe_content(): "harassment": [], "self-harm": [], "violence": [], - } + }, ) - ] + ], ) - - with patch.object(guardrail, 'async_make_request', return_value=mock_response): + + with patch.object(guardrail, "async_make_request", return_value=mock_response): # Test apply_guardrail with safe content using structured_messages inputs = GenericGuardrailAPIInputs( structured_messages=[ {"role": "user", "content": "Hello, how are you today?"} ] ) - + result = await guardrail.apply_guardrail( inputs=inputs, - request_data={"messages": [{"role": "user", "content": "Hello, how are you today?"}]}, - input_type="request" + request_data={ + "messages": [ + {"role": "user", "content": "Hello, how are you today?"} + ] + }, + input_type="request", ) - + # Should return the original inputs unchanged assert result == inputs @@ -145,12 +148,12 @@ async def test_openai_moderation_guardrail_safe_content(): async def test_openai_moderation_guardrail_apply_guardrail(): """Test OpenAI moderation guardrail apply_guardrail method (unified guardrail interface)""" from litellm.types.utils import GenericGuardrailAPIInputs - + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): guardrail = OpenAIModerationGuardrail( guardrail_name="test-openai-moderation", ) - + # Mock safe moderation response mock_response = OpenAIModerationResponse( id="modr-123", @@ -178,37 +181,37 @@ async def test_openai_moderation_guardrail_apply_guardrail(): "harassment": [], "self-harm": [], "violence": [], - } + }, ) - ] + ], ) - - with patch.object(guardrail, 'async_make_request', return_value=mock_response): + + with patch.object(guardrail, "async_make_request", return_value=mock_response): # Test apply_guardrail with texts (embeddings-style input) inputs = GenericGuardrailAPIInputs( texts=["Hello, how are you?", "What is the weather?"] ) - + result = await guardrail.apply_guardrail( inputs=inputs, request_data={}, input_type="request", ) - + # Should return inputs unchanged (moderation doesn't modify, only blocks) assert result == inputs -@pytest.mark.asyncio +@pytest.mark.asyncio async def test_openai_moderation_guardrail_harmful_content(): """Test OpenAI moderation guardrail with harmful content via apply_guardrail""" from litellm.types.utils import GenericGuardrailAPIInputs - + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): guardrail = OpenAIModerationGuardrail( guardrail_name="test-openai-moderation", ) - + # Mock harmful moderation response mock_response = OpenAIModerationResponse( id="modr-123", @@ -236,40 +239,51 @@ async def test_openai_moderation_guardrail_harmful_content(): "harassment": [], "self-harm": [], "violence": [], - } + }, ) - ] + ], ) - - with patch.object(guardrail, 'async_make_request', return_value=mock_response): + + with patch.object(guardrail, "async_make_request", return_value=mock_response): # Test apply_guardrail with harmful content using structured_messages inputs = GenericGuardrailAPIInputs( structured_messages=[ {"role": "user", "content": "This is hateful content"} ] ) - + # Should raise HTTPException from fastapi import HTTPException + with pytest.raises(HTTPException) as exc_info: await guardrail.apply_guardrail( inputs=inputs, - request_data={"messages": [{"role": "user", "content": "This is hateful content"}]}, - input_type="request" + request_data={ + "messages": [ + {"role": "user", "content": "This is hateful content"} + ] + }, + input_type="request", ) - + assert exc_info.value.status_code == 400 assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) @pytest.mark.asyncio async def test_openai_moderation_guardrail_streaming_safe_content(): - """Test OpenAI moderation guardrail with streaming safe content""" + """Test OpenAI moderation guardrail with streaming safe content via UnifiedLLMGuardrails""" + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): guardrail = OpenAIModerationGuardrail( guardrail_name="test-openai-moderation", + event_hook="post_call", ) - + unified_guardrail = UnifiedLLMGuardrails() + # Mock safe moderation response mock_response = OpenAIModerationResponse( id="modr-123", @@ -297,72 +311,85 @@ async def test_openai_moderation_guardrail_streaming_safe_content(): "harassment": [], "self-harm": [], "violence": [], - } + }, ) - ] + ], ) - + # Mock streaming chunks async def mock_stream(): # Simulate streaming chunks with safe content - chunks = [ - MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello "))]), - MagicMock(choices=[MagicMock(delta=MagicMock(content="world"))]), - MagicMock(choices=[MagicMock(delta=MagicMock(content="!"))]) - ] - for chunk in chunks: + chunk1 = MagicMock() + chunk1.model = "gpt-4" + chunk1.choices = [MagicMock()] + chunk1.choices[0].delta = MagicMock() + chunk1.choices[0].delta.content = "Hello " + chunk1.choices[0].finish_reason = None + + chunk2 = MagicMock() + chunk2.model = "gpt-4" + chunk2.choices = [MagicMock()] + chunk2.choices[0].delta = MagicMock() + chunk2.choices[0].delta.content = "world" + chunk2.choices[0].finish_reason = None + + # Last chunk with finish_reason + chunk3 = MagicMock() + chunk3.model = "gpt-4" + chunk3.choices = [MagicMock()] + chunk3.choices[0].delta = MagicMock() + chunk3.choices[0].delta.content = "!" + chunk3.choices[0].finish_reason = "stop" + + for chunk in [chunk1, chunk2, chunk3]: yield chunk - - # Mock the stream_chunk_builder to return a proper ModelResponse + + # Mock for stream_chunk_builder mock_model_response = MagicMock() - mock_model_response.choices = [ - MagicMock(message=MagicMock(content="Hello world!")) - ] - - with patch.object(guardrail, 'async_make_request', return_value=mock_response), \ - patch('litellm.main.stream_chunk_builder', return_value=mock_model_response), \ - patch('litellm.llms.base_llm.base_model_iterator.MockResponseIterator') as mock_iterator: - - # Mock the iterator to yield the original chunks - async def mock_yield_chunks(): - chunks = [ - MagicMock(choices=[MagicMock(delta=MagicMock(content="Hello "))]), - MagicMock(choices=[MagicMock(delta=MagicMock(content="world"))]), - MagicMock(choices=[MagicMock(delta=MagicMock(content="!"))]) - ] - for chunk in chunks: - yield chunk - - mock_iterator.return_value.__aiter__ = lambda self: mock_yield_chunks() - - user_api_key_dict = UserAPIKeyAuth(api_key="test") + mock_model_response.choices = [MagicMock()] + mock_model_response.choices[0].message = MagicMock() + mock_model_response.choices[0].message.content = "Hello world!" + + with patch.object(guardrail, "async_make_request", return_value=mock_response), patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=mock_model_response, + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) request_data = { - "messages": [ - {"role": "user", "content": "Hello, how are you today?"} - ] + "messages": [{"role": "user", "content": "Hello, how are you today?"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": ["test-openai-moderation"]}, } - - # Test streaming hook with safe content + + # Test streaming hook with safe content via UnifiedLLMGuardrails result_chunks = [] - async for chunk in guardrail.async_post_call_streaming_iterator_hook( + async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, response=mock_stream(), - request_data=request_data + request_data=request_data, ): result_chunks.append(chunk) - + # Should return all chunks without blocking assert len(result_chunks) == 3 @pytest.mark.asyncio async def test_openai_moderation_guardrail_streaming_harmful_content(): - """Test OpenAI moderation guardrail with streaming harmful content""" + """Test OpenAI moderation guardrail with streaming harmful content via UnifiedLLMGuardrails""" + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): guardrail = OpenAIModerationGuardrail( guardrail_name="test-openai-moderation", + event_hook="post_call", ) - + unified_guardrail = UnifiedLLMGuardrails() + # Mock harmful moderation response mock_response = OpenAIModerationResponse( id="modr-123", @@ -390,46 +417,74 @@ async def test_openai_moderation_guardrail_streaming_harmful_content(): "harassment": [], "self-harm": [], "violence": [], - } + }, ) - ] + ], ) - + # Mock streaming chunks with harmful content async def mock_stream(): - chunks = [ - MagicMock(choices=[MagicMock(delta=MagicMock(content="This is "))]), - MagicMock(choices=[MagicMock(delta=MagicMock(content="harmful content"))]) - ] - for chunk in chunks: + # First chunk - no finish_reason + chunk1 = MagicMock() + chunk1.model = "gpt-4" + chunk1.choices = [MagicMock()] + chunk1.choices[0].delta = MagicMock() + chunk1.choices[0].delta.content = "This is " + chunk1.choices[0].finish_reason = None + + # Last chunk - with finish_reason to signal end of stream + chunk2 = MagicMock() + chunk2.model = "gpt-4" + chunk2.choices = [MagicMock()] + chunk2.choices[0].delta = MagicMock() + chunk2.choices[0].delta.content = "harmful content" + chunk2.choices[0].finish_reason = "stop" + + for chunk in [chunk1, chunk2]: yield chunk - - # Mock the stream_chunk_builder to return a ModelResponse with harmful content - mock_model_response = MagicMock() - mock_model_response.choices = [ - MagicMock(message=MagicMock(content="This is harmful content")) - ] - - with patch.object(guardrail, 'async_make_request', return_value=mock_response), \ - patch('litellm.main.stream_chunk_builder', return_value=mock_model_response): - - user_api_key_dict = UserAPIKeyAuth(api_key="test") + + # Mock for stream_chunk_builder - use real litellm types so isinstance checks pass + from litellm.types.utils import ModelResponse + import litellm + mock_model_response = ModelResponse( + id="mock-response", + model="gpt-4", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message( + role="assistant", + content="This is harmful content", + ), + finish_reason="stop", + ) + ], + ) + + with patch.object(guardrail, "async_make_request", return_value=mock_response), patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=mock_model_response, + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) request_data = { - "messages": [ - {"role": "user", "content": "Generate harmful content"} - ] + "messages": [{"role": "user", "content": "Generate harmful content"}], + "guardrail_to_apply": guardrail, + "metadata": {"guardrails": ["test-openai-moderation"]}, } - + # Should raise HTTPException when processing streaming harmful content from fastapi import HTTPException + with pytest.raises(HTTPException) as exc_info: result_chunks = [] - async for chunk in guardrail.async_post_call_streaming_iterator_hook( + async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( user_api_key_dict=user_api_key_dict, response=mock_stream(), - request_data=request_data + request_data=request_data, ): result_chunks.append(chunk) - + assert exc_info.value.status_code == 400 - assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) \ No newline at end of file + assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py new file mode 100644 index 00000000000..c77a5d07b3b --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -0,0 +1,172 @@ +import pytest +from unittest.mock import MagicMock, patch +import os +from litellm.proxy.guardrails.guardrail_hooks.openai.moderations import ( + OpenAIModerationGuardrail, +) +from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, +) +from litellm.types.utils import ModelResponseStream, ModelResponse +from litellm.proxy._types import UserAPIKeyAuth + + +@pytest.mark.asyncio +async def test_openai_moderation_guardrail_streaming_latency(): + """ + Test that the OpenAI Moderation guardrail, when run via UnifiedLLMGuardrails, + supports streaming (fast time-to-first-token) instead of buffering. + """ + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + # 1. Initialize the specific guardrail with proper event_hook + openai_guardrail = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + event_hook="post_call", + ) + + # 2. Initialize the Unified Guardrail system (which invokes the specific guardrail) + unified_guardrail = UnifiedLLMGuardrails() + + # Mock safe moderation response + mock_mod_response = MagicMock() + mock_mod_response.results = [] + + # Mock streaming chunks (no artificial delay - test deterministically) + async def mock_stream(): + chunks_data = ["Hello", " ", "world", "!", " Goodbye"] + for i, content in enumerate(chunks_data): + chunk = MagicMock(spec=ModelResponseStream) + chunk.model = "gpt-4" + choice = MagicMock() + choice.delta = MagicMock() + choice.delta.content = content + # Last chunk gets finish_reason + choice.finish_reason = "stop" if i == len(chunks_data) - 1 else None + chunk.choices = [choice] + yield chunk + + # Mock for stream_chunk_builder to return a simple ModelResponse + mock_model_response = MagicMock(spec=ModelResponse) + mock_model_response.choices = [MagicMock()] + mock_model_response.choices[0].message = MagicMock() + mock_model_response.choices[0].message.content = "Hello world! Goodbye" + + # Patch the network call in the specific guardrail + with patch.object( + openai_guardrail, "async_make_request", return_value=mock_mod_response + ), patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=mock_model_response, + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) + request_data = { + "messages": [{"role": "user", "content": "hi"}], + "guardrail_to_apply": openai_guardrail, + "metadata": { + "guardrails": ["test-openai-moderation"], + "guardrail_config": {"streaming_sampling_rate": 1}, + }, # Check every chunk for test + } + + chunks_received = 0 + first_chunk_yielded = False + + # Call the hook on UnifiedLLMGuardrails + async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + if not first_chunk_yielded: + first_chunk_yielded = True + chunks_received += 1 + + # Deterministic assertions (no flaky timing checks) + assert first_chunk_yielded, "Expected at least one chunk to be yielded" + assert chunks_received == 5, f"Expected 5 chunks, got {chunks_received}" + + +@pytest.mark.asyncio +async def test_openai_moderation_guardrail_streaming_harmful_content(): + """ + Test that harmful content is caught during streaming via UnifiedLLMGuardrails + """ + from fastapi import HTTPException + + with patch.dict(os.environ, {"OPENAI_API_KEY": "test-key"}): + openai_guardrail = OpenAIModerationGuardrail( + guardrail_name="test-openai-moderation", + event_hook="post_call", + ) + unified_guardrail = UnifiedLLMGuardrails() + + # Mock harmful moderation response + mock_mod_response = MagicMock() + mock_mod_response.results = [ + MagicMock( + flagged=True, categories={"hate": True}, category_scores={"hate": 0.99} + ) + ] + + async def mock_stream(): + chunks_data = ["This ", "is ", "harmful ", "content"] + for i, content in enumerate(chunks_data): + chunk = MagicMock(spec=ModelResponseStream) + chunk.model = "gpt-4" + choice = MagicMock() + choice.delta = MagicMock() + choice.delta.content = content + # Last chunk gets finish_reason + choice.finish_reason = "stop" if i == len(chunks_data) - 1 else None + chunk.choices = [choice] + yield chunk + + # Mock for stream_chunk_builder - use real litellm types so isinstance checks pass + import litellm + + mock_model_response = ModelResponse( + id="mock-response", + model="gpt-4", + choices=[ + litellm.Choices( + index=0, + message=litellm.Message( + role="assistant", + content="This is harmful content", + ), + finish_reason="stop", + ) + ], + ) + + with patch.object( + openai_guardrail, "async_make_request", return_value=mock_mod_response + ), patch( + "litellm.llms.openai.chat.guardrail_translation.handler.stream_chunk_builder", + return_value=mock_model_response, + ): + user_api_key_dict = UserAPIKeyAuth( + api_key="test", request_route="/chat/completions" + ) + request_data = { + "messages": [{"role": "user", "content": "generate hate"}], + "guardrail_to_apply": openai_guardrail, + "metadata": { + "guardrails": ["test-openai-moderation"], + "guardrail_config": {"streaming_sampling_rate": 1}, + }, + } + + # Should raise HTTPException + with pytest.raises(HTTPException) as exc_info: + async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + pass + + assert exc_info.value.status_code == 400 + assert "Violated OpenAI moderation policy" in str(exc_info.value.detail)