fix(providers/xai): return actual usage when streaming instead of NaN (#5873) (#5891)

Co-authored-by: Lars Grammel <lars.grammel@gmail.com>

Author: samdenty

.changeset/fast-students-turn.md

@@ -0,0 +1,6 @@
+---
+'@ai-sdk/openai-compatible': patch
+'@ai-sdk/xai': patch
+---
+
+fix(providers/xai): return actual usage when streaming instead of NaN

packages/openai-compatible/src/openai-compatible-chat-language-model.test.ts

@@ -655,6 +655,27 @@ describe('doGenerate', () => {
       expect(warnings).toEqual([]);
     });
 
+    it('should respect the includeUsage option', async () => {
+      prepareJsonResponse({ content: '{"value":"test"}' });
+
+      const model = new OpenAICompatibleChatLanguageModel('gpt-4o-2024-08-06', {
+        provider: 'test-provider',
+        url: () => 'https://my.api.com/v1/chat/completions',
+        headers: () => ({}),
+        includeUsage: true,
+      });
+
+      await model.doStream({
+        inputFormat: 'prompt',
+        prompt: TEST_PROMPT,
+      });
+
+      const body = await server.calls[0].requestBody;
+
+      expect(body.stream).toBe(true);
+      expect(body.stream_options).toStrictEqual({ include_usage: true });
+    });
+
     it('should use json_schema & strict with responseFormat json when structuredOutputs are enabled', async () => {
       prepareJsonResponse({ content: '{"value":"Spark"}' });
 
@@ -1708,9 +1729,32 @@ describe('doStream', () => {
       prompt: TEST_PROMPT,
     });
 
-    expect(request).toStrictEqual({
-      body: '{"model":"grok-beta","messages":[{"role":"user","content":"Hello"}],"stream":true}',
-    });
+    expect(request).toMatchInlineSnapshot(`
+      {
+        "body": {
+          "frequency_penalty": undefined,
+          "max_tokens": undefined,
+          "messages": [
+            {
+              "content": "Hello",
+              "role": "user",
+            },
+          ],
+          "model": "grok-beta",
+          "presence_penalty": undefined,
+          "response_format": undefined,
+          "seed": undefined,
+          "stop": undefined,
+          "stream": true,
+          "stream_options": undefined,
+          "temperature": undefined,
+          "tool_choice": undefined,
+          "tools": undefined,
+          "top_p": undefined,
+          "user": undefined,
+        },
+      }
+    `);
   });
 
   describe('usage details in streaming', () => {

packages/openai-compatible/src/openai-compatible-chat-language-model.ts

@@ -41,6 +41,7 @@ export type OpenAICompatibleChatConfig = {
   headers: () => Record<string, string | undefined>;
   url: (options: { modelId: string; path: string }) => string;
   fetch?: FetchFunction;
+  includeUsage?: boolean;
   errorStructure?: ProviderErrorStructure<any>;
   metadataExtractor?: MetadataExtractor;
 
@@ -304,7 +305,16 @@ export class OpenAICompatibleChatLanguageModel implements LanguageModelV2 {
   ): Promise<Awaited<ReturnType<LanguageModelV2['doStream']>>> {
     const { args, warnings } = this.getArgs({ ...options });
 
-    const body = JSON.stringify({ ...args, stream: true });
+    const body = {
+      ...args,
+      stream: true,
+
+      // only include stream_options when in strict compatibility mode:
+      stream_options: this.config.includeUsage
+        ? { include_usage: true }
+        : undefined,
+    };
+
     const metadataExtractor =
       this.config.metadataExtractor?.createStreamExtractor();
 
@@ -314,10 +324,7 @@ export class OpenAICompatibleChatLanguageModel implements LanguageModelV2 {
         modelId: this.modelId,
       }),
       headers: combineHeaders(this.config.headers(), options.headers),
-      body: {
-        ...args,
-        stream: true,
-      },
+      body,
       failedResponseHandler: this.failedResponseHandler,
       successfulResponseHandler: createEventSourceResponseHandler(
         this.chunkSchema,

packages/openai-compatible/src/openai-compatible-completion-language-model.test.ts

@@ -543,6 +543,7 @@ describe('doStream', () => {
           "seed": undefined,
           "stop": undefined,
           "stream": true,
+          "stream_options": undefined,
           "suffix": undefined,
           "temperature": undefined,
           "top_p": undefined,

packages/openai-compatible/src/openai-compatible-completion-language-model.ts

@@ -33,6 +33,7 @@ import {
 
 type OpenAICompatibleCompletionConfig = {
   provider: string;
+  includeUsage?: boolean;
   headers: () => Record<string, string | undefined>;
   url: (options: { modelId: string; path: string }) => string;
   fetch?: FetchFunction;
@@ -220,6 +221,11 @@ export class OpenAICompatibleCompletionLanguageModel
     const body = {
       ...args,
       stream: true,
+
+      // only include stream_options when in strict compatibility mode:
+      stream_options: this.config.includeUsage
+        ? { include_usage: true }
+        : undefined,
     };
 
     const { responseHeaders, value: response } = await postJsonToApi({

packages/xai/src/xai-provider.test.ts

@@ -86,6 +86,21 @@ describe('xAIProvider', () => {
 
       expect(model).toBeInstanceOf(OpenAICompatibleChatLanguageModel);
     });
+
+    it('should pass the includeUsage option to the chat model, to make sure usage is reported while streaming', () => {
+      const provider = createXai();
+      const modelId = 'xai-chat-model';
+
+      const model = provider.chat(modelId);
+
+      expect(model).toBeInstanceOf(OpenAICompatibleChatLanguageModel);
+
+      const constructorCall =
+        OpenAICompatibleChatLanguageModelMock.mock.calls[0];
+
+      expect(constructorCall[0]).toBe(modelId);
+      expect(constructorCall[1].includeUsage).toBe(true);
+    });
   });
 
   describe('imageModel', () => {

packages/xai/src/xai-provider.ts

@@ -98,6 +98,7 @@ export function createXai(options: XaiProviderSettings = {}): XaiProvider {
       fetch: options.fetch,
       errorStructure: xaiErrorStructure,
       supportsStructuredOutputs: structuredOutputs,
+      includeUsage: true,
     });
   };