엔드포인트에 가드레일 지원 추가하기

엔드포인트에 가드레일 지원 추가하기

이 가이드는 새 LiteLLM 엔드포인트(예: Chat Completions, Responses API 등)에 가드레일 변환(translation) 지원을 추가하는 방법을 설명해요.

가드레일 지원을 추가할 때

다음 경우에 가드레일 지원을 추가하세요:

• 새 LiteLLM 엔드포인트(예: 새로운 API 형식)를 만들 때
• 가드레일을 지원하지 않는 기존 엔드포인트에서 가드레일을 활성화하고 싶을 때
• 특정 메시지 형식에 대한 커스텀 텍스트 추출 로직이 필요할 때

디렉터리 구조

가드레일 핸들러는 다음 구조를 따릅니다:

litellm/llms/{provider}/{endpoint}/guardrail_translation/
├── __init__.py          # Exports handler and registers call types
├── handler.py           # Main handler implementation
└── README.md            # Documentation (optional but recommended)

예시 구조

OpenAI Chat Completions:

litellm/llms/openai/chat/guardrail_translation/
├── __init__.py
├── handler.py
└── README.md

OpenAI Responses API:

litellm/llms/openai/responses/guardrail_translation/
├── __init__.py
├── handler.py
└── README.md

Anthropic Messages:

litellm/llms/anthropic/chat/guardrail_translation/
├── __init__.py
└── handler.py

출처: 문서

본문

단계별 구현

1단계: 핸들러 클래스 만들기

BaseTranslation을 상속하는 handler.py를 만드세요:

"""
{Provider} {Endpoint} Handler for Unified Guardrails

This module provides guardrail translation support for {Provider}'s {Endpoint} format.
"""

from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast

from litellm._logging import verbose_proxy_logger
from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation
from litellm.types.utils import GenericGuardrailAPIInputs

if TYPE_CHECKING:
    from litellm.integrations.custom_guardrail import CustomGuardrail
    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
    from litellm.proxy._types import UserAPIKeyAuth
    from litellm.types.utils import ModelResponse  # Or appropriate response type

class MyEndpointHandler(BaseTranslation):
    """
    Handler for processing {Endpoint} with guardrails.

    This class provides methods to:
    1. Process input (pre-call hook)
    2. Process output response (post-call hook)
    """

    async def process_input_messages(
        self,
        data: dict,
        guardrail_to_apply: "CustomGuardrail",
        litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
    ) -> Any:
        """
        Process input by applying guardrails to text content.

        Args:
            data: Request data dictionary
            guardrail_to_apply: The guardrail instance to apply
            litellm_logging_obj: Logging object for the call, forwarded to the guardrail

        Returns:
            Modified data with guardrails applied
        """
        # Your implementation here
        pass

    async def process_output_response(
        self,
        response: Any,  # Use appropriate response type
        guardrail_to_apply: "CustomGuardrail",
        litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
        user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
        request_data: dict | None = None,
    ) -> Any:
        """
        Process output response by applying guardrails to text content.

        Args:
            response: API response object
            guardrail_to_apply: The guardrail instance to apply
            litellm_logging_obj: Logging object for the call, forwarded to the guardrail
            user_api_key_dict: Caller identity, passed separately since the response has none
            request_data: The originating request body

        Returns:
            Modified response with guardrails applied
        """
        # Your implementation here
        pass

2단계: 핵심 메서드 구현

A. 입력 메시지 처리

입력에서 텍스트를 추출하고, 가드레일을 적용하고, 다시 매핑합니다:

async def process_input_messages(
    self,
    data: dict,
    guardrail_to_apply: "CustomGuardrail",
    litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
) -> Any:
    """Process input messages by applying guardrails to text content."""
    # 1. Get input data from request
    messages = data.get("messages")  # or appropriate field
    if messages is None:
        return data

    # 2. Extract text into a flat list, remembering where each item came from
    texts_to_check: List[str] = []
    task_mappings: List[Tuple[int, Optional[int]]] = []

    for msg_idx, message in enumerate(messages):
        self._extract_input_texts(
            message=message,
            msg_idx=msg_idx,
            texts_to_check=texts_to_check,
            task_mappings=task_mappings,
        )

    # 3. Apply the guardrail to everything in one call
    if texts_to_check:
        guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
            inputs=GenericGuardrailAPIInputs(texts=texts_to_check),
            request_data=data,
            input_type="request",
            logging_obj=litellm_logging_obj,
        )

        # 4. Map the returned texts back to the original structure
        await self._apply_guardrail_responses_to_input(
            messages=messages,
            responses=guardrailed_inputs.get("texts", []),
            task_mappings=task_mappings,
        )

    return data

