[langgraph] 그래프에서 LLM 토큰을 스트리밍하는 방법

LangGraph 공식문서를 번역한 내용입니다. 필요한 경우 부연 설명을 추가하였고 이해하기 쉽게 예제를 일부 변경하였습니다. 문제가 되면 삭제하겠습니다.

https://langchain-ai.github.io/langgraph/how-tos/streaming-tokens/

이 예제에서는 에이전트를 구동하는 언어 모델에서 토큰을 스트리밍한다. 예제로는 ReAct 에이전트를 사용한다.

이 가이드는 이 디렉터리의 다른 가이드들과 밀접하게 연관되므로, 아래에서 STREAMING 태그를 사용하여 차이점을 설명할 것이다(그 부분만 검색하고 싶다면 참고하라).

준비

우선, 필요한 패키지를 설치하자.

pip install langgraph langchain_openai langsmith

상태 준비

LangGraph에서 주요 그래프 유형은 StateGraph이다. 이 그래프는 각 노드에 전달되는 State 객체로 매개변수화된다. 각 노드는 그 상태를 업데이트하는 데 사용되는 연산을 반환한다. 이러한 연산은 상태에 특정 속성을 SET하여 기존 값을 덮어쓰거나, 기존 속성에 ADD하여 값을 추가할 수 있다. SET과 ADD의 여부는 그래프를 구성하는 데 사용하는 State 객체에 Annotated를 달아서 표시한다.

이 예제에서는 상태가 메시지 목록이다. 각 노드는 그 목록에 메시지만 추가되므로, 하나의 키(messages)를 가진 TypedDict를 사용하고, 메시지 속성이 "append-only"로 Annotated로 처리된다.

from typing import Annotated

from dotenv import load_dotenv
from typing_extensions import TypedDict

from langgraph.graph.message import add_messages

load_dotenv()


class State(TypedDict):
    messages: Annotated[list, add_messages]

도구 준비

먼저 사용할 도구를 정의한다. 이 예제에서는 Placeholder 검색 엔진을 만들 것이다. 도구를 만드는 것은 매우 쉽다. 어떻게 하는지에 대한 문서는 여기에서 확인할 수 있다.

from langchain_core.tools import tool


@tool
def search(query: str):
    """Call to surf the web."""
    return ["흐려요~"]


tools = [search]

이제 이러한 도구들을 간단한 ToolNode로 래핑할 수 있다. 이 클래스는 도구 호출이 포함된 AIMessages 목록을 입력받아 도구를 실행하고, 결과를 ToolMessages로 반환하는 간단한 클래스이다.

from langgraph.prebuilt import ToolNode

tool_node = ToolNode(tools)

모델 준비

이제 사용할 채팅 모델을 로드해야 한다. 이 모델은 두 가지 기준을 충족해야 한다.

메시지와 함께 작동해야 한다. 왜냐하면 상태는 주로 메시지 목록(채팅 기록)이기 때문이다.
도구 호출과 함께 작동해야 한다. 미리 구축된 ToolNode를 사용하고 있기 때문이다.

참고: 이러한 모델 요구 사항은 LangGraph를 사용하는 데 필수적인 요구 사항이 아니라, 이 특정 예제에서만 필요한 요구 사항이다.

from langchain_openai import ChatOpenAI

model = ChatOpenAI(model="gpt-4o-mini")

이 작업을 마친 후, 모델이 이러한 도구를 호출할 수 있다는 것을 인식하도록 해야 한다. 이를 위해 LangChain 도구들을 함수 호출 형식으로 변환한 후, 모델 클래스에 바인딩할 수 있다.

model = model.bind_tools(tools)

노드 정의

이제 그래프에서 몇 가지 다른 노드를 정의해야 한다. LangGraph에서 노드는 함수나 실행 가능한 객체일 수 있다. 이 예제에서는 두 가지 주요 노드가 필요하다.

에이전트: 어떤 행동을 취할지(있다면)를 결정하는 역할을 한다.
도구를 호출하는 함수: 에이전트가 행동을 취할 것이라고 결정하면, 이 노드는 그 행동을 실행한다.

또한 몇 가지 엣지를 정의해야 한다. 이 엣지들은 조건부일 수 있다. 조건부인 이유는 노드의 출력에 따라 여러 경로 중 하나를 선택해야 하기 때문이다. 경로는 그 노드가 실행되기 전까지 알 수 없다(LLM이 결정한다).

