fix: fix the problem that the usage field in the streaming response was not handled correctly (#2173)

.changeset/weak-ties-listen.md

@@ -0,0 +1,5 @@
+---
+"@llamaindex/openai": patch
+---
+
+fix: fix the problem that the usage field in the streaming response was not handled correctly

packages/providers/openai/src/llm.ts

@@ -370,25 +370,18 @@ export class OpenAI extends ToolCallLLM<OpenAIAdditionalChatOptions> {
     let currentToolCall: PartialToolCall | null = null;
     const toolCallMap = new Map<string, PartialToolCall>();
     for await (const part of stream) {
-      if (part.choices.length === 0) {
+      const choice = part.choices && part.choices[0];
+      const hasValidContent =
+        choice?.delta?.content ||
+        choice?.delta?.tool_calls ||
+        choice?.finish_reason;
+
+      if (!hasValidContent) {
         if (part.usage) {
-          yield {
-            raw: part,
-            delta: "",
-          };
+          yield { raw: part, delta: "" };
         }
         continue;
       }
-      const choice = part.choices[0]!;
-      // skip parts that don't have any content
-      if (
-        !(
-          choice.delta?.content ||
-          choice.delta?.tool_calls ||
-          choice.finish_reason
-        )
-      )
-        continue;
 
       let shouldEmitToolCall: PartialToolCall | null = null;
       if (

packages/providers/openai/tests/chat.test.ts

@@ -283,4 +283,100 @@ describe("OpenAI streamChat", () => {
     expect(chunks[0].options).toEqual({});
     expect(chunks[0].delta).toBe("");
   });
+
+  it("should handle part with undefined choices", async () => {
+    // Create a mock stream that yields a part without choices
+    const mockStream = async function* () {
+      yield {
+        // No choices property defined
+        usage: {
+          prompt_tokens: 10,
+          completion_tokens: 5,
+          total_tokens: 15,
+        },
+      };
+    };
+
+    // Mock the OpenAI session and chat completions
+    const mockSession = {
+      chat: {
+        completions: {
+          create: vi.fn().mockResolvedValue(mockStream()),
+        },
+      },
+    };
+
+    const openai = new OpenAI({
+      model: "gpt-4o-mini",
+      apiKey: "test-key",
+      // @ts-expect-error: mockSession is a mock object for testing purposes
+      session: mockSession,
+    });
+
+    // @ts-expect-error accessing protected method
+    const stream = openai.streamChat({
+      messages: [{ role: "user" as const, content: "Hello" }],
+      stream: true,
+    });
+
+    const chunks: ChatResponseChunk[] = [];
+    for await (const chunk of stream) {
+      chunks.push(chunk);
+    }
+
+    expect(chunks).toHaveLength(1);
+    expect(chunks[0].delta).toBe("");
+    expect(chunks[0].raw).toHaveProperty("usage");
+  });
+
+  it("should handle part with invalid content", async () => {
+    const mockStream = async function* () {
+      yield {
+        choices: [
+          {
+            delta: {
+              role: "assistant",
+              content: "",
+            },
+          },
+        ],
+        usage: {
+          prompt_tokens: 10,
+          completion_tokens: 5,
+          total_tokens: 15,
+        },
+      };
+    };
+
+    // Mock the OpenAI session and chat completions
+    const mockSession = {
+      chat: {
+        completions: {
+          create: vi.fn().mockResolvedValue(mockStream()),
+        },
+      },
+    };
+
+    const openai = new OpenAI({
+      model: "gpt-4o-mini",
+      apiKey: "test-key",
+      // @ts-expect-error: mockSession is a mock object for testing purposes
+      session: mockSession,
+    });
+
+    // @ts-expect-error accessing protected method
+    const stream = openai.streamChat({
+      messages: [{ role: "user" as const, content: "Hello" }],
+      stream: true,
+    });
+
+    const chunks: ChatResponseChunk[] = [];
+    for await (const chunk of stream) {
+      chunks.push(chunk);
+    }
+
+    expect(chunks).toHaveLength(1);
+    expect(chunks[0].delta).toBe("");
+    expect(chunks[0].raw).toHaveProperty("usage");
+  });
 });