From cd3415558febd01eacd0ac5b9fe4722097c898f1 Mon Sep 17 00:00:00 2001 From: Alexis Date: Fri, 16 Oct 2026 23:46:38 +0200 Subject: [PATCH] fix(guardrails): type hook call_type with CallTypesLiteral --- .../guardrail_hooks/aporia_ai/aporia_ai.py | 14 ++----- .../cisco_ai_defense/cisco_ai_defense.py | 27 ++----------- .../guardrail_hooks/custom_guardrail.py | 27 ++----------- .../guardrails_ai/guardrails_ai.py | 13 +------ .../guardrails/guardrail_hooks/lakera_ai.py | 39 ++----------------- .../guardrails/guardrail_hooks/lasso/lasso.py | 25 ++---------- .../guardrail_hooks/pillar/pillar.py | 28 ++----------- 7 files changed, 24 insertions(+), 149 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py index c4615a3cebf..053042ecdfa 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/aporia_ai/aporia_ai.py @@ -11,7 +11,7 @@ import sys sys.path.insert(0, os.path.abspath("../..")) # Adds the parent directory to the system path import json import sys -from typing import TYPE_CHECKING, Final, Literal +from typing import TYPE_CHECKING, Final from fastapi import HTTPException @@ -29,6 +29,7 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CallTypesLiteral GUARDRAIL_NAME: Final = "aporia" @@ -172,16 +173,7 @@ class AporiaGuardrail(CustomGuardrail): self, data: dict, user_api_key_dict: UserAPIKeyAuth, - call_type: Literal[ - "completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - "responses", - "mcp_call", - "anthropic_messages", - ], + call_type: CallTypesLiteral, ): from litellm.proxy.common_utils.callback_utils import ( add_guardrail_to_applied_guardrails_header, diff --git a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py index ab43d32e260..134977790ee 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py +++ b/litellm/proxy/guardrails/guardrail_hooks/cisco_ai_defense/cisco_ai_defense.py @@ -22,7 +22,7 @@ import os from collections.abc import AsyncIterator, Mapping, Sequence from dataclasses import dataclass, replace from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final import httpx from fastapi import HTTPException @@ -45,6 +45,7 @@ from litellm.proxy.common_utils.callback_utils import ( ) from litellm.types.guardrails import GuardrailEventHooks from litellm.types.utils import ( + CallTypesLiteral, Choices, LLMResponseTypes, ModelResponse, @@ -299,18 +300,7 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, - call_type: Literal[ - "completion", - "text_completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - "pass_through_endpoint", - "rerank", - "mcp_call", - "anthropic_messages", - ], + call_type: CallTypesLiteral, ) -> Exception | str | dict | None: # Trust proxy call_type, not caller-controlled request shape. is_mcp: Final = self._is_mcp_call_type(call_type) @@ -350,16 +340,7 @@ class CiscoAIDefenseGuardrail(CiscoAIDefenseMcpMixin, CustomGuardrail): self, data: dict, user_api_key_dict: UserAPIKeyAuth, - call_type: Literal[ - "completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - "responses", - "mcp_call", - "anthropic_messages", - ], + call_type: CallTypesLiteral, ) -> Exception | str | dict | None: is_mcp: Final = self._is_mcp_call_type(call_type) diff --git a/litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py index a9caec7f608..1348b90d13f 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/custom_guardrail.py @@ -1,4 +1,4 @@ -from typing import Final, Literal +from typing import Final import litellm from litellm._logging import verbose_proxy_logger @@ -8,6 +8,7 @@ from litellm.integrations.custom_guardrail import ( log_guardrail_information, ) from litellm.proxy._types import UserAPIKeyAuth +from litellm.types.utils import CallTypesLiteral class myCustomGuardrail(CustomGuardrail): @@ -26,18 +27,7 @@ class myCustomGuardrail(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, - call_type: Literal[ - "completion", - "text_completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - "pass_through_endpoint", - "rerank", - "mcp_call", - "anthropic_messages", - ], + call_type: CallTypesLiteral, ) -> Exception | str | dict | None: """ Runs before the LLM API call @@ -64,16 +54,7 @@ class myCustomGuardrail(CustomGuardrail): self, data: dict, user_api_key_dict: UserAPIKeyAuth, - call_type: Literal[ - "completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - "responses", - "mcp_call", - "anthropic_messages", - ], + call_type: CallTypesLiteral, ): """ Runs in parallel to LLM API call diff --git a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py index cf6592a3e58..57a7c43c9a3 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/guardrails_ai/guardrails_ai.py @@ -22,6 +22,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CallTypesLiteral if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -167,17 +168,7 @@ class GuardrailsAI(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, cache: litellm.DualCache, data: dict, - call_type: Literal[ - "completion", - "text_completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - "pass_through_endpoint", - "rerank", - "mcp_call", - ], + call_type: CallTypesLiteral, ) -> ( Exception | str | dict | None ): # raise exception if invalid, return a str for the user to receive - if rejected, or return a modified dictionary for passing into litellm diff --git a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py index cb1b223fb47..a376e9a8f46 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lakera_ai.py @@ -36,6 +36,7 @@ from litellm.types.guardrails import ( Role, default_roles, ) +from litellm.types.utils import CallTypesLiteral GUARDRAIL_NAME: Final = "lakera_prompt_injection" @@ -116,19 +117,7 @@ class lakeraAI_Moderation(CustomGuardrail): self, data: dict, user_api_key_dict: UserAPIKeyAuth, - call_type: Literal[ - "completion", - "text_completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - "pass_through_endpoint", - "rerank", - "responses", - "mcp_call", - "anthropic_messages", - ], + call_type: CallTypesLiteral, ): if ( await should_proceed_based_on_metadata( @@ -290,18 +279,7 @@ class lakeraAI_Moderation(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, cache: litellm.DualCache, data: dict, - call_type: Literal[ - "completion", - "text_completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - "pass_through_endpoint", - "rerank", - "mcp_call", - "anthropic_messages", - ], + call_type: CallTypesLiteral, ) -> Exception | str | dict | None: from litellm.types.guardrails import GuardrailEventHooks @@ -321,16 +299,7 @@ class lakeraAI_Moderation(CustomGuardrail): self, data: dict, user_api_key_dict: UserAPIKeyAuth, - call_type: Literal[ - "completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - "responses", - "mcp_call", - "anthropic_messages", - ], + call_type: CallTypesLiteral, ): if self.event_hook is None: if self.moderation_check == "pre_call": diff --git a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py index e05db5003fa..9e864db7343 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py +++ b/litellm/proxy/guardrails/guardrail_hooks/lasso/lasso.py @@ -49,6 +49,7 @@ from litellm.proxy.guardrails._content_utils import ( has_non_string_content, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.utils import CallTypesLiteral class LassoResponse(TypedDict): @@ -177,18 +178,7 @@ class LassoGuardrail(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, cache: DualCache, # Deprecated, use global_cache instead (kept to align with CustomGuardrail interface) data: dict, - call_type: Literal[ - "completion", - "text_completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - "pass_through_endpoint", - "rerank", - "mcp_call", - "anthropic_messages", - ], + call_type: CallTypesLiteral, ) -> Exception | str | dict | None: """ Runs before the LLM API call to validate and potentially modify input. @@ -210,16 +200,7 @@ class LassoGuardrail(CustomGuardrail): self, data: dict, user_api_key_dict: UserAPIKeyAuth, - call_type: Literal[ - "completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - "responses", - "mcp_call", - "anthropic_messages", - ], + call_type: CallTypesLiteral, cache: DualCache, ): """ diff --git a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py index 3f025cfc8b3..1bda08bcb40 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py +++ b/litellm/proxy/guardrails/guardrail_hooks/pillar/pillar.py @@ -8,7 +8,7 @@ # Standard library imports import json import os -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol +from typing import TYPE_CHECKING, Any, Final, Protocol from urllib.parse import quote import httpx @@ -37,7 +37,7 @@ from litellm.proxy.common_utils.callback_utils import ( get_metadata_variable_name_from_kwargs, ) from litellm.types.guardrails import GuardrailEventHooks -from litellm.types.utils import LLMResponseTypes +from litellm.types.utils import CallTypesLiteral, LLMResponseTypes if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel @@ -315,18 +315,7 @@ class PillarGuardrail(CustomGuardrail): user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, - call_type: Literal[ - "completion", - "text_completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - "pass_through_endpoint", - "rerank", - "mcp_call", - "anthropic_messages", - ], + call_type: CallTypesLiteral, ) -> Exception | str | dict | None: """ Pre-call hook to scan the request for security threats before sending to LLM. @@ -361,16 +350,7 @@ class PillarGuardrail(CustomGuardrail): self, data: dict, user_api_key_dict: UserAPIKeyAuth, - call_type: Literal[ - "completion", - "embeddings", - "image_generation", - "moderation", - "audio_transcription", - "responses", - "mcp_call", - "anthropic_messages", - ], + call_type: CallTypesLiteral, ) -> Exception | str | dict | None: """ During-call hook to scan the request in parallel with LLM processing.