조건부 엣지: 에이전트가 호출된 후, 다음을 결정한다: a. 에이전트가 행동을 취하라고 하면, 도구를 호출하는 함수가 실행된다. b. 에이전트가 종료했다고 말하면, 프로세스가 종료된다.
일반 엣지: 도구가 호출된 후에는 항상 에이전트로 돌아가서 다음에 무엇을 할지 결정한다.

이제 노드를 정의하고, 어떤 조건부 엣지를 선택할지 결정하는 함수를 정의한다.

STREAMING

각 노드는 비동기 함수로 정의된다.

from langchain_core.runnables import RunnableConfig

from langgraph.graph import END, START, StateGraph


def should_continue(state: State):
    messages = state["messages"]
    last_message = messages[-1]
    if not last_message.tool_calls:
        return END
    else:
        return "tools"


async def call_model(state: State, config: RunnableConfig):
    messages = state["messages"]
    response = await model.ainvoke(messages, config)
    return {"messages": response}

그래프 정의

이제 모든 것을 결합하여 그래프를 정의할 수 있다.

workflow = StateGraph(State)

workflow.add_node("agent", call_model)
workflow.add_node("tools", tool_node)

workflow.add_edge(START, "agent")

workflow.add_conditional_edges(
    "agent",
    should_continue,
    ["tools", END],
)

workflow.add_edge("tools", "agent")

app = workflow.compile()

from IPython.display import Image, display

display(
    Image(
        app.get_graph().draw_mermaid_png(
            output_file_path="how-to-stream-llm-tokens.png"
        )
    )
)

LLM token 스트리밍

각 노드가 생성하는 LLM 토큰에 접근할 수 있다. 이 경우 "agent" 노드만 LLM 토큰을 생성한다. 이 기능이 제대로 작동하려면, 스트리밍을 지원하는 LLM을 사용해야 하며, LLM을 구성할 때 스트리밍 설정을 해야 한다 (예: ChatOpenAI(model="gpt-3.5-turbo-1106", streaming=True)).

from langchain_core.messages import AIMessageChunk, HumanMessage

inputs = [HumanMessage(content="서울 날씨 어때?")]


async def stream_async():
    first = True
    async for msg, metadata in app.astream(
        {"messages": inputs}, stream_mode="messages"
    ):
        if msg.content and not isinstance(msg, HumanMessage):
            print(msg.content, end="|", flush=True)

        if isinstance(msg, AIMessageChunk):
            if first:
                gathered = msg
                first = False
            else:
                gathered = gathered + msg

            if msg.tool_call_chunks:
                print(gathered.tool_calls)


asyncio.run(stream_async())

[{'name': 'search', 'args': {}, 'id': 'call_fRTMOpDtiHT34g51d2Jno32P', 'type': 'tool_call'}]
[{'name': 'search', 'args': {}, 'id': 'call_fRTMOpDtiHT34g51d2Jno32P', 'type': 'tool_call'}]
[{'name': 'search', 'args': {}, 'id': 'call_fRTMOpDtiHT34g51d2Jno32P', 'type': 'tool_call'}]
[{'name': 'search', 'args': {'query': ''}, 'id': 'call_fRTMOpDtiHT34g51d2Jno32P', 'type': 'tool_call'}]
[{'name': 'search', 'args': {'query': '서'}, 'id': 'call_fRTMOpDtiHT34g51d2Jno32P', 'type': 'tool_call'}]
[{'name': 'search', 'args': {'query': '서울'}, 'id': 'call_fRTMOpDtiHT34g51d2Jno32P', 'type': 'tool_call'}]
[{'name': 'search', 'args': {'query': '서울 날'}, 'id': 'call_fRTMOpDtiHT34g51d2Jno32P', 'type': 'tool_call'}]
[{'name': 'search', 'args': {'query': '서울 날씨'}, 'id': 'call_fRTMOpDtiHT34g51d2Jno32P', 'type': 'tool_call'}]
[{'name': 'search', 'args': {'query': '서울 날씨'}, 'id': 'call_fRTMOpDtiHT34g51d2Jno32P', 'type': 'tool_call'}]
["흐려요~"]|서|울|의| 날|씨|는| 흐|린| 상|태|입니다|.|

LangGraph 참고 자료

저작자표시

준비