커스텀 LLM 구현하기

커스텀 LLM 구현하기 (Custom LLM Implementation)

LiteLLM이 지원하지 않는 제공자나, 특별한 인증 방식이 필요한 경우 CrewAI의 BaseLLM 추상 베이스 클래스를 상속해 나만의 LLM을 구현할 수 있어요. 이 글에서는 최소 구현에서부터 함수 호출, 스톱워드 처리, 오류 처리까지 실제로 동작하는 패턴을 살펴볼게요.

출처: 공식문서

본문

CrewAI는 BaseLLM 추상 베이스 클래스를 통한 커스텀 LLM 구현을 지원합니다. 이렇게 하면 LiteLLM에 내장 지원이 없는 제공자를 붙이거나, 커스텀 인증 메커니즘을 구현할 수 있어요.

퀵 스타트 (Quick Start)

가장 단순한 커스텀 LLM 구현은 이렇게 생겼어요.

from crewai import BaseLLM
from typing import Any, Dict, List, Optional, Union
import requests

class CustomLLM(BaseLLM):
    def __init__(self, model: str, api_key: str, endpoint: str, temperature: Optional[float] = None):
        # IMPORTANT: Call super().__init__() with required parameters
        super().__init__(model=model, temperature=temperature)

        self.api_key = api_key
        self.endpoint = endpoint

    def call(
        self,
        messages: Union[str, List[Dict[str, str]]],
        tools: Optional[List[dict]] = None,
        callbacks: Optional[List[Any]] = None,
        available_functions: Optional[Dict[str, Any]] = None,
    ) -> Union[str, Any]:
        """Call the LLM with the given messages."""
        # Convert string to message format if needed
        if isinstance(messages, str):
            messages = [{"role": "user", "content": messages}]

        # Prepare request
        payload = {
            "model": self.model,
            "messages": messages,
            "temperature": self.temperature,
        }

        # Add tools if provided and supported
        if tools and self.supports_function_calling():
            payload["tools"] = tools

        # Make API call
        response = requests.post(
            self.endpoint,
            headers={
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            },
            json=payload,
            timeout=30
        )
        response.raise_for_status()

        result = response.json()
        return result["choices"][0]["message"]["content"]

    def supports_function_calling(self) -> bool:
        """Override if your LLM supports function calling."""
        return True  # Change to False if your LLM doesn't support tools

    def get_context_window_size(self) -> int:
        """Return the context window size of your LLM."""
        return 8192  # Adjust based on your model's actual context window

커스텀 LLM 사용하기

정의한 커스텀 LLM을 에이전트에 연결하면 됩니다.

from crewai import Agent, Task, Crew

# Assuming you have the CustomLLM class defined above
# Create your custom LLM
custom_llm = CustomLLM(
    model="my-custom-model",
    api_key="your-api-key",
    endpoint="https://api.example.com/v1/chat/completions",
    temperature=0.7
)

# Use with an agent
agent = Agent(
    role="Research Assistant",
    goal="Find and analyze information",
    backstory="You are a research assistant.",
    llm=custom_llm
)

# Create and execute tasks
task = Task(
    description="Research the latest developments in AI",
    expected_output="A comprehensive summary",
    agent=agent
)

crew = Crew(agents=[agent], tasks=[task])
result = crew.kickoff()

필수 메서드 (Required Methods)

생성자 __init__() — 필수 파라미터 model, temperature를 꼭 부모 생성자로 넘겨줘야 해요.

def __init__(self, model: str, api_key: str, temperature: Optional[float] = None):
    # REQUIRED: Call parent constructor with model and temperature
    super().__init__(model=model, temperature=temperature)

    # Your custom initialization
    self.api_key = api_key

추상 메서드 call() — 커스텀 LLM 구현의 핵심입니다. 메시지(문자열 또는 role·content를 가진 dict 리스트)를 받아 문자열 응답을 돌려줘야 하고, 지원한다면 도구·함수 호출을 처리하며, 오류는 적절한 예외를 던져야 해요.

선택 메서드들 — 기본값이 있으니 필요할 때만 오버라이드하면 됩니다.

def supports_function_calling(self) -> bool:
    """Return True if your LLM supports function calling."""
    return True  # Default is True

def supports_stop_words(self) -> bool:
    """Return True if your LLM supports stop sequences."""
    return True  # Default is True

def get_context_window_size(self) -> int:
    """Return the context window size."""
    return 4096  # Default is 4096

흔한 패턴 (Common Patterns)

오류 처리 — 요청 타임아웃, 네트워크 오류, 응답 형식 오류를 각각 구분해서 던집니다.

import requests

def call(self, messages, tools=None, callbacks=None, available_functions=None):
    try:
        response = requests.post(
            self.endpoint,
            headers={"Authorization": f"Bearer {self.api_key}"},
            json=payload,
            timeout=30
        )
        response.raise_for_status()
        return response.json()["choices"][0]["message"]["content"]

    except requests.Timeout:
        raise TimeoutError("LLM request timed out")
    except requests.RequestException as e:
        raise RuntimeError(f"LLM request failed: {str(e)}")
    except (KeyError, IndexError) as e:
        raise ValueError(f"Invalid response format: {str(e)}")

