커스텀 API 서버 연결하기

커스텀 API 서버 연결하기 (Custom Format)

우리 회사엔 torch-serve로 띄운 내부 LLM API가 있어요. OpenAI 호환도 아니고, LiteLLM이 정식 지원하는 프로바이더도 아닙니다. 이런 독자 형식의 API를 LiteLLM에 연결하려면 어떻게 해야 할까요? 답은 CustomLLM 클래스로 핸들러를 직접 구현하는 거예요. 이 페이지가 그 전체 과정을 보여줍니다.

출처: 공식문서 - Custom API Server (Custom Format)

어떤 경로를 지원하나요

커스텀 핸들러는 아래 경로를 각각 매핑할 수 있어요.

/v1/chat/completions → litellm.acompletion
/v1/completions → litellm.atext_completion
/v1/embeddings → litellm.aembedding
/v1/images/generations → litellm.aimage_generation
/v1/images/edits → litellm.aimage_edit
/v1/messages → litellm.acompletion

참고로 OpenAI 호환 엔드포인트라면 이 복잡한 과정이 필요 없어요. OpenAI 호환 엔드포인트 페이지를 먼저 보는 게 좋고, 호출 전후를 가로채 바꾸고 싶다면 프록시의 call_hooks 문서를 보세요.

빠른 시작 — 핸들러 등록

CustomLLM을 상속한 클래스를 만들고, 그 인스턴스를 litellm.custom_provider_map에 등록하면 끝입니다. 아래 예제는 어떤 요청이든 "Hi!"를 돌려주는 간단한 핸들러예요.

import litellm
from litellm import CustomLLM, completion, get_llm_provider


class MyCustomLLM(CustomLLM):
    def completion(self, *args, **kwargs) -> litellm.ModelResponse:
        return litellm.completion(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": "Hello world"}],
            mock_response="Hi!",
        )  # type: ignore

my_custom_llm = MyCustomLLM()

litellm.custom_provider_map = [ # 핵심 단계 - 핸들러 등록
        {"provider": "my-custom-llm", "custom_handler": my_custom_llm}
    ]

resp = completion(
        model="my-custom-llm/my-fake-model",
        messages=[{"role": "user", "content": "Hello world!"}],
    )

assert resp.choices[0].message.content == "Hi!"

핵심은 두 부분이에요. CustomLLM 클래스를 상속해 completion 메서드를 구현하는 것, 그리고 그 인스턴스를 custom_provider_map에 {"provider", "custom_handler"} 형태로 등록하는 것입니다. 이후 model="my-custom-llm/my-fake-model"처럼 provider/model 형태로 부르면 LiteLLM이 커스텀 핸들러로 라우팅해요.

OpenAI Proxy에서 쓰기

SDK 없이 Proxy 서버로 같은 핸들러를 노출하는 방식입니다. custom_handler.py 파일에 핸들러를 정의하고, config.yaml에서 파일 이름과 인스턴스 이름을 지정하면 돼요.

1. custom_handler.py 작성

import litellm
from litellm import CustomLLM, completion, get_llm_provider


class MyCustomLLM(CustomLLM):
    def completion(self, *args, **kwargs) -> litellm.ModelResponse:
        return litellm.completion(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": "Hello world"}],
            mock_response="Hi!",
        )  # type: ignore

    async def acompletion(self, *args, **kwargs) -> litellm.ModelResponse:
        return litellm.completion(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": "Hello world"}],
            mock_response="Hi!",
        )  # type: ignore


my_custom_llm = MyCustomLLM()

2. config.yaml에 등록

파일에서 넘겨줄 값은 셋이에요. python_filename은 핸들러 파일명(custom_handler.py), custom_handler_instance_name은 파일 안 인스턴스 이름(my_custom_llm), custom_handler는 점 표기(custom_handler.my_custom_llm)입니다.

model_list:
  - model_name: "test-model"             
    litellm_params:
      model: "openai/text-embedding-ada-002"
  - model_name: "my-custom-model"
    litellm_params:
      model: "my-custom-llm/my-model"

litellm_settings:
  custom_provider_map:
  - {"provider": "my-custom-llm", "custom_handler": custom_handler.my_custom_llm}

litellm --config /path/to/config.yaml

3. 테스트

curl -X POST 'http://0.0.0.0:4000/chat/completions' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer ***' \
-d '{
    "model": "my-custom-model",
    "messages": [{"role": "user", "content": "Say \"this is a test\" in JSON!"}],
}'

응답은 핸들러가 돌려준 내용 그대로 옵니다. 위 예시에서는 content가 "Hi!"로 채워진 OpenAI 형식 응답이 반환돼요.

스트리밍 지원 추가하기

completion과 streaming 모두에서 유닉스 타임스탬프를 돌려주는 예시입니다. streaming/astreaming 메서드에서 GenericStreamingChunk를 만들어 yield하면 됩니다.

import time
from typing import Iterator, AsyncIterator
from litellm.types.utils import GenericStreamingChunk, ModelResponse
from litellm import CustomLLM, completion, acompletion

