Skip to content

Commit 4134774

Browse files
narengoginarengogi
committed
base stream chunk transformer for /v1/messages for bedrock
1 parent 90f8537 commit 4134774

File tree

9 files changed

+395
-46
lines changed

9 files changed

+395
-46
lines changed

‎src/providers/anthropic-base/constants.ts‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,56 @@
1+
import {
2+
AnthropicMessageDeltaEvent,
3+
AnthropicMessageStartEvent,
4+
} from './types';
5+
6+
export const ANTHROPIC_MESSAGE_START_EVENT: AnthropicMessageStartEvent = {
7+
type: 'message_start',
8+
message: {
9+
id: '',
10+
type: 'message',
11+
role: 'assistant',
12+
model: '',
13+
content: [],
14+
stop_reason: null,
15+
stop_sequence: null,
16+
usage: {
17+
input_tokens: 0,
18+
cache_creation_input_tokens: 0,
19+
cache_read_input_tokens: 0,
20+
output_tokens: 0,
21+
},
22+
},
23+
};
24+
25+
export const ANTHROPIC_MESSAGE_DELTA_EVENT: AnthropicMessageDeltaEvent = {
26+
type: 'message_delta',
27+
delta: {
28+
stop_reason: '',
29+
stop_sequence: null,
30+
},
31+
usage: {
32+
input_tokens: 0,
33+
output_tokens: 0,
34+
cache_read_input_tokens: 0,
35+
cache_creation_input_tokens: 0,
36+
},
37+
};
38+
39+
export const ANTHROPIC_MESSAGE_STOP_EVENT = {
40+
type: 'message_stop',
41+
};
42+
43+
export const ANTHROPIC_CONTENT_BLOCK_STOP_EVENT = {
44+
type: 'content_block_stop',
45+
index: 0,
46+
};
47+
48+
export const ANTHROPIC_CONTENT_BLOCK_START_EVENT = {
49+
type: 'content_block_start',
50+
index: 1,
51+
// handle other content block types here
52+
content_block: {
53+
type: 'text',
54+
text: '',
55+
},
56+
};

‎src/providers/anthropic-base/types.ts‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
export interface AnthropicMessageStartEvent {
2+
type: 'message_start';
3+
message: {
4+
id: string;
5+
type: 'message';
6+
role: 'assistant';
7+
model: string;
8+
content: any[];
9+
stop_reason: string | null;
10+
stop_sequence: string | null;
11+
usage?: {
12+
input_tokens: number;
13+
cache_creation_input_tokens: number;
14+
cache_read_input_tokens: number;
15+
output_tokens: number;
16+
};
17+
};
18+
}
19+
20+
export interface AnthropicMessageDeltaEvent {
21+
type: 'message_delta';
22+
delta: {
23+
stop_reason: string;
24+
stop_sequence: string | null;
25+
};
26+
usage: {
27+
input_tokens?: number;
28+
output_tokens: number;
29+
cache_read_input_tokens?: number;
30+
cache_creation_input_tokens?: number;
31+
};
32+
}