GenericGuardrailAPIInputs는 images, tools, tool_calls, structured_messages, model도 담아요. 엔드포인트 형식이 가진 키를 채우고, 가드레일이 반환하는 키를 다시 매핑하세요.

메시지당 한 번이 아니라 모든 것을 한 번의 호출로 보내세요. 그래야 가드레일 프로바이더가 전체 대화를 볼 수 있고, 왕복(round trip)도 여러 번이 아닌 한 번만 지불합니다.

B. 출력 응답 처리

응답에서 텍스트를 추출하고, 가드레일을 적용하고, 업데이트합니다:

async def process_output_response(
    self,
    response: "ModelResponse",
    guardrail_to_apply: "CustomGuardrail",
    litellm_logging_obj: Optional["LiteLLMLoggingObj"] = None,
    user_api_key_dict: Optional["UserAPIKeyAuth"] = None,
    request_data: dict | None = None,
) -> Any:
    """Process output response by applying guardrails to text content."""
    # 1. Check if response has text to process
    if not self._has_text_content(response):
        return response

    # 2. Extract text into a flat list, remembering where each item came from
    texts_to_check: List[str] = []
    task_mappings: List[Tuple[int, Optional[int]]] = []

    for idx, item in enumerate(response.choices):  # or appropriate field
        self._extract_output_texts(
            item=item,
            idx=idx,
            texts_to_check=texts_to_check,
            task_mappings=task_mappings,
        )

    # 3. Apply the guardrail to everything in one call
    if texts_to_check:
        guardrailed_inputs = await guardrail_to_apply.apply_guardrail(
            inputs=GenericGuardrailAPIInputs(texts=texts_to_check),
            request_data=request_data or {},
            input_type="response",
            logging_obj=litellm_logging_obj,
        )

        # 4. Update response with guardrailed text
        await self._apply_guardrail_responses_to_output(
            response=response,
            responses=guardrailed_inputs.get("texts", []),
            task_mappings=task_mappings,
        )

    return response

3단계: 헬퍼 메서드 만들기

텍스트 추출과 매핑을 위한 헬퍼 메서드를 구현하세요:

def _extract_input_texts(
    self,
    message: Dict[str, Any],
    msg_idx: int,
    texts_to_check: List[str],
    task_mappings: List[Tuple[int, Optional[int]]],
) -> None:
    """Extract text content from a message and record where each piece came from."""
    content = message.get("content")
    if content is None:
        return

    if isinstance(content, str):
        # Simple string content
        texts_to_check.append(content)
        task_mappings.append((msg_idx, None))
    elif isinstance(content, list):
        # List content (e.g., multimodal)
        for content_idx, content_item in enumerate(content):
            if isinstance(content_item, dict):
                text_str = content_item.get("text")
                if text_str:
                    texts_to_check.append(text_str)
                    task_mappings.append((msg_idx, int(content_idx)))

async def _apply_guardrail_responses_to_input(
    self,
    messages: List[Dict[str, Any]],
    responses: List[str],
    task_mappings: List[Tuple[int, Optional[int]]],
) -> None:
    """Apply guardrail responses back to input messages."""
    for task_idx, guardrail_response in enumerate(responses):
        msg_idx, content_idx = task_mappings[task_idx]
        
        if content_idx is None:
            # String content
            messages[msg_idx]["content"] = guardrail_response
        else:
            # List content
            messages[msg_idx]["content"][content_idx]["text"] = guardrail_response

def _has_text_content(self, response: Any) -> bool:
    """Check if response has any text content to process."""
    # Implement based on your response structure
    return True  # or appropriate logic

4단계: 핸들러 등록

호출 타입과 함께 핸들러를 등록하는 __init__.py를 만드세요:

"""My Endpoint handler for Unified Guardrails."""

