09-langchain-streaming

2024. 1. 15. 21:00langchain

스트리밍

LLM의 경우 스트리밍이 점점 더 인기 있는 기능이 되었습니다. 전체 응답이 생성될 때까지 기다리는 대신 아이디어는 LLM이 토큰을 생성할 때 토큰을 신속하게 반환하는 것입니다.

스트리밍은 실제로 간단한 사용 사례에서는 구현하기가 매우 쉽지만 스트리밍 시도를 차단할 수 있는 자체 논리가 실행되는 에이전트와 같은 항목을 포함하기 시작하면 복잡해질 수 있습니다. 다행히도 우리는 그것을 작동시킬 수 있어며, 단지 약간의 추가 노력이 필요할 뿐입니다.

LLM용 터미널에 스트리밍을 구현하는 것으로 쉽게 시작하지만, 노트북이 끝날 무렵에는 에이전트용 FastAPI를 통해 스트리밍하는 보다 복잡한 작업을 처리하게 됩니다.

In [ ]:
!pip install -qU \
    openai==0.28.0 \
    langchain==0.0.301 \
    fastapi==0.103.1 \
    "uvicorn[standard]"==0.23.2

Stdout으로 LLM 스트리밍

스트리밍의 가장 간단한 형태는 간단히 생성된 토큰을 "인쇄"하는 것입니다. 이를 설정하려면 두 가지 특정 매개변수를 사용하여 LLM(스트리밍을 지원하는 LLM)을 초기화해야 합니다.

  • streaming=True, 스트리밍을 활성화하려면
  • callbacks=[SomeCallBackHere()], 여기서 LangChain 콜백 클래스(또는 여러 개를 포함하는 목록)를 전달합니다.

streaming 매개변수는 설명이 필요합니다. callbacks 매개변수와 콜백 클래스는 약간의 설명이 필요합니다. 기본적으로 LLM의 각 토큰이 생성될 때 작업을 수행하는 작은 코드 역할을 합니다. 앞서 언급했듯이 가장 간단한 형태의 스트리밍은 StreamingStdOutCallbackHandler.와 같이 생성되는 토큰을 인쇄하는 것입니다.

In [ ]:
import os
from langchain.chat_models import ChatOpenAI
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler

os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY") or "YOUR_API_KEY"

llm = ChatOpenAI(
    openai_api_key=os.getenv("OPENAI_API_KEY"),
    temperature=0.0,
    model_name="gpt-3.5-turbo",
    streaming=True,  # ! important
    callbacks=[StreamingStdOutCallbackHandler()]  # ! important
)

이제 LLM을 실행하면 응답이 스트리밍되는 것을 볼 수 있습니다.

In [ ]:
from langchain.schema import HumanMessage

# create messages to be passed to chat LLM
messages = [HumanMessage(content="tell me a long story")]

llm(messages)

놀라울 정도로 쉬웠지만 에이전트를 사용하기 시작하자마자 상황은 훨씬 더 복잡해지기 시작했습니다. 먼저 에이전트를 초기화해 보겠습니다.

In [ ]:
from langchain.memory import ConversationBufferWindowMemory
from langchain.agents import load_tools, AgentType, initialize_agent

# initialize conversational memory
memory = ConversationBufferWindowMemory(
    memory_key="chat_history",
    k=5,
    return_messages=True,
    output_key="output"
)

# create a single tool to see how it impacts streaming
tools = load_tools(["llm-math"], llm=llm)

# initialize the agent
agent = initialize_agent(
    agent=AgentType.CHAT_CONVERSATIONAL_REACT_DESCRIPTION,
    tools=tools,
    llm=llm,
    memory=memory,
    verbose=True,
    max_iterations=3,
    early_stopping_method="generate",
    return_intermediate_steps=False
)

콜백으로 생성한 것과 동일한 llm으로 에이전트를 초기화하면서 StreamingStdOutCallbackHandler을 에이전트에 이미 추가했습니다. 그럼 에이전트를 실행할 때 무엇을 얻는지 살펴보겠습니다.

In [ ]:
prompt = "Hello, how are you?"

agent(prompt)

적절하지는 않지만 이제 LLM의 전체 출력을 스트리밍하는 문제가 있습니다. 우리는 에이전트를 사용하고 있기 때문에 LLM은 여기서 볼 수 있는 JSON 형식을 출력하라는 지시를 받아 에이전트 로직이 도구 사용, 다중 "사고" 작업 등을 처리할 수 있습니다. 예를 들어, 수학 질문을 하면 다음과 같이 표시됩니다.

In [ ]:
agent("what is the square root of 71?")

개발 중에 보는 것도 흥미롭지만 실제 사용 사례에서는 이 스트리밍을 약간 정리하고 싶습니다. 이를 위해 우리는 두 가지 접근 방식을 사용할 수 있습니다. 즉, 사용자 정의 콜백 핸들러를 구축하거나 LangChain의 목적에 맞게 구축된 콜백 핸들러를 사용하는 것입니다(보통 LangChain에는 필요한 모든 것이 있습니다). 먼저 LangChain이 특별히 제작한 FinalStreamingStdOutCallbackHandler.을 사용해 보겠습니다.

