mirror of
https://github.com/stackblitz/bolt.new
synced 2024-11-30 23:56:07 +00:00
39 lines
1.1 KiB
TypeScript
39 lines
1.1 KiB
TypeScript
|
import { streamText as _streamText, convertToCoreMessages } from 'ai';
|
||
|
import { getAPIKey } from '../llm/api-key';
|
||
|
import { getAnthropicModel } from '../llm/model';
|
||
|
import { MAX_TOKENS } from './constants';
|
||
|
import { getSystemPrompt } from './prompts';
|
||
|
|
||
|
interface ToolResult<Name extends string, Args, Result> {
|
||
|
toolCallId: string;
|
||
|
toolName: Name;
|
||
|
args: Args;
|
||
|
result: Result;
|
||
|
}
|
||
|
|
||
|
interface Message {
|
||
|
role: 'user' | 'assistant';
|
||
|
content: string;
|
||
|
toolInvocations?: ToolResult<string, unknown, unknown>[];
|
||
|
}
|
||
|
|
||
|
export type Messages = Message[];
|
||
|
|
||
|
export type StreamingOptions = Omit<Parameters<typeof _streamText>[0], 'model'>;
|
||
|
|
||
|
export function streamText(messages: Messages, env: Env, options?: StreamingOptions) {
|
||
|
return _streamText({
|
||
|
model: getAnthropicModel(getAPIKey(env)),
|
||
|
system: getSystemPrompt(),
|
||
|
maxTokens: MAX_TOKENS,
|
||
|
headers: {
|
||
|
'anthropic-beta': 'max-tokens-3-5-sonnet-2024-07-15',
|
||
|
},
|
||
|
messages: convertToCoreMessages(messages),
|
||
|
onFinish: ({ finishReason, usage, warnings }) => {
|
||
|
console.log({ finishReason, usage, warnings });
|
||
|
},
|
||
|
...options,
|
||
|
});
|
||
|
}
|