Skip to content

Commit eccbc7c

Browse files
narengoginarengogi
committed
handle stream generation and do not create challow copies of json templates
1 parent ee21746 commit eccbc7c

File tree

4 files changed

+216
-20
lines changed

4 files changed

+216
-20
lines changed

‎src/handlers/responseHandlers.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ import {
1717
import { HookSpan } from '../middlewares/hooks';
1818
import { env } from 'hono/adapter';
1919
import { OpenAIModelResponseJSONToStreamGenerator } from '../providers/open-ai-base/createModelResponse';
20+
import { anthropicMessagesJsonToStreamGenerator } from '../providers/anthropic-base/utils/streamGenerator';
2021

2122
/**
2223
* Handles various types of responses based on the specified parameters
@@ -81,6 +82,9 @@ export async function responseHandler(
8182
responseTransformerFunction =
8283
OpenAIChatCompleteJSONToStreamResponseTransform;
8384
break;
85+
case 'messages':
86+
responseTransformerFunction = anthropicMessagesJsonToStreamGenerator;
87+
break;
8488
case 'createModelResponse':
8589
responseTransformerFunction = OpenAIModelResponseJSONToStreamGenerator;
8690
break;

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

Lines changed: 8 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,4 @@
1-
import {
2-
AnthropicMessageDeltaEvent,
3-
AnthropicMessageStartEvent,
4-
} from './types';
5-
6-
export const ANTHROPIC_MESSAGE_START_EVENT: AnthropicMessageStartEvent = {
1+
export const ANTHROPIC_MESSAGE_START_EVENT = JSON.stringify({
72
type: 'message_start',
83
message: {
94
id: '',
@@ -20,9 +15,9 @@ export const ANTHROPIC_MESSAGE_START_EVENT: AnthropicMessageStartEvent = {
2015
output_tokens: 0,
2116
},
2217
},
23-
};
18+
});
2419

25-
export const ANTHROPIC_MESSAGE_DELTA_EVENT: AnthropicMessageDeltaEvent = {
20+
export const ANTHROPIC_MESSAGE_DELTA_EVENT = JSON.stringify({
2621
type: 'message_delta',
2722
delta: {
2823
stop_reason: '',
@@ -34,23 +29,23 @@ export const ANTHROPIC_MESSAGE_DELTA_EVENT: AnthropicMessageDeltaEvent = {
3429
cache_read_input_tokens: 0,
3530
cache_creation_input_tokens: 0,
3631
},
37-
};
32+
});
3833

3934
export const ANTHROPIC_MESSAGE_STOP_EVENT = {
4035
type: 'message_stop',
4136
};
4237

43-
export const ANTHROPIC_CONTENT_BLOCK_STOP_EVENT = {
38+
export const ANTHROPIC_CONTENT_BLOCK_STOP_EVENT = JSON.stringify({
4439
type: 'content_block_stop',
4540
index: 0,
46-
};
41+
});
4742

48-
export const ANTHROPIC_CONTENT_BLOCK_START_EVENT = {
43+
export const ANTHROPIC_CONTENT_BLOCK_START_EVENT = JSON.stringify({
4944
type: 'content_block_start',
5045
index: 1,
5146
// handle other content block types here
5247
content_block: {
5348
type: 'text',
5449
text: '',
5550
},
56-
};
51+
});

‎src/providers/anthropic-base/utils/streamGenerator.ts‎

Lines changed: 178 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,178 @@
1+
import {
2+
MessagesResponse,
3+
TextBlock,
4+
TextCitation,
5+
ThinkingBlock,
6+
ToolUseBlock,
7+
} from '../../../types/messagesResponse';
8+
9+
const getMessageStartEvent = (response: MessagesResponse): string => {
10+
const message = { ...response, content: [], type: 'message_start' };
11+
return `event: message_start\ndata: ${JSON.stringify({
12+
type: 'message_start',
13+
message,
14+
})}\n\n`;
15+
};
16+
17+
const getMessageDeltaEvent = (response: MessagesResponse): string => {
18+
const messageDeltaEvent = {
19+
type: 'message_delta',
20+
delta: {
21+
stop_reason: response.stop_reason,
22+
stop_sequence: response.stop_sequence,
23+
},
24+
usage: response.usage,
25+
};
26+
return `event: message_delta\ndata: ${JSON.stringify(messageDeltaEvent)}\n\n`;
27+
};
28+
29+
const MESSAGE_STOP_EVENT = `event: message_stop\ndata: {type: 'message_stop'}\n\n`;
30+
31+
const textContentBlockStartEvent = (index: number): string => {
32+
return `event: content_block_start\ndata: ${JSON.stringify({
33+
type: 'content_block_start',
34+
index,
35+
content_block: {
36+
type: 'text',
37+
text: '',
38+
},
39+
})}\n\n`;
40+
};
41+
42+
const textContentBlockDeltaEvent = (
43+
index: number,
44+
textBlock: TextBlock
45+
): string => {
46+
return `event: content_block_delta\ndata: ${JSON.stringify({
47+
type: 'content_block_delta',
48+
index,
49+
delta: {
50+
type: 'text_delta',
51+
text: textBlock.text,
52+
},
53+
})}\n\n`;
54+
};
55+
56+
const toolUseContentBlockStartEvent = (
57+
index: number,
58+
toolUseBlock: ToolUseBlock
59+
): string => {
60+
return `event: content_block_start\ndata: ${JSON.stringify({
61+
type: 'content_block_start',
62+
index,
63+
content_block: {
64+
type: 'tool_use',
65+
tool_use: { ...toolUseBlock, input: {} },
66+
},
67+
})}\n\n`;
68+
};
69+
70+
const toolUseContentBlockDeltaEvent = (
71+
index: number,
72+
toolUseBlock: ToolUseBlock
73+
): string => {
74+
return `event: content_block_delta\ndata: ${JSON.stringify({
75+
type: 'content_block_delta',
76+
index,
77+
delta: {
78+
type: 'input_json_delta',
79+
partial_json: JSON.stringify(toolUseBlock.input),
80+
},
81+
})}\n\n`;
82+
};
83+
84+
const thinkingContentBlockStartEvent = (index: number): string => {
85+
return `event: content_block_start\ndata: ${JSON.stringify({
86+
type: 'content_block_start',
87+
index,
88+
content_block: {
89+
type: 'thinking',
90+
thinking: '',
91+
signature: '',
92+
},
93+
})}\n\n`;
94+
};
95+
96+
const thinkingContentBlockDeltaEvent = (
97+
index: number,
98+
thinkingBlock: ThinkingBlock
99+
): string => {
100+
return `event: content_block_delta\ndata: ${JSON.stringify({
101+
type: 'content_block_delta',
102+
index,
103+
delta: {
104+
type: 'thinking_delta',
105+
thinking: thinkingBlock.thinking,
106+
},
107+
})}\n\n`;
108+
};
109+
110+
const signatureContentBlockDeltaEvent = (
111+
index: number,
112+
thinkingBlock: ThinkingBlock
113+
): string => {
114+
return `event: content_block_delta\ndata: ${JSON.stringify({
115+
type: 'content_block_delta',
116+
index,
117+
delta: {
118+
type: 'signature_delta',
119+
signature: thinkingBlock.signature,
120+
},
121+
})}\n\n`;
122+
};
123+
124+
const citationContentBlockDeltaEvent = (
125+
index: number,
126+
citation: TextCitation
127+
): string => {
128+
return `event: content_block_delta\ndata: ${JSON.stringify({
129+
type: 'content_block_delta',
130+
index,
131+
delta: {
132+
type: 'citations_delta',
133+
citation,
134+
},
135+
})}\n\n`;
136+
};
137+
138+
const contentBlockStopEvent = (index: number): string => {
139+
return `event: content_block_stop\ndata: ${JSON.stringify({
140+
type: 'content_block_stop',
141+
index,
142+
})}\n\n`;
143+
};
144+
145+
export function* anthropicMessagesJsonToStreamGenerator(
146+
response: MessagesResponse
147+
): Generator<string, void, unknown> {
148+
yield getMessageStartEvent(response);
149+
150+
for (const [index, contentBlock] of response.content.entries()) {
151+
switch (contentBlock.type) {
152+
case 'text':
153+
yield textContentBlockStartEvent(index);
154+
yield textContentBlockDeltaEvent(index, contentBlock);
155+
if (contentBlock.citations) {
156+
for (const citation of contentBlock.citations) {
157+
yield citationContentBlockDeltaEvent(index, citation);
158+
}
159+
}
160+
break;
161+
case 'tool_use':
162+
yield toolUseContentBlockStartEvent(index, contentBlock);
163+
yield toolUseContentBlockDeltaEvent(index, contentBlock);
164+
break;
165+
case 'thinking':
166+
yield thinkingContentBlockStartEvent(index);
167+
yield thinkingContentBlockDeltaEvent(index, contentBlock);
168+
yield signatureContentBlockDeltaEvent(index, contentBlock);
169+
break;
170+
}
171+
yield contentBlockStopEvent(index);
172+
}
173+
174+
yield getMessageDeltaEvent(response);
175+
176+
yield MESSAGE_STOP_EVENT;
177+
}
178+
``;

‎src/providers/bedrock/messages.ts‎

Lines changed: 26 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -9,14 +9,23 @@ import {
99
ToolUseBlockParam,
1010
} from '../../types/MessagesRequest';
1111
import { ContentBlock, MessagesResponse } from '../../types/messagesResponse';
12-
import { RawContentBlockDeltaEvent } from '../../types/MessagesStreamResponse';
12+
import {
13+
RawContentBlockDeltaEvent,
14+
RawContentBlockStartEvent,
15+
RawContentBlockStopEvent,
16+
} from '../../types/MessagesStreamResponse';
17+
import { Params } from '../../types/requestBody';
1318
import {
1419
ANTHROPIC_CONTENT_BLOCK_START_EVENT,
1520
ANTHROPIC_CONTENT_BLOCK_STOP_EVENT,
1621
ANTHROPIC_MESSAGE_DELTA_EVENT,
1722
ANTHROPIC_MESSAGE_START_EVENT,
1823
ANTHROPIC_MESSAGE_STOP_EVENT,
1924
} from '../anthropic-base/constants';
25+
import {
26+
AnthropicMessageDeltaEvent,
27+
AnthropicMessageStartEvent,
28+
} from '../anthropic-base/types';
2029
import { ErrorResponse, ProviderConfig } from '../types';
2130
import {
2231
generateInvalidProviderResponseError,
@@ -525,12 +534,16 @@ export const BedrockConverseMessagesStreamChunkTransform = (
525534
) {
526535
let returnChunk = '';
527536
if (streamState.currentContentBlockIndex !== -1) {
528-
const previousBlockStopEvent = { ...ANTHROPIC_CONTENT_BLOCK_STOP_EVENT };
537+
const previousBlockStopEvent: RawContentBlockStopEvent = JSON.parse(
538+
ANTHROPIC_CONTENT_BLOCK_STOP_EVENT
539+
);
529540
previousBlockStopEvent.index = parsedChunk.contentBlockIndex - 1;
530541
returnChunk += `event: content_block_stop\ndata: ${JSON.stringify(previousBlockStopEvent)}\n\n`;
531542
}
532543
streamState.currentContentBlockIndex = parsedChunk.contentBlockIndex;
533-
const contentBlockStartEvent = { ...ANTHROPIC_CONTENT_BLOCK_START_EVENT };
544+
const contentBlockStartEvent: RawContentBlockStartEvent = JSON.parse(
545+
ANTHROPIC_CONTENT_BLOCK_START_EVENT
546+
);
534547
contentBlockStartEvent.index = parsedChunk.contentBlockIndex;
535548
returnChunk += `event: content_block_start\ndata: ${JSON.stringify(contentBlockStartEvent)}\n\n`;
536549
const contentBlockDeltaEvent = transformContentBlock(parsedChunk);
@@ -548,7 +561,9 @@ export const BedrockConverseMessagesStreamChunkTransform = (
548561
}
549562
// message delta and message stop events
550563
if (parsedChunk.usage) {
551-
const messageDeltaEvent = { ...ANTHROPIC_MESSAGE_DELTA_EVENT };
564+
const messageDeltaEvent: AnthropicMessageDeltaEvent = JSON.parse(
565+
ANTHROPIC_MESSAGE_DELTA_EVENT
566+
);
552567
messageDeltaEvent.usage.input_tokens = parsedChunk.usage.inputTokens;
553568
messageDeltaEvent.usage.output_tokens = parsedChunk.usage.outputTokens;
554569
messageDeltaEvent.usage.cache_read_input_tokens =
@@ -558,7 +573,9 @@ export const BedrockConverseMessagesStreamChunkTransform = (
558573
messageDeltaEvent.delta.stop_reason = transformToAnthropicStopReason(
559574
streamState.stopReason
560575
);
561-
const contentBlockStopEvent = { ...ANTHROPIC_CONTENT_BLOCK_STOP_EVENT };
576+
const contentBlockStopEvent: RawContentBlockStopEvent = JSON.parse(
577+
ANTHROPIC_CONTENT_BLOCK_STOP_EVENT
578+
);
562579
contentBlockStopEvent.index = streamState.currentContentBlockIndex;
563580
let returnChunk = `event: content_block_stop\ndata: ${JSON.stringify(contentBlockStopEvent)}\n\n`;
564581
returnChunk += `event: message_delta\ndata: ${JSON.stringify(messageDeltaEvent)}\n\n`;
@@ -567,8 +584,10 @@ export const BedrockConverseMessagesStreamChunkTransform = (
567584
}
568585
};
569586

570-
function getMessageStartEvent(fallbackId: string, gatewayRequest: Params<any>) {
571-
const messageStartEvent = { ...ANTHROPIC_MESSAGE_START_EVENT };
587+
function getMessageStartEvent(fallbackId: string, gatewayRequest: Params) {
588+
const messageStartEvent: AnthropicMessageStartEvent = JSON.parse(
589+
ANTHROPIC_MESSAGE_START_EVENT
590+
);
572591
messageStartEvent.message.id = fallbackId;
573592
messageStartEvent.message.model = gatewayRequest.model as string;
574593
// bedrock does not send usage in the beginning of the stream

0 commit comments

Comments
 (0)