‎src/providers/bedrock/api.ts‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -240,7 +240,8 @@ const BedrockAPIConfig: BedrockAPIConfigInterface = {
240240
case 'messages': {
241241
return endpoint;
242242
}
243-
case 'stream-chatComplete': {
243+
case 'stream-chatComplete':
244+
case 'stream-messages': {
244245
return streamEndpoint;
245246
}
246247
case 'complete': {

‎src/providers/bedrock/chatComplete.ts‎

Lines changed: 6 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,12 @@ import {
2727
BedrockCohereStreamChunk,
2828
} from './complete';
2929
import { BedrockErrorResponse } from './embed';
30-
import { BedrockChatCompletionResponse, BedrockContentItem } from './types';
30+
import {
31+
BedrockChatCompleteStreamChunk,
32+
BedrockChatCompletionResponse,
33+
BedrockContentItem,
34+
BedrockStreamState,
35+
} from './types';
3136
import {
3237
transformAdditionalModelRequestFields,
3338
transformAI21AdditionalModelRequestFields,
@@ -527,48 +532,6 @@ export const BedrockChatCompleteResponseTransform: (
527532
return generateInvalidProviderResponseError(response, BEDROCK);
528533
};
529534

530-
export interface BedrockChatCompleteStreamChunk {
531-
contentBlockIndex?: number;
532-
delta?: {
533-
text: string;
534-
toolUse: {
535-
toolUseId: string;
536-
name: string;
537-
input: object;
538-
};
539-
reasoningContent?: {
540-
text?: string;
541-
signature?: string;
542-
redactedContent?: string;
543-
};
544-
};
545-
start?: {
546-
toolUse: {
547-
toolUseId: string;
548-
name: string;
549-
input?: object;
550-
};
551-
};
552-
stopReason?: string;
553-
metrics?: {
554-
latencyMs: number;
555-
};
556-
usage?: {
557-
inputTokens: number;
558-
outputTokens: number;
559-
totalTokens: number;
560-
cacheReadInputTokenCount?: number;
561-
cacheReadInputTokens?: number;
562-
cacheWriteInputTokenCount?: number;
563-
cacheWriteInputTokens?: number;
564-
};
565-
}
566-
567-
interface BedrockStreamState {
568-
stopReason?: string;
569-
currentToolCallIndex?: number;
570-
}
571-
572535
// refer: https://docs.aws.amazon.com/bedrock/latest/APIReference/API_runtime_ConverseStream.html
573536
export const BedrockChatCompleteStreamChunkTransform: (
574537
response: string,

‎src/providers/bedrock/index.ts‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,7 @@ import { BedrockListFilesResponseTransform } from './listfiles';
7777
import { BedrockDeleteFileResponseTransform } from './deleteFile';
7878
import {
7979
BedrockConverseMessagesConfig,
80+
BedrockConverseMessagesStreamChunkTransform,
8081
BedrockMessagesResponseTransform,
8182
} from './messages';
8283

@@ -109,7 +110,6 @@ const BedrockConfig: ProviderConfigs = {
109110
responseTransforms: {
110111
'stream-complete': BedrockAnthropicCompleteStreamChunkTransform,
111112
complete: BedrockAnthropicCompleteResponseTransform,
112-
messages: BedrockMessagesResponseTransform,
113113
},
114114
};
115115
break;
@@ -206,6 +206,8 @@ const BedrockConfig: ProviderConfigs = {
206206
config.responseTransforms = {
207207
...(config.responseTransforms ?? {}),
208208
'stream-chatComplete': BedrockChatCompleteStreamChunkTransform,
209+
'stream-messages': BedrockConverseMessagesStreamChunkTransform,
210+
messages: BedrockMessagesResponseTransform,
209211
};
210212
}
211213
if (!config.responseTransforms?.chatComplete) {

‎src/providers/bedrock/messages.ts‎

Lines changed: 131 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,14 +13,24 @@ import {
1313
MessagesResponse,
1414
STOP_REASON,
1515
} from '../../types/messagesResponse';
16+
import { RawContentBlockDeltaEvent } from '../../types/MessagesStreamResponse';
17+
import {
18+
ANTHROPIC_CONTENT_BLOCK_START_EVENT,
19+
ANTHROPIC_CONTENT_BLOCK_STOP_EVENT,
20+
ANTHROPIC_MESSAGE_DELTA_EVENT,
21+
ANTHROPIC_MESSAGE_START_EVENT,
22+
ANTHROPIC_MESSAGE_STOP_EVENT,
23+
} from '../anthropic-base/constants';
1624
import { ErrorResponse, ProviderConfig } from '../types';
1725
import { generateInvalidProviderResponseError } from '../utils';
1826
import { BedrockErrorResponseTransform } from './chatComplete';
1927
import { BedrockErrorResponse } from './embed';
2028
import {
29+
BedrockChatCompleteStreamChunk,
2130
BedrockChatCompletionResponse,
2231
BedrockContentItem,
2332
BedrockMessagesParams,
33+
BedrockStreamState,
2434
} from './types';
2535
import {
2636
transformInferenceConfig,
@@ -397,3 +407,124 @@ export const BedrockMessagesResponseTransform = (
397407

398408
return generateInvalidProviderResponseError(response, BEDROCK);
399409
};
410+
411+
const transformContentBlock = (
412+
contentBlock: BedrockChatCompleteStreamChunk
413+
): RawContentBlockDeltaEvent | undefined => {
414+
if (!contentBlock.delta || contentBlock.contentBlockIndex === undefined) {
415+
return undefined;
416+
}
417+
if (contentBlock.delta.text) {
418+
return {
419+
type: 'content_block_delta',
420+
index: contentBlock.contentBlockIndex,
421+
delta: {
422+
type: 'text_delta',
423+
text: contentBlock.delta.text,
424+
},
425+
};
426+
} else if (contentBlock.delta.reasoningContent?.text) {
427+
return {
428+
type: 'content_block_delta',
429+
index: contentBlock.contentBlockIndex,
430+
delta: {
431+
type: 'thinking_delta',
432+
thinking: contentBlock.delta.reasoningContent.text,
433+
},
434+
};
435+
} else if (contentBlock.delta.reasoningContent?.signature) {
436+
return {
437+
type: 'content_block_delta',
438+
index: contentBlock.contentBlockIndex,
439+
delta: {
440+
type: 'signature_delta',
441+
signature: contentBlock.delta.reasoningContent.signature,
442+
},
443+
};
444+
} else if (contentBlock.delta.toolUse) {
445+
return {
446+
type: 'content_block_delta',
447+
index: contentBlock.contentBlockIndex,
448+
delta: {
449+
type: 'input_json_delta',
450+
partial_json: contentBlock.delta.toolUse.input,
451+
},
452+
};
453+
}
454+
return undefined;
455+
};
456+
457+
export const BedrockConverseMessagesStreamChunkTransform = (
458+
responseChunk: string,
459+
fallbackId: string,
460+
streamState: BedrockStreamState,
461+
strictOpenAiCompliance: boolean,
462+
gatewayRequest: Params
463+
) => {
464+
const parsedChunk: BedrockChatCompleteStreamChunk = JSON.parse(responseChunk);
465+
if (streamState.currentContentBlockIndex === undefined) {
466+
streamState.currentContentBlockIndex = -1;
467+
}
468+
if (parsedChunk.stopReason) {
469+
streamState.stopReason = parsedChunk.stopReason;
470+
}
471+
// message start event
472+
if (parsedChunk.role) {
473+
return getMessageStartEvent(fallbackId, gatewayRequest);
474+
}
475+
// content block start and stop events
476+
if (
477+
parsedChunk.contentBlockIndex !== undefined &&
478+
parsedChunk.contentBlockIndex !== streamState.currentContentBlockIndex
479+
) {
480+
let returnChunk = '';
481+
if (streamState.currentContentBlockIndex !== -1) {
482+
const previousBlockStopEvent = { ...ANTHROPIC_CONTENT_BLOCK_STOP_EVENT };
483+
previousBlockStopEvent.index = parsedChunk.contentBlockIndex - 1;
484+
returnChunk += `event: content_block_stop\ndata: ${JSON.stringify(previousBlockStopEvent)}\n\n`;
485+
}
486+
streamState.currentContentBlockIndex = parsedChunk.contentBlockIndex;
487+
const contentBlockStartEvent = { ...ANTHROPIC_CONTENT_BLOCK_START_EVENT };
488+
contentBlockStartEvent.index = parsedChunk.contentBlockIndex;
489+
returnChunk += `event: content_block_start\ndata: ${JSON.stringify(contentBlockStartEvent)}\n\n`;
490+
const contentBlockDeltaEvent = transformContentBlock(parsedChunk);
491+
if (contentBlockDeltaEvent) {
492+
returnChunk += `event: content_block_delta\ndata: ${JSON.stringify(contentBlockDeltaEvent)}\n\n`;
493+
}
494+
return returnChunk;
495+
}
496+
// content block delta event
497+
if (parsedChunk.delta) {
498+
const contentBlockDeltaEvent = transformContentBlock(parsedChunk);
499+
if (contentBlockDeltaEvent) {
500+
return `event: content_block_delta\ndata: ${JSON.stringify(contentBlockDeltaEvent)}\n\n`;
501+
}
502+
}
503+
// message delta and message stop events
504+
if (parsedChunk.usage) {
505+
const messageDeltaEvent = { ...ANTHROPIC_MESSAGE_DELTA_EVENT };
506+
messageDeltaEvent.usage.input_tokens = parsedChunk.usage.inputTokens;
507+
messageDeltaEvent.usage.output_tokens = parsedChunk.usage.outputTokens;
508+
messageDeltaEvent.usage.cache_read_input_tokens =
509+
parsedChunk.usage.cacheReadInputTokens;
510+
messageDeltaEvent.usage.cache_creation_input_tokens =
511+
parsedChunk.usage.cacheWriteInputTokens;
512+
messageDeltaEvent.delta.stop_reason = streamState.stopReason || '';
513+
const contentBlockStopEvent = { ...ANTHROPIC_CONTENT_BLOCK_STOP_EVENT };
514+
contentBlockStopEvent.index = streamState.currentContentBlockIndex;
515+
let returnChunk = `event: content_block_stop\ndata: ${JSON.stringify(contentBlockStopEvent)}\n\n`;
516+
returnChunk += `event: message_delta\ndata: ${JSON.stringify(messageDeltaEvent)}\n\n`;
517+
returnChunk += `event: message_stop\ndata: ${JSON.stringify(ANTHROPIC_MESSAGE_STOP_EVENT)}\n\n`;
518+
return returnChunk;
519+
}
520+
// console.log(JSON.stringify(parsedChunk, null, 2));
521+
};
522+
523+
function getMessageStartEvent(fallbackId: string, gatewayRequest: Params<any>) {
524+
const messageStartEvent = { ...ANTHROPIC_MESSAGE_START_EVENT };
525+
messageStartEvent.message.id = fallbackId;
526+
messageStartEvent.message.model = gatewayRequest.model as string;
527+
// bedrock does not send usage in the beginning of the stream
528+
delete messageStartEvent.message.usage;
529+
return `event: message_start\ndata: ${JSON.stringify(messageStartEvent)}\n\n`;
530+
}

‎src/providers/bedrock/types.ts‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,51 @@ export type BedrockContentItem = {
154154
type: string;
155155
};
156156
};
157+
export interface BedrockStreamState {
158+
stopReason?: string;
159+
currentToolCallIndex?: number;
160+
currentContentBlockIndex?: number;
161+
}
162+
163+
export interface BedrockContentBlockDelta {
164+
text: string;
165+
toolUse: {
166+
toolUseId: string;
167+
name: string;
168+
input: string;
169+
};
170+
reasoningContent?: {
171+
text?: string;
172+
signature?: string;
173+
redactedContent?: string;
174+
};
175+
}
176+
177+
export interface BedrockChatCompleteStreamChunk {
178+
role?: string;
179+
contentBlockIndex?: number;
180+
delta?: BedrockContentBlockDelta;
181+
start?: {
182+
toolUse: {
183+
toolUseId: string;
184+
name: string;
185+
input?: object;
186+
};
187+
};
188+
stopReason?: string;
189+
metrics?: {
190+
latencyMs: number;
191+
};
192+
usage?: {
193+
inputTokens: number;
194+
outputTokens: number;
195+
totalTokens: number;
196+
cacheReadInputTokenCount?: number;
197+
cacheReadInputTokens?: number;
198+
cacheWriteInputTokenCount?: number;
199+
cacheWriteInputTokens?: number;
200+
};
201+
}
157202

158203
// export interface BedrockConverseRequestBody {
159204
// additionalModelRequestFields?: Record<string, any>;

‎src/providers/bedrock/utils/messagesUtils.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ export const transformToolsConfig = (params: BedrockMessagesParams) => {
6868
}
6969
if (params.tools) {
7070
for (const tool of params.tools) {
71-
if (tool.type === 'custom' || tool.type === null) {
71+
if (tool.type === 'custom' || !tool.type) {
7272
tools.push({
7373
toolSpec: {
7474
name: tool.name,

0 commit comments

Comments
 (0)