schema/models.py

from typing import Literal

from pydantic import BaseModel, ConfigDict, Field


CategoryName = Literal[
    "pii",
    "prompt_injection",
    "toxicity",
    "bias",
    "sentiment",
    "topic_ban",
    "competitor",
    "malicious_url",
    "secret",
    "code",
    "language_violation",
    "no_refusal",
    "factual",
    "relevance",
    "substring",
    "gibberish",
]

Severity = Literal["safe", "low", "medium", "high", "critical"]
SupportedLanguage = Literal["en", "hi", "ko", "ja", "ar", "mixed", "unknown"]
ScriptName = Literal[
    "latin",
    "devanagari",
    "hangul",
    "japanese",
    "arabic",
    "mixed",
    "other",
]
ScenarioName = Literal[
    "general",
    "finance",
    "banking",
    "insurance",
    "healthcare",
    "retail",
    "telecom",
    "government",
    "developer",
    "customer_support",
]


class Match(BaseModel):
    model_config = ConfigDict(extra="forbid")

    text: str = Field(
        description="Verbatim evidence copied from the source-language input or response."
    )
    kind: str = Field(default="", description="English subtype label for the match.")


class Category(BaseModel):
    model_config = ConfigDict(extra="forbid")

    name: CategoryName
    matched: bool
    matches: list[Match] = Field(default_factory=list)


class LanguageMetadata(BaseModel):
    model_config = ConfigDict(extra="forbid")

    dominant: SupportedLanguage = Field(
        description="Primary language of the analyzed text."
    )
    alternates: list[SupportedLanguage] = Field(
        default_factory=list,
        description="Secondary language signals for code-mixed or transliterated text.",
    )
    script: ScriptName = Field(description="Primary script used by the text.")
    code_mixed: bool = Field(
        default=False,
        description="True when multiple languages or script systems are materially present.",
    )


class ScenarioMetadata(BaseModel):
    model_config = ConfigDict(extra="forbid")

    name: ScenarioName = Field(
        default="general",
        description="High-level scenario pack that guided the analysis.",
    )
    profile: str = Field(
        default="general",
        description="English scenario profile identifier, such as retail_banking_kyc.",
    )


class Verdict(BaseModel):
    model_config = ConfigDict(extra="forbid")

    overall_blocked: bool
    severity: Severity
    language: LanguageMetadata
    scenario: ScenarioMetadata = Field(default_factory=ScenarioMetadata)
    categories: list[Category]
    reason: str