from litellm.llms.{provider}/{endpoint}/guardrail_translation.handler import (
    MyEndpointHandler,
)
from litellm.types.utils import CallTypes

guardrail_translation_mappings = {
    CallTypes.my_endpoint: MyEndpointHandler,
    CallTypes.amy_endpoint: MyEndpointHandler,  # async version if applicable
}

__all__ = ["guardrail_translation_mappings"]

중요: CallTypes가 litellm/types/utils.py에 정의되어 있는지 확인하세요.

5단계: 문서 추가

사용 예시와 형식 세부 사항이 담긴 README.md를 만드세요:

# {Provider} {Endpoint} Guardrail Translation Handler

Handler for processing {Provider}'s {Endpoint} with guardrails.

## Overview

This handler processes {Endpoint} input/output by:
1. Extracting text from messages/responses
2. Applying guardrails to text content
3. Mapping guardrailed text back to original structure

## Data Format

### Input Format
```json
{
  "field": "value",
  "messages": [...]
}

Output Format

{
  "field": "value",
  "output": [...]
}

Usage

The handler is automatically discovered and applied when guardrails are used with this endpoint.

curl -X POST 'http://localhost:4000/{my_endpoint}' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer ***' \
-d '{
    "model": "gpt-5.6-luna",
    "messages": [{"role": "user", "content": "Hello"}],
    "guardrails": ["test"]
}'

Extension

Override these methods to customize behavior:

• _extract_input_texts() : Custom text extraction
• _apply_guardrail_responses_to_input() : Custom response mapping
• _has_text_content() : Custom content detection

#### 6단계: 유닛 테스트 추가

`tests/test_litellm/llms/{provider}/{endpoint}/`에 철저한 테스트를 만드세요:

```python
"""
Unit tests for {Provider} {Endpoint} Guardrail Translation Handler
"""

import os
import sys
from typing import TYPE_CHECKING, Literal, Optional

import pytest

sys.path.insert(0, os.path.abspath("../../../../../.."))

from litellm.integrations.custom_guardrail import CustomGuardrail
from litellm.llms import get_guardrail_translation_mapping
from litellm.llms.{provider}.{endpoint}.guardrail_translation.handler import (
    MyEndpointHandler,
)
from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs

if TYPE_CHECKING:
    from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj

class MockGuardrail(CustomGuardrail):
    """Mock guardrail for testing"""

    async def apply_guardrail(
        self,
        inputs: GenericGuardrailAPIInputs,
        request_data: dict,
        input_type: Literal["request", "response"],
        logging_obj: Optional["LiteLLMLoggingObj"] = None,
    ) -> GenericGuardrailAPIInputs:
        inputs["texts"] = [f"{text} [GUARDRAILED]" for text in inputs.get("texts", [])]
        return inputs

class TestHandlerDiscovery:
    """Test that the handler is properly discovered"""
    
    def test_handler_discovered(self):
        handler_class = get_guardrail_translation_mapping(CallTypes.my_endpoint)
        assert handler_class == MyEndpointHandler

class TestInputProcessing:
    """Test input processing functionality"""
    
    @pytest.mark.asyncio
    async def test_process_simple_input(self):
        handler = MyEndpointHandler()
        guardrail = MockGuardrail(guardrail_name="test")
        
        data = {"messages": [{"role": "user", "content": "Hello"}]}
        result = await handler.process_input_messages(data, guardrail)
        
        assert result["messages"][0]["content"] == "Hello [GUARDRAILED]"

class TestOutputProcessing:
    """Test output processing functionality"""
    
    @pytest.mark.asyncio
    async def test_process_simple_output(self):
        handler = MyEndpointHandler()
        guardrail = MockGuardrail(guardrail_name="test")
        
        # Create mock response
        response = create_mock_response()
        result = await handler.process_output_response(response, guardrail)
        
        # Assert guardrail was applied
        assert "GUARDRAILED" in get_response_text(result)

지원

질문이나 이슈가 있다면:

• 기존 핸들러 구현을 예시로 확인하세요
• 기본 변환 클래스 문서를 검토하세요
• GitHub에 guardrails 레이블로 이슈를 만들어 주세요

더 알아보기 (Learn more)

• Generic Guardrail API
• Guardrails