[Beta] 커스텀 프롬프트 관리

[Beta] 커스텀 프롬프트 관리 (Custom Prompt Management)

커스텀 훅으로 LiteLLM을 자체 프롬프트 관리 시스템에 연결해요.

출처: 문서

본문

개요 (Overview)

동작 방식 (How it works)

빠른 시작 (Quick Start)

1. 커스텀 프롬프트 관리자 생성 (Create Your Custom Prompt Manager)

프롬프트 검색과 포맷팅을 처리하려면 CustomPromptManagement에서 상속하는 클래스를 만드세요.

예시 구현 (Example Implementation)

custom_prompt.py라는 새 파일을 만들고 이 코드를 추가하세요. 핵심 메서드는 get_chat_completion_prompt이며, prompt_id와 prompt_variables에 따라 프롬프트를 검색·포맷하는 커스텀 로직을 구현할 수 있어요.

from typing import List, Tuple, Optional
from litellm.integrations.custom_prompt_management import CustomPromptManagement
from litellm.types.llms.openai import AllMessageValues
from litellm.types.utils import StandardCallbackDynamicParams

class MyCustomPromptManagement(CustomPromptManagement):
    def get_chat_completion_prompt(
        self,
        model: str,
        messages: List[AllMessageValues],
        non_default_params: dict,
        prompt_id: str,
        prompt_variables: Optional[dict],
        dynamic_callback_params: StandardCallbackDynamicParams,
    ) -> Tuple[str, List[AllMessageValues], dict]:
        """
        Retrieve and format prompts based on prompt_id.

        Returns:
            - model: The model to use
            - messages: The formatted messages
            - non_default_params: Optional parameters like temperature
        """
        # Example matching the diagram: Add system message for prompt_id "1234"
        if prompt_id == "1234":
            # Prepend system message while preserving existing messages
            new_messages = [
                {"role": "system", "content": "Be a good Bot!"},
            ] + messages
            return model, new_messages, non_default_params

        # Default: Return original messages if no prompt_id match
        return model, messages, non_default_params

prompt_management = MyCustomPromptManagement()

2. LiteLLM config.yaml에서 프롬프트 관리자 구성

model_list:
  - model_name: gpt-5.6-terra
    litellm_params:
      model: openai/gpt-5.6-terra
      api_key: os.environ/OPENAI_API_KEY
litellm_settings:
  callbacks: custom_prompt.prompt_management  # sets litellm.callbacks = [prompt_management]

3. LiteLLM 게이트웨이 시작

Docker Run

LiteLLM Docker 컨테이너에 custom_logger.py를 마운트하세요.

docker run -d \
  -p 4000:4000 \
  -e OPENAI_API_KEY=$OPENAI_API_KEY \
  --name my-app \
  -v $(pwd)/my_config.yaml:/app/config.yaml \
  -v $(pwd)/custom_logger.py:/app/custom_logger.py \
  my-app:latest \
  --config /app/config.yaml \
  --port 4000 \
  --detailed_debug \

litellm pip

litellm --config config.yaml --detailed_debug

4. 커스텀 프롬프트 관리자 테스트

prompt_id="1234"를 전달하면 커스텀 프롬프트 관리자가 대화에 system 메시지 "Be a good Bot!"를 추가해요:

OpenAI Python v1.0.0+

from openai import OpenAI
client = OpenAI(
    api_key="sk-<your-litellm-api-key>",
    base_url="http://0.0.0.0:4000")
response = client.chat.completions.create(
    model="gpt-5.6-terra",
    messages=[{"role": "user", "content": "hi"}],
    extra_body={
        "prompt_id": "1234"
    })
print(response.choices[0].message.content)

Langchain

from langchain.chat_models import ChatOpenAI
from langchain.schema import HumanMessage

chat = ChatOpenAI(
    model="gpt-5.6-terra",
    openai_api_key="sk-<your-api-key>",
    openai_api_base="http://0.0.0.0:4000",
    extra_body={
        "prompt_id": "1234"
    })
messages = []
response = chat(messages)
print(response.content)

Curl

curl -X POST http://0.0.0.0:4000/v1/chat/completions \
-H "Content-Type: application/json" \
-H "Authorization: Bearer ***" \
-d '{
    "model": "gpt-5.6-terra",
    "messages": [{"role": "user", "content": "hi"}],
    "prompt_id": "1234"
}'

LiteLLM SDK 직접 사용 (Using the LiteLLM SDK Directly)

Python 스크립트에서 litellm.completion()을 호출한다면(프록시를 거치지 않고), 요청 전에 커스텀 프롬프트 관리자를 등록하세요:

import litellm
from custom_prompt import prompt_management

litellm.callbacks = [prompt_management]
litellm.use_litellm_proxy = True

response = litellm.completion(
    model="gpt-5.6-terra",
    messages=[{"role": "user", "content": "hi"}],
    prompt_id="1234",
    prompt_variables={"user_message": "hi"},
)

참고: SDK 스크립트에는 litellm.callbacks = [prompt_management](동등하게 litellm.logging_callback_manager.add_litellm_callback(prompt_management))가 필요해요. 프록시는 config.yaml에서 callbacks를 자동으로 읽지만, 독립 스크립트는 그렇지 않아요.

요청은 다음과 같이 변환돼요:

원본:

{
    "model": "gpt-5.6-terra",
    "messages": [{"role": "user", "content": "hi"}],
    "prompt_id": "1234"
}

변환 후:

{
    "model": "gpt-5.6-terra",
    "messages": [
        {"role": "system", "content": "Be a good Bot!"},
        {"role": "user", "content": "hi"}
    ]
}