커스텀 함수 실행하기

커스텀 함수 실행하기 (How to run custom functions)

LCEL로 체인을 만들다 보면 LangChain 제공 컴포넌트에 없는 로직이 필요할 때가 있어요. 그럴 땐 아무 함수나 러너블(Runnable)로 쓰면 됩니다. 러너블로 쓰이는 커스텀 함수를 RunnableLambda 라고 불러요. 포맷팅이나 특정 가공이 필요할 때 특히 유용합니다.

출처: 공식문서

사전 준비

이 가이드를 따라가려면 다음 개념에 익숙해져 있으면 좋아요.

LangChain Expression Language (LCEL)
러너블 이어 붙이기 (Chaining runnables)

한 가지 중요한 규칙이 있어요. 이 함수들의 모든 입력은 단일 인자(SINGLE argument) 여야 한다는 점이에요. 여러 인자를 받는 함수라면, 단일 dict 입력을 받아 여러 인자로 풀어주는 래퍼(wrapper)를 만들어야 합니다.

이 가이드에서 다룰 내용이에요.

RunnableLambda 생성자와 편리한 @chain 데코레이터로 커스텀 함수에서 러너블을 명시적으로 만드는 법
체인 안에서 커스텀 함수가 러너블로 자동 변환(coercion)되는 법
커스텀 함수에서 run 메타데이터를 받아 쓰는 법
함수가 제너레이터를 반환하게 해서 스트리밍하는 법

생성자 사용하기

먼저 RunnableLambda 생성자로 커스텀 로직을 명시적으로 감싸볼게요.

%pip install -qU langchain langchain_openai

import os
from getpass import getpass

if "OPENAI_API_KEY" not in os.environ:
    os.environ["OPENAI_API_KEY"] = getpass()

from operator import itemgetter

from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableLambda
from langchain_openai import ChatOpenAI


def length_function(text):
    return len(text)


def _multiple_length_function(text1, text2):
    return len(text1) * len(text2)


def multiple_length_function(_dict):
    return _multiple_length_function(_dict["text1"], _dict["text2"])


model = ChatOpenAI()

prompt = ChatPromptTemplate.from_template("what is {a} + {b}")

chain = (
    {
        "a": itemgetter("foo") | RunnableLambda(length_function),
        "b": {"text1": itemgetter("foo"), "text2": itemgetter("bar")}
        | RunnableLambda(multiple_length_function),
    }
    | prompt
    | model
)

chain.invoke({"foo": "bar", "bar": "gah"})

두 개의 입력 인자를 받는 _multiple_length_function은 multiple_length_function이라는 단일 dict 입력 래퍼로 감싼 점에 주목하세요. 위 규칙처럼 커스텀 함수 입력은 단일 인자여야 하니까요.

편리한 `@chain` 데코레이터

임의의 함수를 @chain 데코레이터를 붙여서 체인으로 만들 수도 있어요. 기능적으로는 위에서 본 RunnableLambda 생성자로 감싼 것과 동일합니다. 예시를 볼게요.

from langchain_core.output_parsers import StrOutputParser
from langchain_core.runnables import chain

prompt1 = ChatPromptTemplate.from_template("Tell me a joke about {topic}")
prompt2 = ChatPromptTemplate.from_template("What is the subject of this joke: {joke}")


@chain
def custom_chain(text):
    prompt_val1 = prompt1.invoke({"topic": text})
    output1 = ChatOpenAI().invoke(prompt_val1)
    parsed_output1 = StrOutputParser().invoke(output1)
    chain2 = prompt2 | ChatOpenAI() | StrOutputParser()
    return chain2.invoke({"joke": parsed_output1})


custom_chain.invoke("bears")

위에서 @chain 데코레이터는 custom_chain을 러너블로 바꿔 주고, 우리는 그걸 .invoke() 메서드로 호출해요. LangSmith로 추적한다면 custom_chain 트레이스가 보이고, 그 아래에 OpenAI 호출들이 중첩되어 나타나는 걸 확인할 수 있어요.

체인에서의 자동 변환 (Coercion)

파이프 연산자(|)로 체인을 만들 때는 RunnableLambda나 @chain 생성자를 생략하고 자동 변환에 맡겨도 돼요. 모델 출력을 받아 첫 다섯 글자를 반환하는 함수를 쓰는 간단한 예시예요.

prompt = ChatPromptTemplate.from_template("tell me a story about {topic}")

model = ChatOpenAI()

chain_with_coerced_function = prompt | model | (lambda x: x.content[:5])

chain_with_coerced_function.invoke({"topic": "bears"})

파이프 연산자 왼쪽의 model이 이미 러너블이라 커스텀 함수 (lambda x: x.content[:5])를 RunnableLambda 생성자로 감쌀 필요가 없어요. 커스텀 함수는 자동으로 coercion되어 러너블이 됩니다. 자세한 내용은 러너블 이어 붙이기의 타입 강제 변환 섹션을 참고하세요.

run 메타데이터 전달하기

Runnable lambda는 선택적으로 RunnableConfig 파라미터를 받을 수 있어요. 이를 통해 중첩된 run에 콜백, 태그, 그 밖의 구성 정보를 전달할 수 있습니다.

