from collections.abc import Mapping
from datetime import datetime
from enum import Enum
from types import MappingProxyType
from typing import Any, Final, Literal

from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from typing_extensions import Required, TypedDict

from litellm.constants import BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS
from litellm.types.proxy.guardrails.guardrail_hooks.akto import (
    AktoConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.block_code_execution import (
    BlockCodeExecutionGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.cisco_ai_defense import (
    CiscoAIDefenseGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.compresr import (
    CompresrGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.enkryptai import (
    EnkryptAIGuardrailConfigs,
)
from litellm.types.proxy.guardrails.guardrail_hooks.grayswan import (
    GraySwanGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.headroom import (
    HeadroomGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.hiddenlayer import (
    HiddenlayerGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.ibm import (
    IBMGuardrailsBaseConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import (
    ContentFilterCategoryConfig,
)
from litellm.types.proxy.guardrails.guardrail_hooks.ovalix import (
    OvalixGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.promptguard import (
    PromptGuardConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.qohash import (
    QostodianNexusConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.qualifire import (
    QualifireGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.repelloai import (
    RepelloAIGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.singulr import (
    SingulrGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.tool_permission import (
    ToolPermissionGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.vigil_guard import (
    VigilGuardGuardrailConfigModel,
)
from litellm.types.proxy.guardrails.guardrail_hooks.xecguard import (
    XecGuardConfigModel,
)

"""
Pydantic object defining how to set guardrails on litellm proxy

guardrails:
  - guardrail_name: "bedrock-pre-guard"
    litellm_params:
      guardrail: bedrock  # supported values: "akto", "aporia", "bedrock", "lakera", "zscaler_ai_guard"
      mode: "during_call"
      guardrailIdentifier: ff6ujrregl1q
      guardrailVersion: "DRAFT"
      default_on: true
"""


class SupportedGuardrailIntegrations(Enum):
    APORIA = "aporia"
    BEDROCK = "bedrock"
    DYNAMOAI = "dynamoai"
    GUARDRAILS_AI = "guardrails_ai"
    LAKERA = "lakera"
    LAKERA_V2 = "lakera_v2"
    PRESIDIO = "presidio"
    HIDE_SECRETS = "hide-secrets"
    HIDDENLAYER = "hiddenlayer"
    AIM = "aim"
    CATO_NETWORKS = "cato_networks"
    PANGEA = "pangea"
    CROWDSTRIKE_AIDR = "crowdstrike_aidr"
    LASSO = "lasso"
    PILLAR = "pillar"
    GRAYSWAN = "grayswan"
    PANW_PRISMA_AIRS = "panw_prisma_airs"
    CISCO_AI_DEFENSE = "cisco_ai_defense"
    AZURE_PROMPT_SHIELD = "azure/prompt_shield"
    AZURE_TEXT_MODERATIONS = "azure/text_moderations"
    MODEL_ARMOR = "model_armor"
    OPENAI_MODERATION = "openai_moderation"
    NOMA = "noma"
    NOMA_V2 = "noma_v2"
    TOOL_PERMISSION = "tool_permission"
    ZSCALER_AI_GUARD = "zscaler_ai_guard"
    JAVELIN = "javelin"
    ENKRYPTAI = "enkryptai"
    IBM_GUARDRAILS = "ibm_guardrails"
    LITELLM_CONTENT_FILTER = "litellm_content_filter"
    MCP_SECURITY = "mcp_security"
    ONYX = "onyx"
    PROMPTGUARD = "promptguard"
    XECGUARD = "xecguard"
    PROMPT_SECURITY = "prompt_security"
    GENERIC_GUARDRAIL_API = "generic_guardrail_api"
    QUALIFIRE = "qualifire"
    CUSTOM_CODE = "custom_code"
    OVALIX = "ovalix"
    MICROSOFT_PURVIEW = "microsoft_purview"
    SEMANTIC_GUARD = "semantic_guard"
    MCP_END_USER_PERMISSION = "mcp_end_user_permission"
    BLOCK_CODE_EXECUTION = "block_code_execution"
    AKTO = "akto"
    MCP_JWT_SIGNER = "mcp_jwt_signer"
    LLM_AS_A_JUDGE = "llm_as_a_judge"
    DEEPKEEP = "deepkeep"
    QOSTODIAN_NEXUS = "qostodian_nexus"
    RUBRIK = "rubrik"
    VIGIL_GUARD = "vigil_guard"
    REPELLOAI = "repelloai"
    SINGULR = "singulr"
    HEADROOM = "headroom"
    COMPRESR = "compresr"
    STRAIKER = "straiker"
    ALICE = "alice"


class Role(Enum):
    SYSTEM = "system"
    ASSISTANT = "assistant"
    USER = "user"


default_roles: Final = [Role.SYSTEM, Role.ASSISTANT, Role.USER]


class GuardrailItemSpec(TypedDict, total=False):
    callbacks: Required[list[str]]
    default_on: bool
    logging_only: bool | None
    enabled_roles: list[Role] | None
    callback_args: dict[str, dict]


class GuardrailItem(BaseModel):
    callbacks: list[str]
    default_on: bool
    logging_only: bool | None
    guardrail_name: str
    callback_args: dict[str, dict]
    enabled_roles: list[Role] | None

    model_config = ConfigDict(use_enum_values=True)

    def __init__(
        self,
        callbacks: list[str],
        guardrail_name: str,
        default_on: bool = False,
        logging_only: bool | None = None,
        enabled_roles: list[Role] | None = default_roles,
        callback_args: dict[str, dict] | None = None,
    ) -> None:
        super().__init__(
            callbacks=callbacks,
            default_on=default_on,
            logging_only=logging_only,
            guardrail_name=guardrail_name,
            enabled_roles=enabled_roles,
            callback_args=callback_args or {},
        )


# Define the TypedDicts
class LakeraCategoryThresholds(TypedDict, total=False):
    prompt_injection: float
    jailbreak: float


class PiiAction(str, Enum):
    BLOCK = "BLOCK"
    MASK = "MASK"


class PiiEntityCategory(str, Enum):
    GENERAL = "General"
    FINANCE = "Finance"
    USA = "USA"
    UK = "UK"
    SPAIN = "Spain"
    ITALY = "Italy"
    POLAND = "Poland"
    SINGAPORE = "Singapore"
    AUSTRALIA = "Australia"
    INDIA = "India"
    FINLAND = "Finland"


class PiiEntityType(str, Enum):
    # General
    CREDIT_CARD = "CREDIT_CARD"
    CRYPTO = "CRYPTO"
    DATE_TIME = "DATE_TIME"
    EMAIL_ADDRESS = "EMAIL_ADDRESS"
    IBAN_CODE = "IBAN_CODE"
    IP_ADDRESS = "IP_ADDRESS"
    NRP = "NRP"
    LOCATION = "LOCATION"
    PERSON = "PERSON"
    PHONE_NUMBER = "PHONE_NUMBER"
    MEDICAL_LICENSE = "MEDICAL_LICENSE"
    URL = "URL"
    # USA
    US_BANK_NUMBER = "US_BANK_NUMBER"
    US_DRIVER_LICENSE = "US_DRIVER_LICENSE"
    US_ITIN = "US_ITIN"
    US_PASSPORT = "US_PASSPORT"
    US_SSN = "US_SSN"
    # UK
    UK_NHS = "UK_NHS"
    UK_NINO = "UK_NINO"
    UK_PASSPORT = "UK_PASSPORT"
    UK_POSTCODE = "UK_POSTCODE"
    UK_VEHICLE_REGISTRATION = "UK_VEHICLE_REGISTRATION"
    # Spain
    ES_NIF = "ES_NIF"
    ES_NIE = "ES_NIE"
    # Italy
    IT_FISCAL_CODE = "IT_FISCAL_CODE"
    IT_DRIVER_LICENSE = "IT_DRIVER_LICENSE"
    IT_VAT_CODE = "IT_VAT_CODE"
    IT_PASSPORT = "IT_PASSPORT"
    IT_IDENTITY_CARD = "IT_IDENTITY_CARD"
    # Poland
    PL_PESEL = "PL_PESEL"
    # Singapore
    SG_NRIC_FIN = "SG_NRIC_FIN"
    SG_UEN = "SG_UEN"
    # Australia
    AU_ABN = "AU_ABN"
    AU_ACN = "AU_ACN"
    AU_TFN = "AU_TFN"
    AU_MEDICARE = "AU_MEDICARE"
    # India
    IN_PAN = "IN_PAN"
    IN_AADHAAR = "IN_AADHAAR"
    IN_VEHICLE_REGISTRATION = "IN_VEHICLE_REGISTRATION"
    IN_VOTER = "IN_VOTER"
    IN_PASSPORT = "IN_PASSPORT"
    # Finland
    FI_PERSONAL_IDENTITY_CODE = "FI_PERSONAL_IDENTITY_CODE"


# Define mappings of PII entity types by category
PII_ENTITY_CATEGORIES_MAP: Final = {
    PiiEntityCategory.GENERAL: [
        PiiEntityType.DATE_TIME,
        PiiEntityType.EMAIL_ADDRESS,
        PiiEntityType.IP_ADDRESS,
        PiiEntityType.NRP,
        PiiEntityType.LOCATION,
        PiiEntityType.PERSON,
        PiiEntityType.PHONE_NUMBER,
        PiiEntityType.MEDICAL_LICENSE,
        PiiEntityType.URL,
    ],
    PiiEntityCategory.FINANCE: [
        PiiEntityType.CREDIT_CARD,
        PiiEntityType.CRYPTO,
        PiiEntityType.IBAN_CODE,
    ],
    PiiEntityCategory.USA: [
        PiiEntityType.US_BANK_NUMBER,
        PiiEntityType.US_DRIVER_LICENSE,
        PiiEntityType.US_ITIN,
        PiiEntityType.US_PASSPORT,
        PiiEntityType.US_SSN,
    ],
    PiiEntityCategory.UK: [
        PiiEntityType.UK_NHS,
        PiiEntityType.UK_NINO,
        PiiEntityType.UK_PASSPORT,
        PiiEntityType.UK_POSTCODE,
        PiiEntityType.UK_VEHICLE_REGISTRATION,
    ],
    PiiEntityCategory.SPAIN: [PiiEntityType.ES_NIF, PiiEntityType.ES_NIE],
    PiiEntityCategory.ITALY: [
        PiiEntityType.IT_FISCAL_CODE,
        PiiEntityType.IT_DRIVER_LICENSE,
        PiiEntityType.IT_VAT_CODE,
        PiiEntityType.IT_PASSPORT,
        PiiEntityType.IT_IDENTITY_CARD,
    ],
    PiiEntityCategory.POLAND: [PiiEntityType.PL_PESEL],
    PiiEntityCategory.SINGAPORE: [PiiEntityType.SG_NRIC_FIN, PiiEntityType.SG_UEN],
    PiiEntityCategory.AUSTRALIA: [
        PiiEntityType.AU_ABN,
        PiiEntityType.AU_ACN,
        PiiEntityType.AU_TFN,
        PiiEntityType.AU_MEDICARE,
    ],
    PiiEntityCategory.INDIA: [
        PiiEntityType.IN_PAN,
        PiiEntityType.IN_AADHAAR,
        PiiEntityType.IN_VEHICLE_REGISTRATION,
        PiiEntityType.IN_VOTER,
        PiiEntityType.IN_PASSPORT,
    ],
    PiiEntityCategory.FINLAND: [PiiEntityType.FI_PERSONAL_IDENTITY_CODE],
}


class PiiEntityCategoryMap(TypedDict):
    category: str
    entities: list[str]


class GuardrailParamUITypes(str, Enum):
    BOOL = "bool"
    STR = "str"
    MULTISELECT = "multiselect"
    PERCENTAGE = "percentage"


class PresidioPresidioConfigModelUserInterface(BaseModel):
    """Configuration parameters for the Presidio PII masking guardrail on LiteLLM UI"""

    presidio_analyzer_api_base: str | None = Field(
        default=None,
        description="Base URL for the Presidio analyzer API",
    )
    presidio_anonymizer_api_base: str | None = Field(
        default=None,
        description="Base URL for the Presidio anonymizer API",
    )
    presidio_filter_scope: Literal["input", "output", "both"] | None = Field(
        default=None,
        description=(
            "Where to apply Presidio checks: 'input' (user -> model), 'output' (model -> user), or 'both' (default)."
        ),
    )
    output_parse_pii: bool | None = Field(
        default=None,
        description="When True, LiteLLM will replace the masked text with the original text in the response",
        # extra param to let the ui know this is a boolean
        json_schema_extra={"ui_type": GuardrailParamUITypes.BOOL},
    )
    presidio_language: str | None = Field(
        default="en",
        description="Language code for Presidio PII analysis (e.g., 'en', 'de', 'es', 'fr')",
    )
    presidio_run_on: Literal["input", "output", "both"] | None = Field(
        default=None,
        description="Where to apply Presidio checks: input, output, or both (default).",
    )


class PresidioConfigModel(PresidioPresidioConfigModelUserInterface):
    """Configuration parameters for the Presidio PII masking guardrail"""

    pii_entities_config: dict[PiiEntityType | str, PiiAction] | None = Field(
        default=None, description="Configuration for PII entity types and actions"
    )

    presidio_score_thresholds: dict[PiiEntityType | str, float] | None = Field(
        default=None,
        description=(
            "Optional per-entity minimum confidence scores for Presidio detections. "
            "Entities below the threshold are ignored."
        ),
    )
    presidio_entities_deny_list: list[PiiEntityType | str] | None = Field(
        default=None,
        description=(
            "List of entity types to exclude from Presidio detection results. "
            "Detections of these types will be silently dropped. "
            "Useful for suppressing false positives (e.g., US_DRIVER_LICENSE on coding routes)."
        ),
    )
    presidio_ad_hoc_recognizers: str | None = Field(
        default=None,
        description="Path to a JSON file containing ad-hoc recognizers for Presidio",
    )
    presidio_analyze_chunk_size_bytes: int | None = Field(
        default=None,
        description=(
            "Maximum UTF-8 bytes of text sent in a single Presidio /analyze call. "
            "Longer texts are split into overlapping chunks of at most this size "
            "and the merged results are remapped onto the original text. "
            "Defaults to 500000; set it below your analyzer deployment's request "
            "body limit, leaving headroom for the rest of the analyze payload."
        ),
    )
    mock_redacted_text: dict | None = Field(default=None, description="Mock redacted text for testing")


BedrockChecksContentFilterCategory = Literal["VIOLENCE", "HATE", "SEXUAL", "MISCONDUCT", "INSULTS"]
BedrockChecksPromptAttackCategory = Literal["JAILBREAK", "PROMPT_INJECTION", "PROMPT_LEAKAGE"]
BedrockChecksSensitiveInformationEntity = Literal[
    "ADDRESS",
    "AGE",
    "AWS_ACCESS_KEY",
    "AWS_SECRET_KEY",
    "CA_HEALTH_NUMBER",
    "CA_SOCIAL_INSURANCE_NUMBER",
    "CREDIT_DEBIT_CARD_CVV",
    "CREDIT_DEBIT_CARD_EXPIRY",
    "CREDIT_DEBIT_CARD_NUMBER",
    "DRIVER_ID",
    "EMAIL",
    "INTERNATIONAL_BANK_ACCOUNT_NUMBER",
    "IP_ADDRESS",
    "LICENSE_PLATE",
    "MAC_ADDRESS",
    "NAME",
    "PASSWORD",
    "PHONE",
    "PIN",
    "SWIFT_CODE",
    "UK_NATIONAL_HEALTH_SERVICE_NUMBER",
    "UK_NATIONAL_INSURANCE_NUMBER",
    "UK_UNIQUE_TAXPAYER_REFERENCE_NUMBER",
    "URL",
    "USERNAME",
    "US_BANK_ACCOUNT_NUMBER",
    "US_BANK_ROUTING_NUMBER",
    "US_INDIVIDUAL_TAX_IDENTIFICATION_NUMBER",
    "US_PASSPORT_NUMBER",
    "US_SOCIAL_SECURITY_NUMBER",
    "VEHICLE_IDENTIFICATION_NUMBER",
]


class BedrockChecksContentFilterCategoryItem(BaseModel):
    category: BedrockChecksContentFilterCategory


class BedrockChecksContentFilterModel(BaseModel):
    categories: list[BedrockChecksContentFilterCategoryItem]


class BedrockChecksPromptAttackCategoryItem(BaseModel):
    category: BedrockChecksPromptAttackCategory


class BedrockChecksPromptAttackModel(BaseModel):
    categories: list[BedrockChecksPromptAttackCategoryItem]


class BedrockChecksSensitiveInformationEntityItem(BaseModel):
    type: BedrockChecksSensitiveInformationEntity


class BedrockChecksSensitiveInformationModel(BaseModel):
    entities: list[BedrockChecksSensitiveInformationEntityItem]


class BedrockChecksConfigModel(BaseModel):
    """Inline `checks` config for the resource-less Bedrock InvokeGuardrailChecks API.

    Include only the checks you want to run; at least one must be set.
    """

    contentFilter: BedrockChecksContentFilterModel | None = None
    promptAttack: BedrockChecksPromptAttackModel | None = None
    sensitiveInformation: BedrockChecksSensitiveInformationModel | None = None

    @model_validator(mode="after")
    def _require_at_least_one_check(self) -> "BedrockChecksConfigModel":
        if self.contentFilter is None and self.promptAttack is None and self.sensitiveInformation is None:
            raise ValueError(
                "Bedrock 'checks' must enable at least one of: contentFilter, promptAttack, sensitiveInformation."
            )
        return self


class BedrockGuardrailConfigModel(BaseModel):
    """Configuration parameters for the AWS Bedrock guardrail"""

    guardrailIdentifier: str | None = Field(default=None, description="The ID of your guardrail on Bedrock")
    guardrailVersion: str | None = Field(
        default=None,
        description="The version of your Bedrock guardrail (e.g., DRAFT or version number)",
    )
    disable_exception_on_block: bool | None = Field(
        default=False,
        description="If True, will not raise an exception when the guardrail is blocked. Useful for OpenWebUI where exceptions can end the chat flow.",
    )
    aws_region_name: str | None = Field(default=None, description="AWS region where your guardrail is deployed")
    aws_access_key_id: str | None = Field(default=None, description="AWS access key ID for authentication")
    aws_secret_access_key: str | None = Field(default=None, description="AWS secret access key for authentication")
    aws_session_token: str | None = Field(default=None, description="AWS session token for temporary credentials")
    aws_session_name: str | None = Field(default=None, description="Name of the AWS session")
    aws_profile_name: str | None = Field(default=None, description="AWS profile name for credential retrieval")
    aws_role_name: str | None = Field(default=None, description="AWS role name for assuming roles")
    aws_web_identity_token: str | None = Field(default=None, description="Web identity token for AWS role assumption")
    aws_sts_endpoint: str | None = Field(default=None, description="AWS STS endpoint URL")
    aws_external_id: str | None = Field(
        default=None, description="External ID required by the target role's trust policy on sts:AssumeRole"
    )
    aws_bedrock_runtime_endpoint: str | None = Field(default=None, description="AWS Bedrock runtime endpoint URL")
    checks: BedrockChecksConfigModel | None = Field(
        default=None,
        description="Inline safeguards for the resource-less InvokeGuardrailChecks API "
        "(contentFilter / promptAttack / sensitiveInformation). When set, the guardrail "
        "calls InvokeGuardrailChecks instead of ApplyGuardrail and no guardrailIdentifier "
        "is required. Mutually exclusive with guardrailIdentifier.",
    )
    content_filter_threshold: float | None = Field(
        default=0.5,
        ge=0.0,
        le=1.0,
        description="InvokeGuardrailChecks: block when any contentFilter severityScore >= "
        "this value (scores are in [0,1]). Set to null to make the content filter "
        "detect-only (logged, never blocks).",
    )
    prompt_attack_threshold: float | None = Field(
        default=0.5,
        ge=0.0,
        le=1.0,
        description="InvokeGuardrailChecks: block when any promptAttack severityScore >= "
        "this value (scores are in [0,1]). Set to null to make prompt-attack detection detect-only.",
    )
    pii_confidence_threshold: float | None = Field(
        default=0.5,
        ge=0.0,
        le=1.0,
        description="InvokeGuardrailChecks: block when any sensitiveInformation confidenceScore "
        ">= this value (scores are in [0,1]). Set to null to make PII detection detect-only.",
    )
    chunk_budget_chars: int = Field(
        default=BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS,
        gt=0,
        description="ApplyGuardrail: batch size, in characters, used to re-send content after AWS "
        "has rejected a request as too large. Requests AWS accepts are always sent in a single "
        "call, so this has no effect until a rejection happens. Defaults to 25,000; a batch AWS "
        "still rejects is bisected automatically, so this value only trades round trips against "
        "batch size and cannot fail a request on its own.",
    )


class BedrockGuardrailStreamingParams(BaseModel):
    streaming_buffer_until_moderated: bool = Field(
        default=True,
        description="If True (default), withhold every streamed chunk until the end-of-stream "
        "ApplyGuardrail scan passes, so no flagged content reaches the client before a block. "
        "If False, chunks stream through unbuffered, so flagged content can reach the client "
        "before the scan finishes; a flagged scan still ends the stream, with a block message "
        "when disable_exception_on_block is true and an in-stream error frame otherwise.",
    )
    streaming_sampling_rate: int = Field(
        default=5,
        ge=1,
        description="When not buffering and not end-of-stream-only, scan the accumulated response "
        "every Nth streamed chunk. Each sampled scan is a full ApplyGuardrail call that delays "
        "that chunk, so lower values add latency and AWS text-unit cost.",
    )
    streaming_end_of_stream_only: bool = Field(
        default=False,
        description="When not buffering, skip per-chunk sampling and run one ApplyGuardrail scan "
        "on the assembled response at end of stream. Combined with "
        "streaming_buffer_until_moderated=false the full response streams live before the scan "
        "and the scan result lands in guardrail_information; a flagged response still ends the "
        "stream with a block message (disable_exception_on_block=true) or an error frame.",
    )

    @classmethod
    def from_extras(cls, extras: Mapping[str, object] | None) -> "BedrockGuardrailStreamingParams":
        if not extras:
            return cls()
        return cls.model_validate(
            MappingProxyType({name: extras[name] for name in cls.model_fields if extras.get(name) is not None})
        )


class LakeraV2GuardrailConfigModel(BaseModel):
    """Configuration parameters for the Lakera AI v2 guardrail"""

    api_key: str | None = Field(default=None, description="API key for the Lakera AI service")
    api_base: str | None = Field(default=None, description="Base URL for the Lakera AI API")
    project_id: str | None = Field(default=None, description="Project ID for the Lakera AI project")
    payload: bool | None = Field(default=True, description="Whether to include payload in the response")
    breakdown: bool | None = Field(default=True, description="Whether to include breakdown in the response")
    metadata: dict | None = Field(default=None, description="Additional metadata to include in the request")
    dev_info: bool | None = Field(
        default=True,
        description="Whether to include developer information in the response",
    )
    on_flagged: Literal["block", "monitor", "inject_system_message"] | None = Field(
        default="block",
        description="Action to take when content is flagged: 'block' (raise exception), 'monitor' (log only), "
        "or 'inject_system_message' (append an advisory system message and let the LLM decide)",
    )
    advisory_system_message: str | None = Field(
        default=None,
        description="Custom advisory message template used when on_flagged='inject_system_message'. "
        "Must contain a {reason} placeholder. Defaults to a generic advisory message if unset.",
    )


class LassoGuardrailConfigModel(BaseModel):
    """Configuration parameters for the Lasso guardrail"""

    lasso_user_id: str | None = Field(default=None, description="User ID for the Lasso guardrail")
    lasso_conversation_id: str | None = Field(default=None, description="Conversation ID for the Lasso guardrail")
    mask: bool | None = Field(default=False, description="Enable content masking using Lasso classifix API")


class DeepKeepGuardrailConfigModel(BaseModel):
    """Configuration parameters for the DeepKeep AI Firewall guardrail"""

    deepkeep_firewall_id: str | None = Field(
        default=None,
        description=(
            "The DeepKeep Firewall ID to use for guardrail evaluation. "
            "If not provided, the `DEEPKEEP_FIREWALL_ID` environment variable is checked."
        ),
    )


class PillarGuardrailConfigModel(BaseModel):
    """Configuration parameters for the Pillar Security guardrail"""

    on_flagged_action: str | None = Field(
        default="monitor",
        description="Action to take when content is flagged: 'block' (raise exception) or 'monitor' (log only)",
    )
    async_mode: bool | None = Field(
        default=None,
        description="Set to True to request asynchronous analysis (sets `plr_async` header). Defaults to provider behaviour when omitted.",
    )
    persist_session: bool | None = Field(
        default=None,
        description="Controls Pillar session persistence (sets `plr_persist` header). Set to False to disable persistence.",
    )
    include_scanners: bool | None = Field(
        default=True,
        description="Include scanner category summaries in responses (sets `plr_scanners` header).",
    )
    include_evidence: bool | None = Field(
        default=True,
        description="Include detailed evidence payloads in responses (sets `plr_evidence` header).",
    )


class NomaGuardrailConfigModel(BaseModel):
    """Configuration parameters for the Noma Security guardrail"""

    use_v2: bool | None = Field(
        default=False,
        description="If True and guardrail='noma', route to the new Noma v2 implementation instead of the legacy implementation.",
    )
    application_id: str | None = Field(
        default=None,
        description="Application ID for Noma Security. Defaults to 'litellm' if not provided",
    )
    monitor_mode: bool | None = Field(
        default=None,
        description="If True, logs violations without blocking. Defaults to False if not provided",
    )
    block_failures: bool | None = Field(
        default=None,
        description="If True, blocks requests on API failures. Defaults to True if not provided",
    )
    anonymize_input: bool | None = Field(
        default=None,
        description="If True, replaces sensitive content with anonymized version when only PII/PCI/secrets are detected. Only applies in blocking mode. Defaults to False if not provided",
    )


class ZscalerAIGuardConfigModel(BaseModel):
    """Configuration parameters for the Zscaler AI Guard guardrail"""

    policy_id: int | None = Field(
        default=None,
        description="Policy ID for Zscaler AI Guard. Can also be set via ZSCALER_AI_GUARD_POLICY_ID environment variable",
    )
    send_user_api_key_alias: bool | None = Field(
        default=False, description="Whether to send user_API_key_alias in headers"
    )
    send_user_api_key_user_id: bool | None = Field(
        default=False, description="Whether to send user_API_key_user_id in headers"
    )
    send_user_api_key_team_id: bool | None = Field(
        default=False, description="Whether to send user_API_key_team_id in headers"
    )


class JavelinGuardrailConfigModel(BaseModel):
    """Configuration parameters for the Javelin guardrail"""

    guard_name: str | None = Field(default=None, description="Name of the Javelin guard to use")
    api_version: str | None = Field(default="v1", description="API version for Javelin service")
    metadata: dict | None = Field(default=None, description="Additional metadata to send with requests")
    application: str | None = Field(default=None, description="Application name for Javelin service")
    config: dict | None = Field(default=None, description="Additional configuration for the guardrail")


class ContentFilterAction(str, Enum):
    """Action to take when content filter detects a match"""

    BLOCK = "BLOCK"
    MASK = "MASK"


class BlockedWord(BaseModel):
    """Represents a blocked word with its action and optional description"""

    keyword: str = Field(description="The keyword to block or mask")
    action: ContentFilterAction = Field(description="Action to take when keyword is detected (BLOCK or MASK)")
    description: str | None = Field(
        default=None,
        description="Optional description explaining why this keyword is sensitive",
    )


class ContentFilterPattern(BaseModel):
    """Represents a content filter pattern (prebuilt or custom regex)"""

    pattern_type: Literal["prebuilt", "regex"] = Field(
        description="Type of pattern: 'prebuilt' for predefined patterns or 'regex' for custom"
    )
    pattern_name: str | None = Field(
        default=None,
        description="Name of prebuilt pattern (e.g., 'us_ssn', 'credit_card'). Required if pattern_type is 'prebuilt'",
    )
    pattern: str | None = Field(
        default=None,
        description="Custom regex pattern. Required if pattern_type is 'regex'",
    )
    name: str | None = Field(
        default=None,
        description="Name for this pattern (used in logging and error messages)",
    )
    action: ContentFilterAction = Field(description="Action to take when pattern matches (BLOCK or MASK)")


class ContentFilterConfigModel(BaseModel):
    """Configuration parameters for the content filter guardrail"""

    patterns: list[ContentFilterPattern] | None = Field(
        default=None,
        description="List of patterns (prebuilt or custom regex) to detect",
    )
    blocked_words: list[BlockedWord] | None = Field(
        default=None, description="List of blocked words with individual actions"
    )
    blocked_words_file: str | None = Field(default=None, description="Path to YAML file containing blocked_words list")
    categories: list[ContentFilterCategoryConfig] | None = Field(
        default=None,
        description="List of prebuilt categories to enable (harmful_*, bias_*)",
    )
    severity_threshold: str | None = Field(
        default=None,
        description="Minimum severity to block (high, medium, low)",
    )
    pattern_redaction_format: str | None = Field(
        default=None,
        description="Format string for pattern redaction (use {pattern_name} placeholder)",
    )
    keyword_redaction_tag: str | None = Field(
        default=None,
        description="Tag to use for keyword redaction",
    )


class BaseLitellmParams(ContentFilterConfigModel):  # works for new and patch update guardrails
    api_key: str | None = Field(default=None, description="API key for the guardrail service")
    api_base: str | None = Field(default=None, description="Base URL for the guardrail service API")

    experimental_use_latest_role_message_only: bool | None = Field(
        default=False,
        description="When True, guardrails only receive the latest message for the relevant role (e.g., newest user input pre-call, newest assistant output post-call)",
    )

    only_scan_new_messages: bool | None = Field(
        default=False,
        description=(
            "When True, the guardrail only scans messages that have not already been scanned "
            "earlier in the same session (identified by litellm_session_id / session_id). "
            "Message content is hashed per session and cached; only the diff (new or edited "
            "messages) is sent to the guardrail provider on follow-up calls. Falls back to a "
            "full scan when the request has no session id or the cache is unavailable. Intended "
            "for blocking/detection guardrails; not applied when mask_request_content is set."
        ),
    )

    skip_system_message_in_guardrail: bool | None = Field(
        default=None,
        description=(
            "When True, unified guardrails skip system-role messages when building "
            "evaluation inputs (texts and structured_messages). When False, system "
            "messages are included even if litellm_settings sets a global skip. When "
            "None, use the global litellm.skip_system_message_in_guardrail setting. "
            "For Anthropic /v1/messages, the flag applies only to the trusted top-level "
            "system prompt. In-sequence system entries are untrusted client input and remain "
            "in texts and structured_messages."
        ),
    )

    skip_tool_message_in_guardrail: bool | None = Field(
        default=None,
        description=(
            "When True, unified guardrails skip tool-role messages when building "
            "evaluation inputs (texts and structured_messages). When False, tool "
            "messages are included even if litellm_settings sets a global skip. When "
            "None, use the global litellm.skip_tool_message_in_guardrail setting."
        ),
    )

    scan_only_tool_results: bool | None = Field(
        default=None,
        description=(
            "When True, unified guardrails only evaluate tool results, the untrusted data an "
            "agent feeds back into the model, and skip system, user, and assistant content. "
            "Intended for agent harnesses whose own prompt scaffolding is trusted but often "
            "trips prompt-attack detectors."
        ),
    )

    # Lakera specific params
    category_thresholds: LakeraCategoryThresholds | None = Field(
        default=None,
        description="Threshold configuration for Lakera guardrail categories",
    )

    # hide secrets params
    detect_secrets_config: dict | None = Field(default=None, description="Configuration for detect-secrets guardrail")

    # guardrails ai params
    guard_name: str | None = Field(default=None, description="Name of the guardrail in guardrails.ai")
    default_on: bool | None = Field(default=None, description="Whether the guardrail is enabled by default")

    ################## PII control params #################
    ########################################################
    mask_request_content: bool | None = Field(
        default=None,
        description="Will mask request content if guardrail makes any changes",
    )
    mask_response_content: bool | None = Field(
        default=None,
        description="Will mask response content if guardrail makes any changes",
    )

    # pangea params
    pangea_input_recipe: str | None = Field(default=None, description="Recipe for input (LLM request)")

    pangea_output_recipe: str | None = Field(default=None, description="Recipe for output (LLM response)")

    model: str | None = Field(
        default=None,
        description="Optional field if guardrail requires a 'model' parameter",
    )

    violation_message_template: str | None = Field(
        default=None,
        description="Custom message when a guardrail blocks an action. Supports placeholders like {tool_name}, {rule_id}, and {default_message}.",
    )

    ################## Realtime API params ################
    ########################################################
    end_session_after_n_fails: int | None = Field(
        default=None,
        description="For /v1/realtime sessions: automatically close the session after this many guardrail violations.",
    )
    on_violation: Literal["warn", "end_session"] | None = Field(
        default=None,
        description="For /v1/realtime sessions: 'warn' speaks the violation message and continues; 'end_session' speaks the message and closes the connection.",
    )
    realtime_violation_message: str | None = Field(
        default=None,
        description="The message the bot speaks aloud when a /v1/realtime guardrail fires. Falls back to violation_message_template if not set.",
    )

    # Model Armor params
    template_id: str | None = Field(default=None, description="The ID of your Model Armor template")
    location: str | None = Field(default=None, description="Google Cloud location/region (e.g., us-central1)")
    credentials: str | None = Field(
        default=None,
        description="Path to Google Cloud credentials JSON file or JSON string",
    )
    api_endpoint: str | None = Field(default=None, description="Optional custom API endpoint for Model Armor")
    fail_on_error: bool | None = Field(
        default=True,
        description=(
            "Whether to fail the request if the guardrail encounters an error. "
            "Implemented by guardrail='model_armor', 'generic_guardrail_api' and 'crowdstrike_aidr'. "
            "True (default) raises the error. False logs a critical error and lets the request proceed, "
            "so only a valid guardrail response can block or modify it."
        ),
    )
    skip_unscannable_attachments: bool | None = Field(
        default=False,
        description=(
            "Implemented by guardrail='model_armor'. When True, attachment references that carry no "
            "inline bytes (file_id, gs://, or http(s) URLs) pass through unscanned instead of blocking, "
            "while fail_on_error still governs real Model Armor API errors. Default False blocks them."
        ),
    )
    sanitize_error_detail: bool | None = Field(
        default=True,
        description=(
            "For guardrail='model_armor': omit the raw Model Armor response from "
            "caller-facing errors and logs by default. Set False to restore verbose output."
        ),
    )

    additional_provider_specific_params: dict[str, Any] | None = Field(
        default=None,
        description="Additional provider-specific parameters for generic guardrail APIs",
    )

    unreachable_fallback: Literal["fail_closed", "fail_open"] = Field(
        default="fail_closed",
        description=(
            "Behavior when a guardrail endpoint is unreachable due to network errors. "
            "Implemented by guardrail='generic_guardrail_api', 'akto', 'vigil_guard', 'repelloai', 'headroom', and 'compresr'. "
            "'fail_closed' raises an error (default). 'fail_open' logs a critical error and allows the request to proceed."
        ),
    )

    extra_headers: list[str] | None = Field(
        default=None,
        description=(
            "Header names to forward from the client request to the guardrail (e.g. x-request-id). "
            "Only these headers' values are sent; others may be omitted or sent as [present]. "
            "Used by generic_guardrail_api (similar to MCP extra_headers)."
        ),
    )

    # Custom code guardrail params
    custom_code: str | None = Field(
        default=None,
        description="Python-like code containing the apply_guardrail function for custom guardrail logic",
    )

    timeout: float | None = Field(
        default=None,
        description=(
            "Per-request timeout for the guardrail provider API call (seconds). "
            "Accepts int, float, or numeric string; coerced to float on load. "
            "Each guardrail handler chooses its own default when unset."
        ),
    )

    on_sensitive_data: Literal["block", "route"] | None = Field(
        default=None,
        description=(
            "Action to take when sensitive data is detected. "
            "'block' raises an exception (default behavior). "
            "'route' reroutes the request to the model specified in sensitive_data_route_to_model."
        ),
    )

    sensitive_data_route_to_model: str | None = Field(
        default=None,
        description=(
            "Model to route requests to when sensitive data is detected and on_sensitive_data='route'. "
            "This is typically an on-premise model for data privacy. "
            "The routing decision persists for the entire session."
        ),
    )

    sticky_session_routing: bool | None = Field(
        default=True,
        description=(
            "When True (default), after sensitive data is detected and routed, all subsequent "
            "requests in the same session will continue routing to the same model."
        ),
    )

    run_in_parallel: bool | None = Field(
        default=None,
        description=(
            "When True, this pre_call or post_call guardrail runs concurrently with other opted-in "
            "guardrails of the same hook, after the sequential guardrails have run. Use only for "
            "block-only guardrails that inspect and reject; do not enable it for guardrails that "
            "modify the request or response (e.g. PII masking or sensitive-data routing), since "
            "parallel runs share one snapshot and their mutations would race."
        ),
    )

    scan_raw_request: bool | None = Field(
        default=None,
        description=(
            "When True, this pre_call guardrail always evaluates the request as it was before any "
            "guardrail in this hook ran, regardless of its position in the guardrails list -- so the "
            "YAML order of guardrails can never change whether this one blocks. Use only for "
            "block-only guardrails: any data this guardrail returns is discarded, same contract as "
            "run_in_parallel, since an earlier guardrail's masking must not be undone by this one."
        ),
    )

    @field_validator(
        "mode",
        "default_action",
        "on_disallowed_action",
        "unreachable_fallback",
        "on_sensitive_data",
        mode="before",
        check_fields=False,
    )
    @classmethod
    def normalize_lowercase(cls, v):
        """Normalize string and list fields to lowercase for ALL guardrail types."""
        if isinstance(v, str):
            return v.lower()
        if isinstance(v, list):
            return [x.lower() if isinstance(x, str) else x for x in v]
        return v

    @model_validator(mode="after")
    def validate_sensitive_data_routing(self) -> "BaseLitellmParams":
        if self.on_sensitive_data == "route" and not self.sensitive_data_route_to_model:
            raise ValueError("sensitive_data_route_to_model must be set when on_sensitive_data='route'")
        return self

    model_config = ConfigDict(extra="allow", protected_namespaces=())


class Mode(BaseModel):
    tags: dict[str, str | list[str]] = Field(description="Tags for the guardrail mode")
    default: str | list[str] | None = Field(default=None, description="Default mode when no tags match")


class LitellmParams(  # pyright: ignore[reportIncompatibleVariableOverride]  # on_flagged literal diverges across mixins
    CiscoAIDefenseGuardrailConfigModel,
    PresidioConfigModel,
    BedrockGuardrailConfigModel,
    LakeraV2GuardrailConfigModel,
    HeadroomGuardrailConfigModel,
    CompresrGuardrailConfigModel,
    RepelloAIGuardrailConfigModel,
    LassoGuardrailConfigModel,
    DeepKeepGuardrailConfigModel,
    PillarGuardrailConfigModel,
    GraySwanGuardrailConfigModel,
    NomaGuardrailConfigModel,
    PromptGuardConfigModel,
    XecGuardConfigModel,
    ToolPermissionGuardrailConfigModel,
    ZscalerAIGuardConfigModel,
    AktoConfigModel,
    JavelinGuardrailConfigModel,
    BaseLitellmParams,
    EnkryptAIGuardrailConfigs,
    IBMGuardrailsBaseConfigModel,
    OvalixGuardrailConfigModel,
    QualifireGuardrailConfigModel,
    BlockCodeExecutionGuardrailConfigModel,
    HiddenlayerGuardrailConfigModel,
    QostodianNexusConfigModel,
    VigilGuardGuardrailConfigModel,
    SingulrGuardrailConfigModel,
):
    guardrail: str = Field(description="The type of guardrail integration to use")
    mode: str | list[str] | Mode = Field(
        description="When to apply the guardrail (pre_call, post_call, during_call, logging_only)"
    )

    @field_validator("timeout", mode="before", check_fields=False)
    @classmethod
    def coerce_timeout(cls, v):
        """Accept string-valued timeouts (dashboard UI sends JSON strings)
        and coerce to float before any handler reads the value."""
        if v is None or v == "":
            return None
        try:
            return float(v)
        except (TypeError, ValueError) as e:
            raise ValueError(f"timeout must be numeric, got {v!r}") from e

    def __init__(self, **kwargs) -> None:
        default_on: Final = kwargs.pop("default_on", None)
        if default_on is not None:
            kwargs["default_on"] = default_on
        else:
            kwargs["default_on"] = False

        super().__init__(**kwargs)

    def __contains__(self, key) -> bool:
        # Define custom behavior for the 'in' operator
        return hasattr(self, key)

    def get(self, key, default=None):
        # Custom .get() method to access attributes with a default value if the attribute doesn't exist
        return getattr(self, key, default)

    def __getitem__(self, key):
        # Allow dictionary-style access to attributes
        return getattr(self, key)


class Guardrail(TypedDict, total=False):
    guardrail_id: str | None
    guardrail_name: Required[str]
    litellm_params: Required[LitellmParams]
    guardrail_info: dict | None
    policy_template: str | None
    created_at: datetime | None
    updated_at: datetime | None


class guardrailConfig(TypedDict):
    guardrails: list[Guardrail]


class GuardrailEventHooks(str, Enum):
    pre_call = "pre_call"
    post_call = "post_call"
    during_call = "during_call"
    logging_only = "logging_only"
    pre_mcp_call = "pre_mcp_call"
    during_mcp_call = "during_mcp_call"
    post_mcp_call = "post_mcp_call"
    realtime_input_transcription = "realtime_input_transcription"


class DynamicGuardrailParams(TypedDict):
    extra_body: dict[str, Any]


class GUARDRAIL_DEFINITION_LOCATION(str, Enum):
    DB = "db"
    CONFIG = "config"


class GuardrailInfoResponse(BaseModel):
    guardrail_id: str | None = None
    guardrail_name: str
    litellm_params: BaseLitellmParams | None = None
    guardrail_info: dict | None = None
    created_at: datetime | None = None
    updated_at: datetime | None = None
    guardrail_definition_location: GUARDRAIL_DEFINITION_LOCATION = GUARDRAIL_DEFINITION_LOCATION.CONFIG

    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)


class ListGuardrailsResponse(BaseModel):
    guardrails: list[GuardrailInfoResponse]


class GuardrailUIAddGuardrailSettings(BaseModel):
    supported_entities: list[str]
    supported_actions: list[str]
    supported_modes: list[str]
    supported_modes_by_provider: dict[str, list[str]]
    pii_entity_categories: list[PiiEntityCategoryMap]
    content_filter_settings: dict[str, Any] | None = None


class PresidioPerRequestConfig(BaseModel):
    """
    presdio params that can be controlled per request, api key
    """

    language: str | None = None
    entities: list[PiiEntityType] | None = None


class ApplyGuardrailRequest(BaseModel):
    guardrail_name: str
    text: str
    language: str | None = None
    entities: list[PiiEntityType] | None = None
    input_type: str = "request"
    messages: list[dict[str, Any]] | None = None
    metadata: dict[str, Any] | None = None


class ApplyGuardrailResponse(BaseModel):
    response_text: str


class PatchGuardrailRequest(BaseModel):
    guardrail_name: str | None = None
    litellm_params: BaseLitellmParams | None = None
    guardrail_info: dict[str, Any] | None = None
