ai/core: expose raw response headers (#1417)

Author: lgrammel

.changeset/short-seas-flash.md

@@ -0,0 +1,11 @@
+---
+'@ai-sdk/provider-utils': patch
+'@ai-sdk/anthropic': patch
+'@ai-sdk/provider': patch
+'@ai-sdk/mistral': patch
+'@ai-sdk/google': patch
+'@ai-sdk/openai': patch
+'ai': patch
+---
+
+ai/core: add support for getting raw response headers.

examples/ai-core/src/stream-text/openai-response-headers.ts

@@ -0,0 +1,24 @@
+import { openai } from '@ai-sdk/openai';
+import { experimental_streamText } from 'ai';
+import dotenv from 'dotenv';
+
+dotenv.config();
+
+async function main() {
+  const result = await experimental_streamText({
+    model: openai('gpt-3.5-turbo'),
+    maxTokens: 512,
+    temperature: 0.3,
+    maxRetries: 5,
+    prompt: 'Invent a new holiday and describe its traditions.',
+  });
+
+  console.log(`Request ID: ${result.rawResponse?.headers?.['x-request-id']}`);
+  console.log();
+
+  for await (const textPart of result.textStream) {
+    process.stdout.write(textPart);
+  }
+}
+
+main().catch(console.error);

packages/anthropic/src/anthropic-messages-language-model.test.ts

@@ -11,10 +11,7 @@ const TEST_PROMPT: LanguageModelV1Prompt = [
   { role: 'user', content: [{ type: 'text', text: 'Hello' }] },
 ];
 
-const provider = createAnthropic({
-  apiKey: 'test-api-key',
-});
-
+const provider = createAnthropic({ apiKey: 'test-api-key' });
 const model = provider.chat('claude-3-haiku-20240307');
 
 describe('doGenerate', () => {
@@ -181,6 +178,28 @@ describe('doGenerate', () => {
     });
   });
 
+  it('should expose the raw response headers', async () => {
+    prepareJsonResponse({});
+
+    server.responseHeaders = {
+      'test-header': 'test-value',
+    };
+
+    const { rawResponse } = await model.doGenerate({
+      inputFormat: 'prompt',
+      mode: { type: 'regular' },
+      prompt: TEST_PROMPT,
+    });
+
+    expect(rawResponse?.headers).toStrictEqual({
+      // default headers:
+      'content-type': 'application/json',
+
+      // custom header
+      'test-header': 'test-value',
+    });
+  });
+
   it('should pass the model and the messages', async () => {
     prepareJsonResponse({});
 
@@ -279,6 +298,30 @@ describe('doStream', () => {
     ]);
   });
 
+  it('should expose the raw response headers', async () => {
+    prepareStreamResponse({ content: [] });
+
+    server.responseHeaders = {
+      'test-header': 'test-value',
+    };
+
+    const { rawResponse } = await model.doStream({
+      inputFormat: 'prompt',
+      mode: { type: 'regular' },
+      prompt: TEST_PROMPT,
+    });
+
+    expect(rawResponse?.headers).toStrictEqual({
+      // default headers:
+      'content-type': 'text/event-stream',
+      'cache-control': 'no-cache',
+      connection: 'keep-alive',
+
+      // custom header
+      'test-header': 'test-value',
+    });
+  });
+
   it('should pass the messages and the model', async () => {
     prepareStreamResponse({ content: [] });
 

packages/anthropic/src/anthropic-messages-language-model.ts

@@ -164,7 +164,7 @@ export class AnthropicMessagesLanguageModel implements LanguageModelV1 {
   ): Promise<Awaited<ReturnType<LanguageModelV1['doGenerate']>>> {
     const { args, warnings } = this.getArgs(options);
 
-    const response = await postJsonToApi({
+    const { responseHeaders, value: response } = await postJsonToApi({
       url: `${this.config.baseURL}/messages`,
       headers: this.config.headers(),
       body: args,
@@ -210,6 +210,7 @@ export class AnthropicMessagesLanguageModel implements LanguageModelV1 {
         completionTokens: response.usage.output_tokens,
       },
       rawCall: { rawPrompt, rawSettings },
+      rawResponse: { headers: responseHeaders },
       warnings,
     };
   }
@@ -219,7 +220,7 @@ export class AnthropicMessagesLanguageModel implements LanguageModelV1 {
   ): Promise<Awaited<ReturnType<LanguageModelV1['doStream']>>> {
     const { args, warnings } = this.getArgs(options);
 
-    const response = await postJsonToApi({
+    const { responseHeaders, value: response } = await postJsonToApi({
       url: `${this.config.baseURL}/messages`,
       headers: this.config.headers(),
       body: {
@@ -296,6 +297,7 @@ export class AnthropicMessagesLanguageModel implements LanguageModelV1 {
         }),
       ),
       rawCall: { rawPrompt, rawSettings },
+      rawResponse: { headers: responseHeaders },
       warnings,
     };
   }

packages/core/core/generate-object/generate-object.ts

@@ -94,6 +94,7 @@ Default and recommended: 'auto' (best mode for the model).
   let finishReason: LanguageModelV1FinishReason;
   let usage: Parameters<typeof calculateTokenUsage>[0];
   let warnings: LanguageModelV1CallWarning[] | undefined;
+  let rawResponse: { headers?: Record<string, string> } | undefined;
   let logprobs: LanguageModelV1LogProbs | undefined;
 
   switch (mode) {
@@ -122,6 +123,7 @@ Default and recommended: 'auto' (best mode for the model).
       finishReason = generateResult.finishReason;
       usage = generateResult.usage;
       warnings = generateResult.warnings;
+      rawResponse = generateResult.rawResponse;
       logprobs = generateResult.logprobs;
 
       break;
@@ -152,6 +154,7 @@ Default and recommended: 'auto' (best mode for the model).
       finishReason = generateResult.finishReason;
       usage = generateResult.usage;
       warnings = generateResult.warnings;
+      rawResponse = generateResult.rawResponse;
       logprobs = generateResult.logprobs;
 
       break;
@@ -192,6 +195,7 @@ Default and recommended: 'auto' (best mode for the model).
       finishReason = generateResult.finishReason;
       usage = generateResult.usage;
       warnings = generateResult.warnings;
+      rawResponse = generateResult.rawResponse;
       logprobs = generateResult.logprobs;
 
       break;
@@ -218,6 +222,7 @@ Default and recommended: 'auto' (best mode for the model).
     finishReason,
     usage: calculateTokenUsage(usage),
     warnings,
+    rawResponse,
     logprobs,
   });
 }
@@ -246,6 +251,16 @@ Warnings from the model provider (e.g. unsupported settings)
    */
   readonly warnings: LanguageModelV1CallWarning[] | undefined;
 
+  /**
+Optional raw response data.
+   */
+  rawResponse?: {
+    /**
+Response headers.
+ */
+    headers?: Record<string, string>;
+  };
+
   /**
 Logprobs for the completion. 
 `undefined` if the mode does not support logprobs or if was not enabled
@@ -257,12 +272,16 @@ Logprobs for the completion.
     finishReason: LanguageModelV1FinishReason;
     usage: TokenUsage;
     warnings: LanguageModelV1CallWarning[] | undefined;
+    rawResponse?: {
+      headers?: Record<string, string>;
+    };
     logprobs: LanguageModelV1LogProbs | undefined;
   }) {
     this.object = options.object;
     this.finishReason = options.finishReason;
     this.usage = options.usage;
     this.warnings = options.warnings;
+    this.rawResponse = options.rawResponse;
     this.logprobs = options.logprobs;
   }
 }

packages/core/core/generate-object/stream-object.ts

@@ -220,6 +220,7 @@ Default and recommended: 'auto' (best mode for the model).
   return new StreamObjectResult({
     stream: result.stream.pipeThrough(new TransformStream(transformer)),
     warnings: result.warnings,
+    rawResponse: result.rawResponse,
   });
 }
 
@@ -259,15 +260,30 @@ Warnings from the model provider (e.g. unsupported settings)
    */
   readonly warnings: LanguageModelV1CallWarning[] | undefined;
 
+  /**
+Optional raw response data.
+   */
+  rawResponse?: {
+    /**
+Response headers.
+ */
+    headers?: Record<string, string>;
+  };
+
   constructor({
     stream,
     warnings,
+    rawResponse,
   }: {
     stream: ReadableStream<string | ObjectStreamPartInput>;
     warnings: LanguageModelV1CallWarning[] | undefined;
+    rawResponse?: {
+      headers?: Record<string, string>;
+    };
   }) {
     this.originalStream = stream;
     this.warnings = warnings;
+    this.rawResponse = rawResponse;
   }
 
   get partialObjectStream(): AsyncIterableStream<DeepPartial<T>> {

packages/core/core/generate-text/generate-text.ts

@@ -116,6 +116,7 @@ The tools that the model can call. The model needs to support calling tools.
     finishReason: modelResponse.finishReason,
     usage: calculateTokenUsage(modelResponse.usage),
     warnings: modelResponse.warnings,
+    rawResponse: modelResponse.rawResponse,
     logprobs: modelResponse.logprobs,
   });
 }
@@ -188,6 +189,16 @@ Warnings from the model provider (e.g. unsupported settings)
    */
   readonly warnings: LanguageModelV1CallWarning[] | undefined;
 
+  /**
+Optional raw response data.
+   */
+  rawResponse?: {
+    /**
+Response headers.
+   */
+    headers?: Record<string, string>;
+  };
+
   /**
 Logprobs for the completion. 
 `undefined` if the mode does not support logprobs or if was not enabled
@@ -201,6 +212,9 @@ Logprobs for the completion.
     finishReason: LanguageModelV1FinishReason;
     usage: TokenUsage;
     warnings: LanguageModelV1CallWarning[] | undefined;
+    rawResponse?: {
+      headers?: Record<string, string>;
+    };
     logprobs: LanguageModelV1LogProbs | undefined;
   }) {
     this.text = options.text;
@@ -209,6 +223,7 @@ Logprobs for the completion.
     this.finishReason = options.finishReason;
     this.usage = options.usage;
     this.warnings = options.warnings;
+    this.rawResponse = options.rawResponse;
     this.logprobs = options.logprobs;
   }
 }

packages/core/core/generate-text/stream-text.test.ts

@@ -4,9 +4,8 @@ import { convertArrayToReadableStream } from '../test/convert-array-to-readable-
 import { convertAsyncIterableToArray } from '../test/convert-async-iterable-to-array';
 import { convertReadableStreamToArray } from '../test/convert-readable-stream-to-array';
 import { MockLanguageModelV1 } from '../test/mock-language-model-v1';
-import { experimental_streamText } from './stream-text';
-import { ServerResponse } from 'node:http';
 import { createMockServerResponse } from '../test/mock-server-response';
+import { experimental_streamText } from './stream-text';
 
 describe('result.textStream', () => {
   it('should send text deltas', async () => {

packages/core/core/generate-text/stream-text.ts

@@ -85,7 +85,7 @@ The tools that the model can call. The model needs to support calling tools.
   }): Promise<StreamTextResult<TOOLS>> {
   const retry = retryWithExponentialBackoff({ maxRetries });
   const validatedPrompt = getValidatedPrompt({ system, prompt, messages });
-  const { stream, warnings } = await retry(() =>
+  const { stream, warnings, rawResponse } = await retry(() =>
     model.doStream({
       mode: {
         type: 'regular',
@@ -112,6 +112,7 @@ The tools that the model can call. The model needs to support calling tools.
       generatorStream: stream,
     }),
     warnings,
+    rawResponse,
   });
 }
 
@@ -152,15 +153,30 @@ Warnings from the model provider (e.g. unsupported settings)
    */
   readonly warnings: LanguageModelV1CallWarning[] | undefined;
 
+  /**
+Optional raw response data.
+   */
+  rawResponse?: {
+    /**
+Response headers.
+     */
+    headers?: Record<string, string>;
+  };
+
   constructor({
     stream,
     warnings,
+    rawResponse,
   }: {
     stream: ReadableStream<TextStreamPart<TOOLS>>;
     warnings: LanguageModelV1CallWarning[] | undefined;
+    rawResponse?: {
+      headers?: Record<string, string>;
+    };
   }) {
     this.originalStream = stream;
     this.warnings = warnings;
+    this.rawResponse = rawResponse;
   }
 
   /**