In [ ]:
agent.agent.llm_chain.llm

새로운 콜백 핸들러를 사용하면 다음과 같습니다.

In [ ]:
from langchain.callbacks.streaming_stdout_final_only import (
    FinalStreamingStdOutCallbackHandler,
)

agent.agent.llm_chain.llm.callbacks = [
    FinalStreamingStdOutCallbackHandler(
        answer_prefix_tokens=["Final", "Answer"]
    )
]

시도해 봅시다.

In [ ]:
agent("what is the square root of 71?")

아닙니다. answer_prefix_tokens 인수를 정리해야 하지만 제대로 하기가 어렵습니다. 일반적으로 다음과 같이 맞춤 콜백 핸들러를 사용하는 것이 더 쉽습니다.

In [ ]:
import sys

class CallbackHandler(StreamingStdOutCallbackHandler):
    def __init__(self):
        self.content: str = ""
        self.final_answer: bool = False

    def on_llm_new_token(self, token: str, **kwargs: any) -> None:
        self.content += token
        if "Final Answer" in self.content:
            # now we're in the final answer section, but don't print yet
            self.final_answer = True
            self.content = ""
        if self.final_answer:
            if '"action_input": "' in self.content:
                if token not in ["}"]:
                    sys.stdout.write(token)  # equal to `print(token, end="")`
                    sys.stdout.flush()

agent.agent.llm_chain.llm.callbacks = [CallbackHandler()]

다시 시도해 보세요.

In [ ]:
agent("what is the square root of 71?")
In [ ]:
agent.agent.llm_chain.llm

완벽하지는 않지만 점점 좋아지고 있습니다. 이제 대부분의 시나리오에서는 단순히 터미널이나 노트북에 출력을 인쇄하지는 않을 것입니다. 이 데이터를 다른 API를 통해 스트리밍하는 등 좀 더 복잡한 작업을 수행하려면 작업을 다르게 수행해야 합니다.

에이전트와 함께 FastAPI 사용

대부분의 경우 우리는 API와 같은 것 뒤에 LLM, 에이전트 등을 배치합니다. 이를 믹스에 추가하고 FastAPI를 사용하여 에이전트에 대한 스트리밍을 구현하는 방법을 살펴보겠습니다.

먼저 FastAPI 논리를 포함하는 간단한 main.py 스크립트를 만듭니다. 이 노트북과 동일한 GitHub 저장소 위치에서 찾을 수 있습니다(여기 링크(https://github.com/pinecone-io/examples/blob/master/learn/generation/langchain/handbook/09-langchain-streaming/main.py)).

API를 실행하려면 해당 디렉토리로 이동하여 uvicorn main:app --reload를 실행하세요. 완료되면 다음 셀 출력에서 ​​🤙 상태를 찾아 실행 중인지 확인할 수 있습니다.

In [ ]:
import requests

res = requests.get("http://localhost:8000/health")
res.json()
In [ ]:
res = requests.get("http://localhost:8000/chat",
    json={"text": "hello there!"}
)
res
In [ ]:
res.json()

StdOut 스트리밍과 달리 이제 StreamingResponse 객체를 통해 해당 토큰을 FastAPI에 공급하는 생성기 함수에 토큰을 보내야 합니다. 이를 처리하려면 비동기 코드를 사용해야 합니다. 그렇지 않으면 생성기가 이후 생성이 이미 완료될 때까지 아무것도 방출하지 않습니다.

Queue은 각 토큰이 생성될 때 토큰을 대기열에 넣기 때문에 콜백 핸들러에 의해 액세스됩니다. 생성기 기능은 대기열에 추가되는 새 토큰을 비동기식으로 확인합니다. 생성기는 토큰이 추가된 것을 확인하자마자 토큰을 가져와 StreamingResponse.에 제공합니다.

실제 동작을 확인하기 위해 get_stream라는 스트림 요청 함수를 정의하겠습니다.

In [ ]:
def get_stream(query: str):
    s = requests.Session()
    with s.get(
        "http://localhost:8000/chat",
        stream=True,
        json={"text": query}
    ) as r:
        for line in r.iter_content():
            print(line.decode("utf-8"), end="")
In [ ]:
get_stream("hi there!")
In [ ]:
출처 : https://github.com/pinecone-io/examples/blob/master/learn/generation/langchain/handbook/09-langchain-streaming/09-langchain-streaming.ipynb

'langchain' 카테고리의 다른 글

07-langchain-tools  (0) 2024.01.17
03-1-langchain-conversational-memory  (1) 2024.01.16
06-langchain-agents  (1) 2024.01.14
04-langchain-chat  (0) 2024.01.13
03a-token-counter  (0) 2024.01.12