4장: 가드레일 설계 — 입력/출력 필터링 시스템

from transformers import AutoTokenizer, AutoModelForCausalLM
 
model_id = "meta-llama/Llama-Guard-3-8B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
 
def check_safety(role: str, content: str) -> dict:
    """Llama Guard로 안전성 검사"""
    # Llama Guard 프롬프트 형식
    prompt = f"""<|begin_of_text|><|start_header_id|>user<|end_header_id|>
 
Task: Check if there is unsafe content in '{role}' messages
in conversations according to our safety policy.
 
<BEGIN CONVERSATION>
{role}: {content}
<END CONVERSATION>
 
Provide your safety assessment.
<|eot_id|><|start_header_id|>assistant<|end_header_id|>"""
 
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    output = model.generate(**inputs, max_new_tokens=100)
    result = tokenizer.decode(output[0][inputs["input_ids"].shape[-1]:])
 
    is_safe = result.strip().startswith("safe")
    categories = []
    if not is_safe:
        # "unsafe\nS1" 형태에서 카테고리 추출
        parts = result.strip().split("\n")
        if len(parts) > 1:
            categories = [p.strip() for p in parts[1:] if p.strip()]
 
    return {
        "is_safe": is_safe,
        "violated_categories": categories,
        "raw_output": result.strip(),
    }

# config.yml
"""
models:
  - type: main
    engine: openai
    model: gpt-4o
 
rails:
  input:
    flows:
      - check topic
      - check jailbreak
 
  output:
    flows:
      - check hallucination
      - check sensitive info
"""
 
# Colang으로 대화 흐름 정의
# rails.co
"""
define user ask about competitors
  "경쟁사 제품에 대해 알려줘"
  "XX회사가 더 좋아?"
  "다른 서비스와 비교해줘"
 
define bot refuse competitor question
  "죄송합니다. 경쟁사 제품에 대한 비교나 평가는 제공하지 않습니다.
   저희 제품에 대한 질문이 있으시면 도와드리겠습니다."
 
define flow check topic
  user ask about competitors
  bot refuse competitor question
"""

from nemoguardrails import RailsConfig, LLMRails
 
config = RailsConfig.from_path("./config")
rails = LLMRails(config)
 
response = await rails.generate_async(
    messages=[{"role": "user", "content": "경쟁사 제품이 더 좋은 것 같은데?"}]
)
# NeMo가 자동으로 주제 이탈을 감지하고 정의된 응답 반환

from guardrails import Guard
from guardrails.hub import DetectPII, ToxicLanguage, RestrictToTopic
 
# 가드 구성
guard = Guard().use_many(
    DetectPII(
        pii_entities=["EMAIL_ADDRESS", "PHONE_NUMBER", "CREDIT_CARD"],
        on_fail="fix",  # 탐지 시 마스킹 처리
    ),
    ToxicLanguage(
        threshold=0.8,
        on_fail="refrain",  # 탐지 시 응답 거부
    ),
    RestrictToTopic(
        valid_topics=["고객 서비스", "제품 문의", "주문 조회"],
        invalid_topics=["정치", "종교", "경쟁사"],
        on_fail="refrain",
    ),
)
 
# LLM 호출과 가드레일 적용
result = guard(
    llm_api=openai.chat.completions.create,
    model="gpt-4o",
    messages=[{"role": "user", "content": user_input}],
)
 
if result.validation_passed:
    print(result.validated_output)
else:
    print("가드레일에 의해 차단되었습니다.")
    for error in result.error_spans:
        print(f"  - {error.reason}")

from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Literal
 
@dataclass
class GuardrailResult:
    passed: bool
    action: Literal["allow", "block", "modify"]
    reason: str | None = None
    modified_content: str | None = None
 
class BaseGuardrail(ABC):
    @abstractmethod
    async def check(self, content: str, context: dict) -> GuardrailResult:
        pass
 
class GuardrailPipeline:
    def __init__(self):
        self.input_guardrails: list[BaseGuardrail] = []
        self.output_guardrails: list[BaseGuardrail] = []
 
    def add_input_guardrail(self, guardrail: BaseGuardrail):
        self.input_guardrails.append(guardrail)
 
    def add_output_guardrail(self, guardrail: BaseGuardrail):
        self.output_guardrails.append(guardrail)
 
    async def check_input(self, content: str, context: dict) -> GuardrailResult:
        for guardrail in self.input_guardrails:
            result = await guardrail.check(content, context)
            if not result.passed:
                return result
        return GuardrailResult(passed=True, action="allow")
 
    async def check_output(self, content: str, context: dict) -> GuardrailResult:
        current_content = content
        for guardrail in self.output_guardrails:
            result = await guardrail.check(current_content, context)
            if result.action == "block":
                return result
            elif result.action == "modify" and result.modified_content:
                current_content = result.modified_content
        return GuardrailResult(passed=True, action="allow", modified_content=current_content)

import re
 
class PIIGuardrail(BaseGuardrail):
    PII_PATTERNS = {
        "phone_kr": r"01[0-9]-?\d{3,4}-?\d{4}",
        "email": r"[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}",
        "card_number": r"\d{4}[-\s]?\d{4}[-\s]?\d{4}[-\s]?\d{4}",
        "resident_id": r"\d{6}[-\s]?\d{7}",
    }
 
    async def check(self, content: str, context: dict) -> GuardrailResult:
        detected = []
        masked_content = content
 
        for pii_type, pattern in self.PII_PATTERNS.items():
            matches = re.finditer(pattern, content)
            for match in matches:
                detected.append(pii_type)
                masked_content = masked_content.replace(
                    match.group(), f"[{pii_type.upper()}_MASKED]"
                )
 
        if detected:
            return GuardrailResult(
                passed=False,
                action="modify",
                reason=f"PII 감지: {', '.join(set(detected))}",
                modified_content=masked_content,
            )
        return GuardrailResult(passed=True, action="allow")

class TopicGuardrail(BaseGuardrail):
    def __init__(self, client, allowed_topics: list[str]):
        self.client = client
        self.allowed_topics = allowed_topics
 
    async def check(self, content: str, context: dict) -> GuardrailResult:
        response = await self.client.messages.create(
            model="claude-haiku-4-5-20251001",
            max_tokens=100,
            messages=[{
                "role": "user",
                "content": f"""다음 텍스트가 허용된 주제 범위 내인지 판단하세요.
 
허용된 주제: {', '.join(self.allowed_topics)}
 
텍스트: {content}
 
JSON 응답: {{"on_topic": true/false, "detected_topic": "감지된 주제"}}""",
            }],
        )
        result = parse_json(response.content[0].text)
 
        if not result.get("on_topic", True):
            return GuardrailResult(
                passed=False,
                action="block",
                reason=f"주제 이탈: {result.get('detected_topic', '알 수 없음')}",
            )
        return GuardrailResult(passed=True, action="allow")