@lobehub/chat
Version:
Lobe Chat - an open-source, high-performance chatbot framework that supports speech synthesis, multimodal, and extensible Function Call plugin system. Supports one-click free deployment of your private ChatGPT/LLM web application.
262 lines (223 loc) • 7.66 kB
text/typescript
import Anthropic from '@anthropic-ai/sdk';
import type { Stream } from '@anthropic-ai/sdk/streaming';
import { CitationItem, ModelTokensUsage } from '@/types/message';
import { ChatStreamCallbacks } from '../../types';
import {
StreamContext,
StreamProtocolChunk,
StreamProtocolToolCallChunk,
StreamToolCallChunkData,
convertIterableToStream,
createCallbacksTransformer,
createSSEProtocolTransformer,
createTokenSpeedCalculator,
} from './protocol';
export const transformAnthropicStream = (
chunk: Anthropic.MessageStreamEvent,
context: StreamContext,
): StreamProtocolChunk | StreamProtocolChunk[] => {
// maybe need another structure to add support for multiple choices
switch (chunk.type) {
case 'message_start': {
context.id = chunk.message.id;
context.returnedCitationArray = [];
let totalInputTokens = chunk.message.usage?.input_tokens;
if (
chunk.message.usage?.cache_creation_input_tokens ||
chunk.message.usage?.cache_read_input_tokens
) {
totalInputTokens =
chunk.message.usage?.input_tokens +
(chunk.message.usage.cache_creation_input_tokens || 0) +
(chunk.message.usage.cache_read_input_tokens || 0);
}
context.usage = {
inputCacheMissTokens: chunk.message.usage?.input_tokens,
inputCachedTokens: chunk.message.usage?.cache_read_input_tokens || undefined,
inputWriteCacheTokens: chunk.message.usage?.cache_creation_input_tokens || undefined,
totalInputTokens,
totalOutputTokens: chunk.message.usage?.output_tokens,
};
return { data: chunk.message, id: chunk.message.id, type: 'data' };
}
case 'content_block_start': {
switch (chunk.content_block.type) {
case 'redacted_thinking': {
return {
data: chunk.content_block.data,
id: context.id,
type: 'flagged_reasoning_signature',
};
}
case 'text': {
return { data: chunk.content_block.text, id: context.id, type: 'data' };
}
case 'server_tool_use':
case 'tool_use': {
const toolChunk = chunk.content_block;
// if toolIndex is not defined, set it to 0
if (typeof context.toolIndex === 'undefined') {
context.toolIndex = 0;
}
// if toolIndex is defined, increment it
else {
context.toolIndex += 1;
}
const toolCall: StreamToolCallChunkData = {
function: {
arguments: '',
name: toolChunk.name,
},
id: toolChunk.id,
index: context.toolIndex,
type: 'function',
};
context.tool = { id: toolChunk.id, index: context.toolIndex, name: toolChunk.name };
return { data: [toolCall], id: context.id, type: 'tool_calls' };
}
/*
case 'web_search_tool_result': {
const citations = chunk.content_block.content;
return [
{
data: {
citations: (citations as any[]).map(
(item) =>
({
title: item.title,
url: item.url,
}) as CitationItem,
),
},
id: context.id,
type: 'grounding',
},
];
}
*/
case 'thinking': {
const thinkingChunk = chunk.content_block;
// if there is signature in the thinking block, return both thinking and signature
if (!!thinkingChunk.signature) {
return [
{ data: thinkingChunk.thinking, id: context.id, type: 'reasoning' },
{ data: thinkingChunk.signature, id: context.id, type: 'reasoning_signature' },
];
}
if (typeof thinkingChunk.thinking === 'string')
return { data: thinkingChunk.thinking, id: context.id, type: 'reasoning' };
return { data: thinkingChunk, id: context.id, type: 'data' };
}
default: {
break;
}
}
return { data: chunk, id: context.id, type: 'data' };
}
case 'content_block_delta': {
switch (chunk.delta.type) {
case 'text_delta': {
return { data: chunk.delta.text, id: context.id, type: 'text' };
}
case 'input_json_delta': {
const delta = chunk.delta.partial_json;
const toolCall: StreamToolCallChunkData = {
function: { arguments: delta },
index: context.toolIndex || 0,
type: 'function',
};
return {
data: [toolCall],
id: context.id,
type: 'tool_calls',
} as StreamProtocolToolCallChunk;
}
case 'signature_delta': {
return {
data: chunk.delta.signature,
id: context.id,
type: 'reasoning_signature',
};
}
case 'thinking_delta': {
return {
data: chunk.delta.thinking,
id: context.id,
type: 'reasoning',
};
}
case 'citations_delta': {
const citations = (chunk as any).delta.citation;
if (context.returnedCitationArray) {
context.returnedCitationArray.push({
title: citations.title,
url: citations.url,
} as CitationItem);
}
return { data: null, id: context.id, type: 'text' };
}
default: {
break;
}
}
return { data: chunk, id: context.id, type: 'data' };
}
case 'message_delta': {
const totalOutputTokens =
chunk.usage?.output_tokens + (context.usage?.totalOutputTokens || 0);
const totalInputTokens = context.usage?.totalInputTokens || 0;
const totalTokens = totalInputTokens + totalOutputTokens;
if (totalTokens > 0) {
return [
{ data: chunk.delta.stop_reason, id: context.id, type: 'stop' },
{
data: {
...context.usage,
totalInputTokens,
totalOutputTokens,
totalTokens,
} as ModelTokensUsage,
id: context.id,
type: 'usage',
},
];
}
return { data: chunk.delta.stop_reason, id: context.id, type: 'stop' };
}
case 'message_stop': {
return [
...(context.returnedCitationArray?.length
? [
{
data: { citations: context.returnedCitationArray },
id: context.id,
type: 'grounding',
},
]
: []),
{ data: 'message_stop', id: context.id, type: 'stop' },
] as any;
}
default: {
return { data: chunk, id: context.id, type: 'data' };
}
}
};
export interface AnthropicStreamOptions {
callbacks?: ChatStreamCallbacks;
inputStartAt?: number;
}
export const AnthropicStream = (
stream: Stream<Anthropic.MessageStreamEvent> | ReadableStream,
{ callbacks, inputStartAt }: AnthropicStreamOptions = {},
) => {
const streamStack: StreamContext = { id: '' };
const readableStream =
stream instanceof ReadableStream ? stream : convertIterableToStream(stream);
return readableStream
.pipeThrough(
createTokenSpeedCalculator(transformAnthropicStream, { inputStartAt, streamStack }),
)
.pipeThrough(createSSEProtocolTransformer((c) => c, streamStack))
.pipeThrough(createCallbacksTransformer(callbacks));
};