OpenAI API 스트리밍 응답에서 토큰 사용량 추적하기

문제 상황

사내 AI 챗봇 서비스에서 OpenAI API를 스트리밍 모드로 사용 중이었다. 응답 속도는 만족스러웠지만, 토큰 사용량을 실시간으로 파악할 수 없어 비용 추적이 어려웠다.

일반 API 호출에서는 usage 필드로 토큰 사용량을 확인할 수 있지만, 스트리밍 응답에서는 이 정보가 마지막 청크에만 포함된다는 점을 간과했다.

해결 방법

스트리밍 응답의 마지막 청크에 usage 정보가 포함되어 있었다. stream_options 파라미터를 설정하면 된다.

const stream = await openai.chat.completions.create({
  model: 'gpt-4o',
  messages: [{ role: 'user', content: prompt }],
  stream: true,
  stream_options: { include_usage: true }
});

let totalTokens = 0;

for await (const chunk of stream) {
  if (chunk.choices[0]?.delta?.content) {
    process.stdout.write(chunk.choices[0].delta.content);
  }
  
  if (chunk.usage) {
    totalTokens = chunk.usage.total_tokens;
    console.log(`\n사용 토큰: ${totalTokens}`);
  }
}

추가 고려사항

토큰 사용량을 DB에 기록하여 사용자별, 일별 사용 패턴을 분석할 수 있게 했다. 특정 임계값을 넘으면 알림을 보내는 로직도 추가했다.

await db.tokenUsage.create({
  data: {
    userId,
    model: 'gpt-4o',
    promptTokens: chunk.usage.prompt_tokens,
    completionTokens: chunk.usage.completion_tokens,
    totalTokens: chunk.usage.total_tokens,
    timestamp: new Date()
  }
});

간단한 설정으로 비용 관리가 훨씬 수월해졌다.