import json

from langchain_core.runnables import RunnableConfig


def parse_or_fix(text: str, config: RunnableConfig):
    fixing_chain = (
        ChatPromptTemplate.from_template(
            "Fix the following text:\n\n```text\n{input}\n```\nError: {error}"
            " Don't narrate, just respond with the fixed data."
        )
        | model
        | StrOutputParser()
    )
    for _ in range(3):
        try:
            return json.loads(text)
        except Exception as e:
            text = fixing_chain.invoke({"input": text, "error": e}, config)
    return "Failed to parse"


from langchain_community.callbacks import get_openai_callback

with get_openai_callback() as cb:
    output = RunnableLambda(parse_or_fix).invoke(
        "{foo: bar}", {"tags": ["my-tag"], "callbacks": [cb]}
    )
    print(output)
    print(cb)

from langchain_community.callbacks import get_openai_callback

with get_openai_callback() as cb:
    output = RunnableLambda(parse_or_fix).invoke(
        "{foo: bar}", {"tags": ["my-tag"], "callbacks": [cb]}
    )
    print(output)
    print(cb)

parse_or_fix 함수가 config: RunnableConfig 인자를 받고, 그 config를 fixing_chain.invoke(..., config)로 다시 전달하는 걸 볼 수 있어요. 이렇게 해서 with get_openai_callback() as cb로 잡은 콜백이 함수 안의 중첩된 LLM 호출에도 전달됩니다.

스트리밍

참고: RunnableLambda는 스트리밍을 지원하지 않아도 되는 코드에 가장 잘 맞아요. 스트리밍을 지원해야 한다면(즉 입력 청크들을 처리하고 출력 청크들을 산출해야 한다면) 아래 예시처럼 RunnableGenerator를 쓰세요.

체인 안에서 제너레이터 함수(즉 yield 키워드를 쓰고 이터레이터처럼 동작하는 함수)를 쓸 수 있어요. 이런 제너레이터의 시그니처는 Iterator[Input] -> Iterator[Output]이고, 비동기 제너레이터는 AsyncIterator[Input] -> AsyncIterator[Output]이에요.

이건 다음과 같은 경우에 유용해요.

커스텀 출력 파서 구현
스트리밍 능력을 유지하면서 이전 단계의 출력 수정

쉼표로 구분된 리스트를 만드는 커스텀 출력 파서 예시를 볼게요. 먼저 그런 리스트를 텍스트로 생성하는 체인을 만들어요.

from typing import Iterator, List

prompt = ChatPromptTemplate.from_template(
    "Write a comma-separated list of 5 animals similar to: {animal}. Do not include numbers"
)

str_chain = prompt | model | StrOutputParser()

for chunk in str_chain.stream({"animal": "bear"}):
    print(chunk, end="", flush=True)

다음으로, 현재까지 스트리밍된 출력을 모았다가 모델이 리스트에서 다음 쉼표를 만들 때 산출하는 커스텀 함수를 정의해요.

# This is a custom parser that splits an iterator of llm tokens
# into a list of strings separated by commas
def split_into_list(input: Iterator[str]) -> Iterator[List[str]]:
    # hold partial input until we get a comma
    buffer = ""
    for chunk in input:
        # add current chunk to buffer
        buffer += chunk
        # while there are commas in the buffer
        while "," in buffer:
            # split buffer on comma
            comma_index = buffer.index(",")
            # yield everything before the comma
            yield [buffer[:comma_index].strip()]
            # save the rest for the next iteration
            buffer = buffer[comma_index + 1 :]
    # yield the last chunk
    yield [buffer.strip()]


list_chain = str_chain | split_into_list

for chunk in list_chain.stream({"animal": "bear"}):
    print(chunk, flush=True)

invoke 하면 전체 값 배열을 돌려받아요.

list_chain.invoke({"animal": "bear"})

비동기 버전

async 환경에서 작업한다면, 위 예시의 비동기 버전은 이렇게 생겼어요.

from typing import AsyncIterator


async def asplit_into_list(
    input: AsyncIterator[str],
) -> AsyncIterator[List[str]]:  # async def
    buffer = ""
    async for (
        chunk
    ) in input:  # `input` is a `async_generator` object, so use `async for`
        buffer += chunk
        while "," in buffer:
            comma_index = buffer.index(",")
            yield [buffer[:comma_index].strip()]
            buffer = buffer[comma_index + 1 :]
    yield [buffer.strip()]


list_chain = str_chain | asplit_into_list

async for chunk in list_chain.astream({"animal": "bear"}):
    print(chunk, flush=True)

await list_chain.ainvoke({"animal": "bear"})

input이 async_generator 객체이므로 async for로 순회하는 점이 달라요. str_chain의 결과를 asplit_into_list에 |로 이어 붙이고, astream/ainvoke로 호출합니다.

다음 단계

이제 체인 안에서 커스텀 로직을 쓰는 몇 가지 방법과 스트리밍 구현 방법을 배웠어요. 더 배우고 싶다면 이 섹션의 다른 러너블 how-to 가이드를 확인해 보세요.

커스텀 함수 실행하기