class UnixTimeLLM(CustomLLM):
    def completion(self, *args, **kwargs) -> ModelResponse:
        return completion(
            model="test/unixtime",
            mock_response=str(int(time.time())),
        )  # type: ignore

    async def acompletion(self, *args, **kwargs) -> ModelResponse:
        return await acompletion(
            model="test/unixtime",
            mock_response=str(int(time.time())),
        )  # type: ignore

    def streaming(self, *args, **kwargs) -> Iterator[GenericStreamingChunk]:
        generic_streaming_chunk: GenericStreamingChunk = {
            "finish_reason": "stop",
            "index": 0,
            "is_finished": True,
            "text": str(int(time.time())),
            "tool_use": None,
            "usage": {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
        }
        return generic_streaming_chunk # type: ignore

    async def astreaming(self, *args, **kwargs) -> AsyncIterator[GenericStreamingChunk]:
        generic_streaming_chunk: GenericStreamingChunk = {
            "finish_reason": "stop",
            "index": 0,
            "is_finished": True,
            "text": str(int(time.time())),
            "tool_use": None,
            "usage": {"completion_tokens": 0, "prompt_tokens": 0, "total_tokens": 0},
        }
        yield generic_streaming_chunk # type: ignore

unixtime = UnixTimeLLM()

스트리밍에서 중요한 점은, 각 청크에 is_finished와 finish_reason, 그리고 text 청크 내용을 담아야 한다는 거예요. LiteLLM의 GenericStreamingChunk가 그 규격을 정해 주므로 이를 맞춰 주면 됩니다.

이미지 생성 & 편집

aimage_generation, aimage_edit 메서드를 구현해서 이미지 생성/편집도 커스텀하게 처리할 수 있어요. 예를 들어 Stability AI나 Black Forest Labs를 내가 원하는 방식으로 호출하고 싶다면 여기서 직접 하면 됩니다.

import litellm
from litellm import CustomLLM
from litellm.types.utils import ImageResponse, ImageObject


class MyCustomLLM(CustomLLM):
    async def aimage_generation(self, model: str, prompt: str, model_response: ImageResponse, optional_params: dict, logging_obj: Any, timeout: Optional[Union[float, httpx.Timeout]] = None, client: Optional[AsyncHTTPHandler] = None,) -> ImageResponse:
        return ImageResponse(
            created=int(time.time()),
            data=[ImageObject(url="https://example.com/image.png")],
        )

my_custom_llm = MyCustomLLM()

config.yaml에 모델을 추가하고 띄운 뒤,

model_list:
  - model_name: "test-model"             
    litellm_params:
      model: "openai/text-embedding-ada-002"
  - model_name: "my-custom-model"
    litellm_params:
      model: "my-custom-llm/my-model"

litellm_settings:
  custom_provider_map:
  - {"provider": "my-custom-llm", "custom_handler": custom_handler.my_custom_llm}

litellm --config /path/to/config.yaml

이렇게 테스트합니다.

curl -X POST 'http://0.0.0.0:4000/v1/images/generations' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer ***' \
-d '{
    "model": "my-custom-model",
    "prompt": "A cute baby sea otter",
}'

Anthropic `/v1/messages` 지원

커스텀 핸들러의 .acompletion만 작성하면 LiteLLM이 그걸 /v1/messages(Anthropic 형식)로 변환해 줘요. 핸들러는 OpenAI 형식으로 구현하고, 경로 변환은 LiteLLM이 담당합니다. 등록·테스트 흐름은 앞선 채팅 예시와 동일하고, 테스트는 아래처럼 하면 됩니다.

curl -L -X POST 'http://0.0.0.0:4000/v1/messages' \
-H 'anthropic-version: 2023-06-01' \
-H 'content-type: application/json' \
-H 'Authorization: Bearer ***' \
-d '{
   "model": "my-custom-model",
     "max_tokens": 1024,
     "messages": [{
         "role": "user",
         "content": [
         {
             "type": "text",
             "text": "What are the key findings in this document 12?"
         }]
     }]
}'

추가 파라미터 넘기기

호출 시 넘긴 추가 파라미터는 핸들러의 optional_params 키로 전달돼요. SDK에서는 아래처럼 확인할 수 있습니다.

import litellm
from litellm import CustomLLM, completion, get_llm_provider


class MyCustomLLM(CustomLLM):
    def completion(self, *args, **kwargs) -> litellm.ModelResponse:
        assert kwargs["optional_params"] == {"my_custom_param": "my-custom-param"} # 확인 지점
        return litellm.completion(
            model="gpt-3.5-turbo",
            messages=[{"role": "user", "content": "Hello world"}],
            mock_response="Hi!",
        )  # type: ignore

my_custom_llm = MyCustomLLM()

litellm.custom_provider_map = [ # 핵심 단계 - 핸들러 등록
        {"provider": "my-custom-llm", "custom_handler": my_custom_llm}
    ]

resp = completion(model="my-custom-llm/my-model", my_custom_param="my-custom-param")

Custom Handler 스펙 — 무엇을 구현해야 하나

CustomLLM은 BaseLLM을 상속하고, 기본적으로 모든 메서드가 "구현 안 됨" 예외를 던집니다. 그래서 우리가 필요한 것만 골라 구현하면 돼요.

completion / acompletion — 채팅 완성 (동기/비동기)
streaming / astreaming — 스트리밍 청크 생성
image_generation / aimage_generation — 이미지 생성
image_edit / aimage_edit — 이미지 편집

예외는 CustomLLMError(status_code, message)를 만들어 쓰는 게 관례예요. 각 메서드 시그니처는 공식 문서의 Custom Handler Spec에서 확인할 수 있고, LiteLLM은 정의한 메서드에 필요한 인자를 통째로 전달해 주므로 내부에서 그대로 활용하면 됩니다.

정리 — 어느 방식을 언제 쓰나요

OpenAI 호환이면 → openai/ 프리픽스 또는 JSON 등록 (간단)
독자 형식이면 → CustomLLM 핸들러로 직접 구현 (이 페이지)
프로바이더로 "정식" 등록하고 싶다면 → 다음 페이지의 Python 설정 클래스 작성법을 따라가세요

더 알아보기

OpenAI 호환 엔드포인트 연결 — OpenAI 호환 서버일 때
프로바이더 등록 가이드 — Python 설정 클래스로 정식 프로바이더 추가
OpenAI 호환 프로바이더 추가(JSON) — JSON 한 파일로 등록

커스텀 API 서버 연결하기