프로바이더로 등록하기 — Python 설정 클래스 작성 가이드

프로바이더로 등록하기 — Python 설정 클래스 작성 가이드

지금까지는 "기존 프로바이더를 호출"하는 법을 봤어요. 이제 반대로 생각해 봅시다. 우리 회사의 모델 API를 LiteLLM의 정식 프로바이더로 등록해서, LiteLLM을 쓰는 모든 코드가 우리 모델을 한 줄로 부르게 하려면 어떻게 해야 할까요? JSON 한 파일로 끝나는 OpenAI 호환 케이스부터, 복잡한 변환이 필요한 케이스까지 — 프로바이더 등록 전체 과정을 이 페이지에서 다룹니다.

출처: 공식문서 - Integrate as a Model Provider

시작 전에 — 두 갈래 길

  • 내 API가 OpenAI 호환이라면, 이 페이지의 복잡한 과정 없이 OpenAI 호환 프로바이더 추가(JSON)로 JSON 한 파일만 고치면 끝나요. 이 페이지는 그보다 복잡한 경우를 위한 가이드예요.
  • 이 가이드는 **채팅 프로바이더(chat provider)**로 등록하는 과정에 집중합니다. 임베딩이나 이미지 생성 프로바이더가 필요하면 이 구조를 이해한 뒤 코드베이스의 기존 구현을 참고하면 됩니다.

전체 그림 — LiteLLM이 프로바이더를 보는 법

프로바이더 입장에서 LiteLLM의 동작은 단순해요.

LiteLLM은 래퍼(wrapper)로 동작합니다. OpenAI 요청을 받아서 우리 API로 라우팅하고, 우리가 돌려준 응답을 표준 출력 형식으로 바꿔 줘요. 그래서 우리가 할 일은 그 사이에 끼는 어댑터 모듈을 작성하는 것입니다. 그 모듈은 설정(config)이자 요청·응답 변환기 역할을 겸해요.

모듈에 반드시 들어가야 할 메서드들은 다음과 같아요.

  • 요청을 검증 (validate)
  • OpenAI 형식 요청을 우리 API로 변환 (transform request)
  • 우리 API 응답을 LiteLLM이 기대하는 표준 출력으로 변환 (transform response)
  • 그 외 몇 가지

1단계 — 설정 클래스 만들기

우리 프로바이더 이름으로 새 디렉토리를 만듭니다. litellm/llms/your_provider_name_here

그 안에 채팅 설정 파일을 추가합니다. litellm/llms/your_provider_name_here/chat/transformation.py

transformation.py의 설정 클래스가 "우리 API가 LiteLLM API에 어떻게 끼워질지"를 정의해요. BaseConfig를 상속해 정의하고, 추상 메서드는 나중에 채웁니다.

from litellm.llms.base_llm.chat.transformation import BaseConfig

class MyProviderChatConfig(BaseConfig):
    def __init__(self):
        ...

2단계 — 코드베이스 여기저기에 나를 등록하기

아직 이 과정은 수동입니다. 등록해야 할 지점이 여럿 있어요.

litellm/__init__.py — 파일 상단의 키 목록에 우리 키를 옵션으로 추가합니다.

azure_key: Optional[str] = None
anthropic_key: Optional[str] = None
replicate_key: Optional[str] = None
bytez_key: Optional[str] = None
cohere_key: Optional[str] = None
infinity_key: Optional[str] = None
clarifai_key: Optional[str] = None

그리고 설정 클래스를 import 합니다.

from .llms.bytez.chat.transformation import BytezChatConfig
from .llms.custom_llm import CustomLLM
from .llms.bedrock.chat.converse_transformation import AmazonConverseConfig
from .llms.openai_like.chat.handler import OpenAILikeChatConfig

litellm/main.py — 요청이 우리 설정 클래스로 라우팅되도록 등록합니다.

from .llms.bedrock.chat import BedrockConverseLLM, BedrockLLM
from .llms.bedrock.embed.embedding import BedrockEmbedding
from .llms.bedrock.image.image_handler import BedrockImageGeneration
from .llms.bytez.chat.transformation import BytezChatConfig
from .llms.codestral.completion.handler import CodestralTextCompletion
from .llms.cohere.embed import handler as cohere_embed
from .llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler

base_llm_http_handler = BaseLLMHTTPHandler()
base_llm_aiohttp_handler = BaseLLMAIOHTTPHandler()
sagemaker_chat_completion = SagemakerChatHandler()
bytez_transformation = BytezChatConfig()

그리고 파일 아래쪽의 분기 체인에 우리 프로바이더를 추가합니다. bytez를 예로 들면:

