6장: 스트리밍 응답 인터페이스 설계

event: message
data: {"key": "value"}
id: msg-001
retry: 3000
 
event: message
data: {"key": "value2"}
id: msg-002
 
event: done
data: [DONE]

from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
import json
import asyncio
 
 
@app.post("/api/v1/chat/completions")
async def create_chat_completion(
    request: CompletionRequest,
    http_request: Request,
):
    if not request.stream:
        return await sync_completion(request)
    
    return StreamingResponse(
        stream_completion(request, http_request),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-cache",
            "Connection": "keep-alive",
            "X-Accel-Buffering": "no",  # Nginx 버퍼링 비활성화
        },
    )
 
 
async def stream_completion(
    request: CompletionRequest,
    http_request: Request,
):
    request_id = f"chatcmpl-{uuid.uuid4().hex[:12]}"
    total_prompt_tokens = count_tokens(request.messages)
    total_completion_tokens = 0
    
    try:
        async for token in inference_engine.stream(request):
            # 클라이언트 연결 확인
            if await http_request.is_disconnected():
                break
            
            total_completion_tokens += 1
            
            chunk = {
                "id": request_id,
                "object": "chat.completion.chunk",
                "created": int(time.time()),
                "model": request.model,
                "choices": [{
                    "index": 0,
                    "delta": {"content": token},
                    "finish_reason": None,
                }],
            }
            
            yield f"data: {json.dumps(chunk, ensure_ascii=False)}\n\n"
        
        # 완료 청크
        final_chunk = {
            "id": request_id,
            "object": "chat.completion.chunk",
            "created": int(time.time()),
            "model": request.model,
            "choices": [{
                "index": 0,
                "delta": {},
                "finish_reason": "stop",
            }],
            "usage": {
                "prompt_tokens": total_prompt_tokens,
                "completion_tokens": total_completion_tokens,
                "total_tokens": total_prompt_tokens + total_completion_tokens,
            },
        }
        yield f"data: {json.dumps(final_chunk, ensure_ascii=False)}\n\n"
        
        # 종료 신호
        yield "data: [DONE]\n\n"
        
    except Exception as e:
        error_chunk = {
            "error": {
                "type": "server_error",
                "message": str(e),
                "code": "inference_error",
            }
        }
        yield f"data: {json.dumps(error_chunk)}\n\n"
        yield "data: [DONE]\n\n"

// 일반 텍스트 스트리밍 청크
interface StreamChunk {
  id: string;
  object: "chat.completion.chunk";
  created: number;
  model: string;
  choices: StreamChoice[];
  usage?: TokenUsage;  // 마지막 청크에서만 포함
}
 
interface StreamChoice {
  index: number;
  delta: MessageDelta;
  finish_reason: string | null;
}
 
interface MessageDelta {
  role?: string;     // 첫 번째 청크에서만
  content?: string;  // 토큰 내용
  tool_calls?: ToolCallDelta[];
}
 
// 도구 호출 스트리밍
interface ToolCallDelta {
  index: number;
  id?: string;       // 첫 번째 청크에서만
  type?: string;     // 첫 번째 청크에서만
  function?: {
    name?: string;   // 첫 번째 청크에서만
    arguments?: string;  // 점진적으로 누적
  };
}

data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1711234567,"model":"claude-4","choices":[{"index":0,"delta":{"role":"assistant","content":""},"finish_reason":null}]}
 
data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1711234567,"model":"claude-4","choices":[{"index":0,"delta":{"content":"API"},"finish_reason":null}]}
 
data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1711234567,"model":"claude-4","choices":[{"index":0,"delta":{"content":" 설계"},"finish_reason":null}]}
 
data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1711234567,"model":"claude-4","choices":[{"index":0,"delta":{"content":"에서"},"finish_reason":null}]}
 
data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1711234567,"model":"claude-4","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":25,"completion_tokens":8,"total_tokens":33}}
 
data: [DONE]

data: {"id":"chatcmpl-def","choices":[{"index":0,"delta":{"role":"assistant","content":null,"tool_calls":[{"index":0,"id":"call_xyz","type":"function","function":{"name":"get_weather","arguments":""}}]},"finish_reason":null}]}
 
