fix: type

himself65 · Apr 10, 2024 · 71cc4cf · 71cc4cf
1 parent 6277441
commit 71cc4cf
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 9 deletions.
diff --git a/examples/toolsStream.ts b/examples/toolsStream.ts
@@ -1,4 +1,8 @@
-import { ChatResponseChunk, OpenAI } from "llamaindex";
+import {
+  ChatResponseChunk,
+  OpenAI,
+  OpenAIAdditionalChatOptions,
+} from "llamaindex";
 
 async function main() {
   const llm = new OpenAI({ model: "gpt-4-turbo" });
@@ -34,11 +38,11 @@ async function main() {
   };
 
   const stream = await llm.chat({ ...args, stream: true });
-  let chunk: ChatResponseChunk | null = null;
+  let chunk: ChatResponseChunk<OpenAIAdditionalChatOptions>;
   for await (chunk of stream) {
     process.stdout.write(chunk.delta);
   }
-  console.log(chunk?.additionalKwargs?.toolCalls[0]);
+  console.log(chunk.options.toolCalls[0]);
 }
 
 (async function () {

diff --git a/packages/core/src/llm/open_ai.ts b/packages/core/src/llm/open_ai.ts
@@ -377,7 +377,7 @@ export class OpenAI extends BaseLLM<
   @wrapEventCaller
   protected async *streamChat(
     baseRequestParams: OpenAILLM.Chat.ChatCompletionCreateParams,
-  ): AsyncIterable<ChatResponseChunk> {
+  ): AsyncIterable<ChatResponseChunk<OpenAIAdditionalMessageOptions>> {
     const stream: AsyncIterable<OpenAILLM.Chat.ChatCompletionChunk> =
       await this.session.openai.chat.completions.create({
         ...baseRequestParams,
@@ -403,7 +403,7 @@ export class OpenAI extends BaseLLM<
 
       yield {
         // add tool calls to final chunk
-        options: toolCalls.length > 0 ? { toolCalls: toolCalls } : undefined,
+        options: toolCalls.length > 0 ? { toolCalls: toolCalls } : {},
         delta: choice.delta.content ?? "",
       };
     }

diff --git a/packages/core/src/llm/types.ts b/packages/core/src/llm/types.ts
@@ -130,10 +130,21 @@ export interface ChatResponse<
   raw: object;
 }
 
-export interface ChatResponseChunk {
-  delta: string;
-  options?: Record<string, any>;
-}
+export type ChatResponseChunk<
+  AdditionalMessageOptions extends Record<string, unknown> = Record<
+    string,
+    unknown
+  >,
+> =
+  AdditionalMessageOptions extends Record<string, unknown>
+    ? {
+        delta: string;
+        options?: AdditionalMessageOptions;
+      }
+    : {
+        delta: string;
+        options: AdditionalMessageOptions;
+      };
 
 export interface CompletionResponse {
   text: string;