LangChain에서 출력 스트리밍하기

스트리밍이 중요한 이유

.stream()을 사용한 기본 스트리밍

from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser

chain = (
    ChatPromptTemplate.from_template('Explain {topic} in detail.')
    | ChatOpenAI(model='gpt-4o-mini')
    | StrOutputParser()
)

# Stream tokens to stdout
for chunk in chain.stream({'topic': 'quantum entanglement'}):
    print(chunk, end='', flush=True)
print()  # final newline

.astream()을 사용한 비동기 스트리밍

import asyncio
from langchain_openai import ChatOpenAI
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate

chain = (
    ChatPromptTemplate.from_template('Write a poem about {subject}')
    | ChatOpenAI(model='gpt-4o-mini')
    | StrOutputParser()
)

async def stream_response():
    async for chunk in chain.astream({'subject': 'the ocean'}):
        print(chunk, end='', flush=True)

asyncio.run(stream_response())

StreamingResponse를 사용한 FastAPI 스트리밍

from fastapi import FastAPI
from fastapi.responses import StreamingResponse

app = FastAPI()

async def generate_stream(topic: str):
    async for chunk in chain.astream({'topic': topic}):
        yield chunk

@app.get('/stream')
async def stream_endpoint(topic: str):
    return StreamingResponse(
        generate_stream(topic),
        media_type='text/plain'
    )

# SSE format for frontend EventSource
async def sse_stream(topic: str):
    async for chunk in chain.astream({'topic': topic}):
        yield f'data: {chunk}\n\n'

중간 단계를 거치는 스트리밍

from langchain_core.output_parsers import JsonOutputParser

# This chain does NOT stream token by token
# JsonOutputParser must buffer the full response before parsing JSON
json_chain = (
    ChatPromptTemplate.from_template('Return JSON: {task}')
    | ChatOpenAI(model='gpt-4o-mini')
    | JsonOutputParser()  # buffers until complete
)

# But partial JSON streaming IS possible with streaming_json_parser
for partial in json_chain.stream({'task': 'list 3 colors'}):
    print(partial)  # prints partial dict as it fills in

세밀한 제어를 위한 astream_events

async def stream_with_events(question: str):
    async for event in chain.astream_events(
        {'question': question},
        version='v2'
    ):
        kind = event['event']
        if kind == 'on_llm_stream':
            chunk = event['data']['chunk'].content
            print(chunk, end='', flush=True)
        elif kind == 'on_chain_end':
            print('\n[Done]')
        elif kind == 'on_tool_start':
            print(f'\n[Tool: {event["name"]}]')

스트리밍 출력 버퍼링

async def stream_and_capture(question: str) -> str:
    full_response = []
    async for chunk in chain.astream({'question': question}):
        print(chunk, end='', flush=True)  # stream to user
        full_response.append(chunk)        # also collect
    print()  # newline
    complete = ''.join(full_response)
    await log_response(question, complete)  # log full text
    return complete

도구 호출을 사용한 스트리밍

from openai import AsyncOpenAI

client = AsyncOpenAI()

async def stream_with_tools(prompt: str):
    tool_call_buffer = {}
    async with client.chat.completions.stream(
        model='gpt-4o-mini',
        messages=[{'role': 'user', 'content': prompt}],
        tools=[weather_tool_schema]
    ) as stream:
        async for chunk in stream:
            delta = chunk.choices[0].delta
            if delta.tool_calls:
                for tc in delta.tool_calls:
                    idx = tc.index
                    if idx not in tool_call_buffer:
                        tool_call_buffer[idx] = ''
                    if tc.function.arguments:
                        tool_call_buffer[idx] += tc.function.arguments

스트리밍 취소와 시간 제한

import asyncio

async def stream_with_timeout(question: str, timeout: float = 30.0):
    async def _stream():
        async for chunk in chain.astream({'question': question}):
            yield chunk

    try:
        async for chunk in asyncio.timeout(_stream(), timeout):
            print(chunk, end='', flush=True)
    except asyncio.TimeoutError:
        print('\n[Stream timed out after 30 seconds]')
    except asyncio.CancelledError:
        print('\n[Stream cancelled by client disconnect]')

JavaScript를 사용한 클라이언트 측 SSE

// Frontend JavaScript (not Python)
const source = new EventSource('/stream?topic=quantum+computing');
const outputDiv = document.getElementById('output');

source.onmessage = (event) => {
    outputDiv.textContent += event.data;
};

source.onerror = () => {
    source.close();
    outputDiv.textContent += ' [done]';
};

// Alternative: fetch with ReadableStream
const response = await fetch('/stream?topic=ai');
const reader = response.body.getReader();
while (true) {
    const {done, value} = await reader.read();
    if (done) break;
    outputDiv.textContent += new TextDecoder().decode(value);
}

스트리밍 모범 사례

# Complete SSE endpoint with DONE sentinel
async def sse_generator(question: str):
    try:
        async for chunk in chain.astream({'question': question}):
            # Escape any newlines in the chunk
            safe_chunk = chunk.replace('\n', ' ')
            yield f'data: {safe_chunk}\n\n'
    finally:
        yield 'data: [DONE]\n\n'

@app.get('/chat/stream')
async def chat_stream(question: str):
    return StreamingResponse(
        sse_generator(question),
        media_type='text/event-stream',
        headers={'Cache-Control': 'no-cache', 'X-Accel-Buffering': 'no'}
    )

빠른 확인

레슨 요약