Guardrails

import re
from dataclasses import dataclass
from langchain_anthropic import ChatAnthropic
from pydantic import BaseModel

@dataclass
class GuardrailResult:
    passed: bool
    reason: str = ""
    sanitized_input: str = ""

# ─── 1. PII 마스킹 ──────────────────────────────────
PII_PATTERNS = {
    "email":        (r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}",        "[EMAIL]"),
    "phone_kr":     (r"01[016789]-?\d{3,4}-?\d{4}",                          "[PHONE]"),
    "rrn":          (r"\d{6}-[1-4]\d{6}",                                     "[RRN]"),   # 주민번호
    "credit_card":  (r"\d{4}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}",             "[CARD]"),
    "ip_address":   (r"\b(?:\d{1,3}\.){3}\d{1,3}\b",                        "[IP]"),
}

def mask_pii(text: str) -> tuple[str, list[str]]:
    """PII 패턴을 마스킹하고 발견된 PII 타입 목록 반환"""
    detected = []
    for pii_type, (pattern, mask) in PII_PATTERNS.items():
        if re.search(pattern, text):
            detected.append(pii_type)
            text = re.sub(pattern, mask, text)
    return text, detected

# ─── 2. 프롬프트 인젝션 탐지 ────────────────────────
INJECTION_PATTERNS = [
    r"ignore (previous|all|above) instructions?",
    r"you are now (a|an)?\s+\w+",
    r"forget (your|all) (instructions?|rules?|guidelines?)",
    r"act as (if you are|a)?\s+",
    r"(system|admin|root) (prompt|message|override)",
    r"jailbreak",
    r"DAN (mode|persona)",
]

def detect_injection(text: str) -> bool:
    text_lower = text.lower()
    return any(re.search(p, text_lower) for p in INJECTION_PATTERNS)

# ─── 3. 토픽 경계 검사 (LLM 기반) ──────────────────
class TopicCheck(BaseModel):
    on_topic: bool
    topic: str
    reason: str

async def check_topic(
    text: str,
    allowed_topics: list[str],
    llm: ChatAnthropic
) -> bool:
    structured = llm.with_structured_output(TopicCheck)
    result = await structured.ainvoke([{
        "role": "user",
        "content": f"""사용자 메시지가 허용된 주제에 해당하는지 판단하세요.

허용 주제: {', '.join(allowed_topics)}
사용자 메시지: {text}

허용 주제 외 내용(금융 조언, 의학 진단, 법적 조언, 성인 콘텐츠 등)이면 on_topic=false"""
    }])
    return result.on_topic

# ─── 통합 입력 Guardrail ────────────────────────────
async def input_guardrail(
    user_input: str,
    allowed_topics: list[str] = [],
    max_length: int = 4000
) -> GuardrailResult:

    # 1. 길이 제한
    if len(user_input) > max_length:
        return GuardrailResult(False, f"입력이 너무 깁니다 ({len(user_input)}/{max_length}자)")

    # 2. 프롬프트 인젝션 탐지
    if detect_injection(user_input):
        return GuardrailResult(False, "허용되지 않는 지시가 감지됐습니다")

    # 3. PII 마스킹
    sanitized, detected_pii = mask_pii(user_input)
    if detected_pii:
        print(f"[PII] 감지된 개인정보: {detected_pii}")  # 감사 로그

    # 4. 토픽 경계 (설정된 경우)
    if allowed_topics:
        llm = ChatAnthropic(model="claude-haiku-4-5-20251001")
        if not await check_topic(sanitized, allowed_topics, llm):
            return GuardrailResult(False, "해당 주제는 지원하지 않습니다")

    return GuardrailResult(True, sanitized_input=sanitized)

class OutputSafetyCheck(BaseModel):
    safe: bool
    issues: list[str]
    severity: str  # "low" | "medium" | "high"

async def output_guardrail(
    agent_output: str,
    original_question: str
) -> GuardrailResult:

    # 1. PII 재마스킹 (Agent가 개인정보를 노출했을 경우)
    sanitized, leaked_pii = mask_pii(agent_output)
    if leaked_pii:
        print(f"[WARN] 출력에 PII 감지됨: {leaked_pii}")

    # 2. LLM 기반 안전성 검사
    llm = ChatAnthropic(model="claude-haiku-4-5-20251001")
    structured = llm.with_structured_output(OutputSafetyCheck)

    check = await structured.ainvoke([{
        "role": "user",
        "content": f"""AI 응답의 안전성을 검사하세요.

원본 질문: {original_question}
AI 응답: {sanitized[:1000]}

다음을 확인하세요:
- 유해하거나 위험한 정보 포함 여부
- 명백한 사실 오류 여부
- 의도치 않은 민감 정보 노출 여부
- 질문과 무관한 내용 여부"""
    }])

    if not check.safe and check.severity == "high":
        return GuardrailResult(
            False,
            f"출력 안전성 문제: {', '.join(check.issues)}"
        )

    return GuardrailResult(True, sanitized_input=sanitized)

# LangGraph 노드로 통합
async def guardrail_node(state) -> dict:
    user_msg = state["messages"][0].content

    # 입력 검증
    in_result = await input_guardrail(
        user_msg,
        allowed_topics=["기술 지원", "제품 문의", "일반 정보"]
    )
    if not in_result.passed:
        from langchain_core.messages import AIMessage
        return {"messages": [AIMessage(
            content=f"죄송합니다. {in_result.reason}"
        )], "blocked": True}

    return {
        "messages": [{"role": "user", "content": in_result.sanitized_input}],
        "blocked": False
    }