data: {"id":"chatcmpl-def","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"{\"ci"}}]},"finish_reason":null}]}
 
data: {"id":"chatcmpl-def","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":"ty\":"}}]},"finish_reason":null}]}
 
data: {"id":"chatcmpl-def","choices":[{"index":0,"delta":{"tool_calls":[{"index":0,"function":{"arguments":" \"서울\"}"}}]},"finish_reason":null}]}
 
data: {"id":"chatcmpl-def","choices":[{"index":0,"delta":{},"finish_reason":"tool_calls"}]}
 
data: [DONE]

class ToolCallAccumulator {
  private toolCalls: Map<number, {
    id: string;
    type: string;
    function: { name: string; arguments: string };
  }> = new Map();
 
  processDelta(delta: MessageDelta): void {
    if (!delta.tool_calls) return;
 
    for (const tc of delta.tool_calls) {
      if (!this.toolCalls.has(tc.index)) {
        this.toolCalls.set(tc.index, {
          id: tc.id ?? "",
          type: tc.type ?? "function",
          function: {
            name: tc.function?.name ?? "",
            arguments: "",
          },
        });
      }
 
      const existing = this.toolCalls.get(tc.index)!;
      if (tc.function?.arguments) {
        existing.function.arguments += tc.function.arguments;
      }
    }
  }
 
  getToolCalls(): ToolCall[] {
    return Array.from(this.toolCalls.values()).map((tc) => ({
      ...tc,
      function: {
        ...tc.function,
        arguments: JSON.parse(tc.function.arguments),
      },
    }));
  }
}

async def stream_completion(
    request: CompletionRequest,
    http_request: Request,
):
    """클라이언트 연결 해제 시 추론을 즉시 중단합니다."""
    inference_task = None
    
    try:
        async for token in inference_engine.stream(request):
            # 주기적으로 클라이언트 연결 확인
            if await http_request.is_disconnected():
                logger.info(
                    f"클라이언트 연결 해제: {request.model}, "
                    f"생성된 토큰: {token_count}"
                )
                # 추론 리소스 즉시 해제
                await inference_engine.cancel(request_id)
                return
            
            yield format_chunk(token)
            
    except asyncio.CancelledError:
        logger.info("스트리밍 취소됨")
        await inference_engine.cancel(request_id)
    except Exception as e:
        yield format_error_chunk(e)
    finally:
        # 사용량 기록 (부분 응답도 과금)
        await usage_tracker.record(
            request_id=request_id,
            prompt_tokens=prompt_tokens,
            completion_tokens=token_count,
        )

class StreamingClient {
  private controller: AbortController | null = null;
 
  async streamCompletion(
    request: CompletionRequest,
    onToken: (token: string) => void,
    onComplete: (usage: TokenUsage) => void,
    onError: (error: Error) => void,
  ): Promise<void> {
    this.controller = new AbortController();
    const signal = this.controller.signal;
 
    // 타임아웃 설정 (60초)
    const timeoutId = setTimeout(() => this.cancel(), 60000);
 
    try {
      const response = await fetch("/api/v1/chat/completions", {
        method: "POST",
        headers: { "Content-Type": "application/json" },
        body: JSON.stringify({ ...request, stream: true }),
        signal,
      });
 
      if (!response.ok) {
        throw new Error(`HTTP ${response.status}`);
      }
 
      const reader = response.body!.getReader();
      const decoder = new TextDecoder();
      let buffer = "";
 
      while (true) {
        const { done, value } = await reader.read();
        if (done) break;
 
        buffer += decoder.decode(value, { stream: true });
        const lines = buffer.split("\n");
        buffer = lines.pop() ?? "";
 
        for (const line of lines) {
          if (!line.startsWith("data: ")) continue;
          const data = line.slice(6);
 
          if (data === "[DONE]") {
            return;
          }
 
          const chunk = JSON.parse(data);
          const delta = chunk.choices[0]?.delta;
 
          if (delta?.content) {
            onToken(delta.content);
          }
 
          if (chunk.usage) {
            onComplete(chunk.usage);
          }
        }
      }
    } catch (error) {
      if (signal.aborted) return; // 의도적 취소
      onError(error as Error);
    } finally {
      clearTimeout(timeoutId);
      this.controller = null;
    }
  }
 
