From 77c64d386e3723501abd6606dc92c819f6cac43b Mon Sep 17 00:00:00 2001 From: Gaspar Garcia Date: Fri, 17 May 2024 16:40:35 -0700 Subject: [PATCH 1/4] Add doRawStream --- .../openai/src/openai-chat-language-model.ts | 35 +++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/packages/openai/src/openai-chat-language-model.ts b/packages/openai/src/openai-chat-language-model.ts index 521b61af12fc..6fb52d0112a4 100644 --- a/packages/openai/src/openai-chat-language-model.ts +++ b/packages/openai/src/openai-chat-language-model.ts @@ -188,6 +188,41 @@ export class OpenAIChatLanguageModel implements LanguageModelV1 { }; } + async doRawStream( + options: Parameters[0], + ): Promise>> { + const args = this.getArgs(options); + + const { responseHeaders, value: response } = await postJsonToApi({ + url: `${this.config.baseURL}/chat/completions`, + headers: this.config.headers(), + body: { + ...args, + stream: true, + + // only include stream_options when in strict compatibility mode: + stream_options: + this.config.compatibility === 'strict' + ? { include_usage: true } + : undefined, + }, + failedResponseHandler: openaiFailedResponseHandler, + successfulResponseHandler: createEventSourceResponseHandler( + openaiChatChunkSchema, + ), + abortSignal: options.abortSignal, + }); + + const { messages: rawPrompt, ...rawSettings } = args; + + return { + stream: response, + rawCall: { rawPrompt, rawSettings }, + rawResponse: { headers: responseHeaders }, + warnings: [], + }; + } + async doStream( options: Parameters[0], ): Promise>> { From 5830c6667ae2ea3494fa266eca8aadea22452c88 Mon Sep 17 00:00:00 2001 From: Gaspar Garcia Date: Fri, 17 May 2024 16:45:09 -0700 Subject: [PATCH 2/4] skip ts on src/openai-chat-language-model.ts --- packages/openai/src/openai-chat-language-model.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/openai/src/openai-chat-language-model.ts b/packages/openai/src/openai-chat-language-model.ts index 6fb52d0112a4..a2b9ee0a77cd 100644 --- a/packages/openai/src/openai-chat-language-model.ts +++ b/packages/openai/src/openai-chat-language-model.ts @@ -1,3 +1,4 @@ +// @ts-nocheck import { InvalidResponseDataError, LanguageModelV1, From 52f89c88fee86b2d3ef04b120cb50958e1b633ca Mon Sep 17 00:00:00 2001 From: Gaspar Garcia Date: Fri, 17 May 2024 17:37:20 -0700 Subject: [PATCH 3/4] add streamResponse method --- .../core/core/generate-text/stream-text.ts | 59 +++++++++++++++++++ .../openai/src/openai-chat-language-model.ts | 14 +++-- .../provider-utils/src/response-handler.ts | 15 +++++ .../language-model/v1/language-model-v1.ts | 33 +++++++++++ 4 files changed, 116 insertions(+), 5 deletions(-) diff --git a/packages/core/core/generate-text/stream-text.ts b/packages/core/core/generate-text/stream-text.ts index e52134a350fc..61ea39aa5370 100644 --- a/packages/core/core/generate-text/stream-text.ts +++ b/packages/core/core/generate-text/stream-text.ts @@ -21,6 +21,7 @@ import { runToolsTransformation } from './run-tools-transformation'; import { TokenUsage } from './token-usage'; import { ToToolCall } from './tool-call'; import { ToToolResult } from './tool-result'; +import { LanguageModelV1CallWarning } from '@ai-sdk/provider'; /** Generate a text and call tools for a given prompt using a language model. @@ -110,6 +111,64 @@ The tools that the model can call. The model needs to support calling tools. }); } +export async function streamResponse>({ + model, + tools, + system, + prompt, + messages, + maxRetries, + abortSignal, + ...settings +}: CallSettings & + Prompt & { + /** +The language model to use. + */ + model: LanguageModel; + + /** +The tools that the model can call. The model needs to support calling tools. + */ + tools?: TOOLS; + }): Promise<{ + stream: ReadableStream; + warnings: LanguageModelV1CallWarning[] | undefined; + rawResponse: + | { + headers?: Record; + } + | undefined; +}> { + const retry = retryWithExponentialBackoff({ maxRetries }); + const validatedPrompt = getValidatedPrompt({ system, prompt, messages }); + const { stream, warnings, rawResponse } = await retry(() => { + if (!model.doRawStream) { + throw new Error('The model does not support raw streaming.'); + } + return model.doRawStream({ + mode: { + type: 'regular', + tools: + tools == null + ? undefined + : Object.entries(tools).map(([name, tool]) => ({ + type: 'function', + name, + description: tool.description, + parameters: convertZodToJSONSchema(tool.parameters), + })), + }, + ...prepareCallSettings(settings), + inputFormat: validatedPrompt.type, + prompt: convertToLanguageModelPrompt(validatedPrompt), + abortSignal, + }); + }); + + return { stream, warnings, rawResponse }; +} + export type TextStreamPart> = | { type: 'text-delta'; diff --git a/packages/openai/src/openai-chat-language-model.ts b/packages/openai/src/openai-chat-language-model.ts index a2b9ee0a77cd..a1949ce911e5 100644 --- a/packages/openai/src/openai-chat-language-model.ts +++ b/packages/openai/src/openai-chat-language-model.ts @@ -1,4 +1,3 @@ -// @ts-nocheck import { InvalidResponseDataError, LanguageModelV1, @@ -9,6 +8,7 @@ import { } from '@ai-sdk/provider'; import { ParseResult, + createEventSourcePassThroughHandler, createEventSourceResponseHandler, createJsonResponseHandler, generateId, @@ -191,10 +191,14 @@ export class OpenAIChatLanguageModel implements LanguageModelV1 { async doRawStream( options: Parameters[0], - ): Promise>> { + ): Promise< + Omit>, 'stream'> & { + stream: ReadableStream; + } + > { const args = this.getArgs(options); - const { responseHeaders, value: response } = await postJsonToApi({ + const { responseHeaders, value: responseBody } = await postJsonToApi({ url: `${this.config.baseURL}/chat/completions`, headers: this.config.headers(), body: { @@ -208,7 +212,7 @@ export class OpenAIChatLanguageModel implements LanguageModelV1 { : undefined, }, failedResponseHandler: openaiFailedResponseHandler, - successfulResponseHandler: createEventSourceResponseHandler( + successfulResponseHandler: createEventSourcePassThroughHandler( openaiChatChunkSchema, ), abortSignal: options.abortSignal, @@ -217,7 +221,7 @@ export class OpenAIChatLanguageModel implements LanguageModelV1 { const { messages: rawPrompt, ...rawSettings } = args; return { - stream: response, + stream: responseBody, rawCall: { rawPrompt, rawSettings }, rawResponse: { headers: responseHeaders }, warnings: [], diff --git a/packages/provider-utils/src/response-handler.ts b/packages/provider-utils/src/response-handler.ts index 63b5b5a8dad1..f61001c80f31 100644 --- a/packages/provider-utils/src/response-handler.ts +++ b/packages/provider-utils/src/response-handler.ts @@ -118,6 +118,21 @@ export const createEventSourceResponseHandler = }; }; +export const createEventSourcePassThroughHandler = + (chunkSchema: ZodSchema): ResponseHandler> => + async ({ response }: { response: Response }) => { + const responseHeaders = extractResponseHeaders(response); + + if (response.body == null) { + throw new EmptyResponseBodyError({}); + } + + return { + responseHeaders, + value: response.body, + }; + }; + export const createJsonResponseHandler = (responseSchema: ZodSchema): ResponseHandler => async ({ response, url, requestBodyValues }) => { diff --git a/packages/provider/src/language-model/v1/language-model-v1.ts b/packages/provider/src/language-model/v1/language-model-v1.ts index 3d3c09a49285..4cbc73cfa725 100644 --- a/packages/provider/src/language-model/v1/language-model-v1.ts +++ b/packages/provider/src/language-model/v1/language-model-v1.ts @@ -145,6 +145,39 @@ Response headers. warnings?: LanguageModelV1CallWarning[]; }>; + + doRawStream?: (options: LanguageModelV1CallOptions) => PromiseLike<{ + stream: ReadableStream; + + /** +Raw prompt and setting information for observability provider integration. + */ + rawCall: { + /** +Raw prompt after expansion and conversion to the format that the +provider uses to send the information to their API. + */ + rawPrompt: unknown; + + /** +Raw settings that are used for the API call. Includes provider-specific +settings. + */ + rawSettings: Record; + }; + + /** +Optional raw response data. + */ + rawResponse?: { + /** +Response headers. + */ + headers?: Record; + }; + + warnings?: LanguageModelV1CallWarning[]; + }>; }; export type LanguageModelV1StreamPart = From 85f9a6350aab323e939fdb15491d20cea6fc86a0 Mon Sep 17 00:00:00 2001 From: Gaspar Garcia Date: Fri, 17 May 2024 17:46:59 -0700 Subject: [PATCH 4/4] Add changeset --- .changeset/clean-brooms-fold.md | 8 ++++++++ 1 file changed, 8 insertions(+) create mode 100644 .changeset/clean-brooms-fold.md diff --git a/.changeset/clean-brooms-fold.md b/.changeset/clean-brooms-fold.md new file mode 100644 index 000000000000..50d6a0b6bacf --- /dev/null +++ b/.changeset/clean-brooms-fold.md @@ -0,0 +1,8 @@ +--- +'@ai-sdk/provider-utils': patch +'@ai-sdk/provider': patch +'@ai-sdk/openai': patch +'ai': patch +--- + +Prototype Raw Response