커스텀 인증 — 인증 헤더 형식이 특이한 제공자를 지원할 수 있어요.

from crewai import BaseLLM
from typing import Optional

class CustomAuthLLM(BaseLLM):
    def __init__(self, model: str, auth_token: str, endpoint: str, temperature: Optional[float] = None):
        super().__init__(model=model, temperature=temperature)
        self.auth_token = auth_token
        self.endpoint = endpoint

    def call(self, messages, tools=None, callbacks=None, available_functions=None):
        headers = {
            "Authorization": f"Custom {self.auth_token}",  # Custom auth format
            "Content-Type": "application/json"
        }
        # Rest of implementation...

스톱워드 지원 — CrewAI는 에이전트 동작 제어를 위해 "\nObservation:"을 스톱워드로 자동 추가합니다. LLM이 스톱워드를 지원하면 API 호출에 stop을 포함하고, 지원하지 않으면 직접 잘라내야 해요.

def call(self, messages, tools=None, callbacks=None, available_functions=None):
    payload = {
        "model": self.model,
        "messages": messages,
        "stop": self.stop  # Include stop words in API call
    }
    # Make API call...

def supports_stop_words(self) -> bool:
    return True  # Your LLM supports stop sequences

스톱워드를 네이티브로 지원하지 않는다면 응답에서 직접 잘라냅니다.

def call(self, messages, tools=None, callbacks=None, available_functions=None):
    response = self._make_api_call(messages, tools)
    content = response["choices"][0]["message"]["content"]

    # Manually truncate at stop words
    if self.stop:
        for stop_word in self.stop:
            if stop_word in content:
                content = content.split(stop_word)[0]
                break

    return content

def supports_stop_words(self) -> bool:
    return False  # Tell CrewAI we handle stop words manually

함수 호출 (Function Calling)

LLM이 함수 호출을 지원한다면 전체 흐름을 구현합니다. 응답에 tool_calls가 있으면 available_functions에서 해당 함수를 찾아 실행하고, 결과를 메시지 이력에 담아 다시 LLM에 넘겨요.

import json

def call(self, messages, tools=None, callbacks=None, available_functions=None):
    # Convert string to message format
    if isinstance(messages, str):
        messages = [{"role": "user", "content": messages}]

    # Make API call
    response = self._make_api_call(messages, tools)
    message = response["choices"][0]["message"]

    # Check for function calls
    if "tool_calls" in message and available_functions:
        return self._handle_function_calls(
            message["tool_calls"], messages, tools, available_functions
        )

    return message["content"]

def _handle_function_calls(self, tool_calls, messages, tools, available_functions):
    """Handle function calling with proper message flow."""
    for tool_call in tool_calls:
        function_name = tool_call["function"]["name"]

        if function_name in available_functions:
            # Parse and execute function
            function_args = json.loads(tool_call["function"]["arguments"])
            function_result = available_functions[function_name](**function_args)

            # Add function call and result to message history
            messages.append({
                "role": "assistant",
                "content": None,
                "tool_calls": [tool_call]
            })
            messages.append({
                "role": "tool",
                "tool_call_id": tool_call["id"],
                "name": function_name,
                "content": str(function_result)
            })

            # Call LLM again with updated context
            return self.call(messages, tools, None, available_functions)

    return "Function call failed"

문제 해결 (Troubleshooting)

생성자 오류 — 부모 생성자에 필수 파라미터를 빠뜨리면 안 됩니다.

# ❌ Wrong - missing required parameters
def __init__(self, api_key: str):
    super().__init__()

# ✅ Correct
def __init__(self, model: str, api_key: str, temperature: Optional[float] = None):
    super().__init__(model=model, temperature=temperature)

함수 호출이 동작하지 않을 때

  • supports_function_calling()이 True를 반환하는지 확인
  • 응답에서 tool_calls를 처리하는지 확인
  • available_functions 파라미터가 올바르게 쓰였는지 확인

인증 실패 — API 키 형식과 권한, 인증 헤더 형식, 엔드포인트 URL을 점검합니다.

응답 파싱 오류 — 중첩 필드 접근 전에 응답 구조를 검증하고, content가 None일 수 있는 경우를 처리하며, 잘못된 응답에 대한 오류 처리를 추가해요.

커스텀 LLM 테스트하기

from crewai import Agent, Task, Crew

def test_custom_llm():
    llm = CustomLLM(
        model="test-model",
        api_key="test-key",
        endpoint="https://api.test.com"
    )

    # Test basic call
    result = llm.call("Hello, world!")
    assert isinstance(result, str)
    assert len(result) > 0

    # Test with CrewAI agent
    agent = Agent(
        role="Test Agent",
        goal="Test custom LLM",
        backstory="A test agent.",
        llm=llm
    )

    task = Task(
        description="Say hello",
        expected_output="A greeting",
        agent=agent
    )

    crew = Crew(agents=[agent], tasks=[task])
    result = crew.kickoff()
    assert "hello" in result.raw.lower()

더 알아보기