  cancel(): void {
    this.controller?.abort();
  }
}

@app.post("/api/v1/chat/completions")
async def create_streaming_completion(
    request: CompletionRequest,
    http_request: Request,
):
    last_event_id = http_request.headers.get("Last-Event-ID")
    
    if last_event_id:
        # 이전 스트리밍 세션 이어서 전송
        return StreamingResponse(
            resume_stream(request, last_event_id),
            media_type="text/event-stream",
        )
    
    return StreamingResponse(
        stream_completion(request, http_request),
        media_type="text/event-stream",
    )
 
 
async def stream_completion(request, http_request):
    event_id = 0
    
    async for token in inference_engine.stream(request):
        event_id += 1
        # 각 이벤트에 ID를 부여하여 재개 가능하게
        chunk = format_chunk(token, request_id)
        yield f"id: {event_id}\ndata: {json.dumps(chunk)}\n\n"
        
        # 스트리밍 상태를 캐시에 저장
        await cache.set(
            f"stream:{request_id}:{event_id}",
            chunk,
            ttl=300,  # 5분간 보관
        )

import { useState, useCallback, useRef } from "react";
 
interface UseStreamingOptions {
  onToken?: (token: string) => void;
  onComplete?: (usage: TokenUsage) => void;
  onError?: (error: Error) => void;
  timeout?: number;
}
 
function useStreaming(options: UseStreamingOptions = {}) {
  const [content, setContent] = useState("");
  const [isStreaming, setIsStreaming] = useState(false);
  const [usage, setUsage] = useState<TokenUsage | null>(null);
  const [error, setError] = useState<Error | null>(null);
  const controllerRef = useRef<AbortController | null>(null);
 
  const startStream = useCallback(
    async (request: CompletionRequest) => {
      setContent("");
      setIsStreaming(true);
      setUsage(null);
      setError(null);
 
      controllerRef.current = new AbortController();
 
      try {
        const response = await fetch("/api/v1/chat/completions", {
          method: "POST",
          headers: { "Content-Type": "application/json" },
          body: JSON.stringify({ ...request, stream: true }),
          signal: controllerRef.current.signal,
        });
 
        const reader = response.body!.getReader();
        const decoder = new TextDecoder();
        let buffer = "";
 
        while (true) {
          const { done, value } = await reader.read();
          if (done) break;
 
          buffer += decoder.decode(value, { stream: true });
          const lines = buffer.split("\n");
          buffer = lines.pop() ?? "";
 
          for (const line of lines) {
            if (!line.startsWith("data: ")) continue;
            const data = line.slice(6);
            if (data === "[DONE]") break;
 
            const chunk = JSON.parse(data);
            const token = chunk.choices[0]?.delta?.content;
 
            if (token) {
              setContent((prev) => prev + token);
              options.onToken?.(token);
            }
 
            if (chunk.usage) {
              setUsage(chunk.usage);
              options.onComplete?.(chunk.usage);
            }
          }
        }
      } catch (err) {
        if (controllerRef.current?.signal.aborted) return;
        const error = err as Error;
        setError(error);
        options.onError?.(error);
      } finally {
        setIsStreaming(false);
        controllerRef.current = null;
      }
    },
    [options]
  );
 
  const cancel = useCallback(() => {
    controllerRef.current?.abort();
    setIsStreaming(false);
  }, []);
 
  return { content, isStreaming, usage, error, startStream, cancel };
}

function ChatComponent() {
  const { content, isStreaming, usage, startStream, cancel } =
    useStreaming({
      onComplete: (usage) => {
        console.log(`토큰 사용: ${usage.total_tokens}`);
      },
    });
 
  const handleSubmit = async (message: string) => {
    await startStream({
      model: "claude-4",
      messages: [{ role: "user", content: message }],
    });
  };
 
  return (
    <div>
      <div className="prose">{content}</div>
      {isStreaming && (
        <button onClick={cancel}>응답 중지</button>
      )}
      {usage && (
        <span className="text-sm text-muted">
          {usage.total_tokens} 토큰
        </span>
      )}
    </div>
  );
}