elif custom_llm_provider == "bytez":
    api_key = (
        api_key
        or litellm.bytez_key
        or get_secret_str("BYTEZ_API_KEY")
        or litellm.api_key
    )

    response = base_llm_http_handler.completion(
        model=model,
        messages=messages,
        headers=headers,
        model_response=model_response,
        api_key=api_key,
        api_base=api_base,
        acompletion=acompletion,
        logging_obj=logging,
        optional_params=optional_params,
        litellm_params=litellm_params,
        timeout=timeout,  # type: ignore
        client=client,
        custom_llm_provider=custom_llm_provider,
        encoding=encoding,
        stream=stream,
    )

    pass

여기서 잘 짚고 갈 점: LiteLLM은 .completion() 호출을 통해 우리 설정에 필요한 각종 인자를 그대로 넘겨주므로, 우리는 해당 인자를 편하게 받아 쓰면 됩니다.

litellm/constants.py — LITELLM_CHAT_PROVIDERS 목록에 우리를 추가합니다.

LITELLM_CHAT_PROVIDERS = [
    "openai",
    "openai_like",
    "bytez",
    "xai",
    "custom_openai",
    "text-completion-openai",
    # ...
]

litellm/litellm_core_utils/get_llm_provider_logic.py — 모델 프리픽스 분기 체인에 프로바이더를 추가합니다.

if model == "*":
    custom_llm_provider = "openai"
# bytez models
elif model.startswith("bytez/"):
    custom_llm_provider = "bytez"
if not custom_llm_provider:
    if litellm.suppress_debug_info is False:
        print()  # noqa

litellm/litellm_core_utils/streaming_handler.py — 스트리밍을 커스텀하게 처리한다면 청크 핸들러를 추가합니다.

    def handle_bytez_chunk(self, chunk):
        try:
            is_finished = False
            finish_reason = ""

            return {
                "text": chunk,
                "is_finished": is_finished,
                "finish_reason": finish_reason,
            }
        except Exception as e:
            raise e

3단계 — 반복 개발용 테스트 파일 작성

tests/test_litellm/llms/my_provider/chat/test.py 에 테스트를 추가합니다.

import os
from litellm import completion

os.environ["MY_PROVIDER_KEY"] = "KEY_GOES_HERE"

completion(model="my_provider/your-model", messages=[...], api_key="...")

VSCode 디버거로 실행하고 싶다면 .vscode/launch.json 설정을 쓰면 편리해요. MY_PROVIDER_API_KEY 환경변수를 잡아두면 테스트 스크립트 안의 os.environ["MY_PROVIDER_KEY"] = ... 줄은 지워도 됩니다.

4단계 — 필수 메서드 구현하기

litellm/llms/custom_httpx/llm_http_handler.py의 completion()을 따라가 보면 기본 클래스의 각 메서드가 어떻게 호출되는지 알 수 있어요. 디버거가 가장 좋은 선생님입니다. 주요 메서드들은 이렇게 생겼습니다.

validate_environment — 헤더 구성, 키/모델 검증:

def validate_environment(*args, **kwargs):
    headers.update({
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json"
    })
    return headers

get_complete_url — 최종 요청 URL 반환:

def get_complete_url(*args, **kwargs):
    return f"{api_base}/{model}"

transform_request — OpenAI 스타일 입력을 프로바이더 고유 형식으로 변환:

def transform_request(*args, **kwargs):
    data = {"messages": messages, "params": optional_params}
    return data

transform_response — 프로바이더 원시 응답을 매핑/가공:

def transform_response(*args, **kwargs):
    json = raw_response.json()
    model_response.model = model
    model_response.choices[0].message.content = json.get("output")
    return model_response

get_sync_custom_stream_wrapper / get_async_custom_stream_wrapper — 스트리밍 커스터마이즈가 필요할 때 사용합니다. CustomStreamWrapper + httpx 스트리밍 클라이언트로 콘텐츠를 yield하면 됩니다. 참고로 litellm/llms/sagemaker/chat/transformation.py나 litellm/llms/bytez/chat/transformation.py 구현이 좋은 예시예요.

테스트 & 마무리

tests/test_litellm/llms/my_provider/chat/test.py에서 MY_PROVIDER_KEY를 실제로 바꿔가며 각 메서드를 반복 테스트하세요. 막히면 기존 프로바이더 구현을 찾아보는 게 제일 빠릅니다. ctrl + shift + f와 ctrl + p가 큰 도움이 돼요. 그래도 안 풀리면 LiteLLM Discord 피드백 채널에 질문할 수 있습니다.

더 알아보기