Add streaming support for LLM and Chat Model calls with multiple prompts or completions (#1760)

* Add streaming support for LLM and Chat Model calls with multiple prompts or completions

- multiple prompts are batch calls to generate()
- multiple completions are calls with n>1

* Undo unrelated change

* Update docstring

* Rename

* Lint

Author: nfcampos

langchain/src/callbacks/base.ts

@@ -22,6 +22,11 @@ export interface BaseCallbackHandlerInput {
   ignoreAgent?: boolean;
 }
 
+export interface NewTokenIndices {
+  prompt: number;
+  completion: number;
+}
+
 abstract class BaseCallbackHandlerMethodsClass {
   /**
    * Called at the start of an LLM or Chat Model run, with the prompt(s)
@@ -41,6 +46,13 @@ abstract class BaseCallbackHandlerMethodsClass {
    */
   handleLLMNewToken?(
     token: string,
+    /**
+     * idx.prompt is the index of the prompt that produced the token
+     *   (if there are multiple prompts)
+     * idx.completion is the index of the completion that produced the token
+     *   (if multiple completions per prompt are requested)
+     */
+    idx: NewTokenIndices,
     runId: string,
     parentRunId?: string
   ): Promise<void> | void;

langchain/src/callbacks/index.ts

@@ -2,6 +2,7 @@ export {
   BaseCallbackHandler,
   CallbackHandlerMethods,
   BaseCallbackHandlerInput,
+  NewTokenIndices,
 } from "./base.js";
 
 export { Run, RunType, BaseTracer } from "./handlers/tracer.js";

langchain/src/callbacks/manager.ts

@@ -6,7 +6,11 @@ import {
   ChainValues,
   LLMResult,
 } from "../schema/index.js";
-import { BaseCallbackHandler, CallbackHandlerMethods } from "./base.js";
+import {
+  BaseCallbackHandler,
+  CallbackHandlerMethods,
+  NewTokenIndices,
+} from "./base.js";
 import { ConsoleCallbackHandler } from "./handlers/console.js";
 import {
   getTracingCallbackHandler,
@@ -79,14 +83,18 @@ export class CallbackManagerForLLMRun
   extends BaseRunManager
   implements BaseCallbackManagerMethods
 {
-  async handleLLMNewToken(token: string): Promise<void> {
+  async handleLLMNewToken(
+    token: string,
+    idx: NewTokenIndices = { prompt: 0, completion: 0 }
+  ): Promise<void> {
     await Promise.all(
       this.handlers.map((handler) =>
         consumeCallback(async () => {
           if (!handler.ignoreLLM) {
             try {
               await handler.handleLLMNewToken?.(
                 token,
+                idx,
                 this.runId,
                 this._parentRunId
               );

langchain/src/chat_models/base.ts

@@ -90,7 +90,11 @@ export abstract class BaseChatModel extends BaseLanguageModel {
     // generate results
     const results = await Promise.allSettled(
       messages.map((messageList, i) =>
-        this._generate(messageList, parsedOptions, runManagers?.[i])
+        this._generate(
+          messageList,
+          { ...parsedOptions, promptIndex: i },
+          runManagers?.[i]
+        )
       )
     );
     // handle results

langchain/src/chat_models/openai.ts

@@ -86,6 +86,7 @@ export interface ChatOpenAICallOptions extends OpenAICallOptions {
   function_call?: CreateChatCompletionRequestFunctionCall;
   functions?: ChatCompletionFunctions[];
   tools?: StructuredTool[];
+  promptIndex?: number;
 }
 
 /**
@@ -113,7 +114,15 @@ export class ChatOpenAI
   declare CallOptions: ChatOpenAICallOptions;
 
   get callKeys(): (keyof ChatOpenAICallOptions)[] {
-    return ["stop", "signal", "timeout", "options", "functions", "tools"];
+    return [
+      "stop",
+      "signal",
+      "timeout",
+      "options",
+      "functions",
+      "tools",
+      "promptIndex",
+    ];
   }
 
   lc_serializable = true;
@@ -223,10 +232,6 @@ export class ChatOpenAI
 
     this.streaming = fields?.streaming ?? false;
 
-    if (this.streaming && this.n > 1) {
-      throw new Error("Cannot stream results when n > 1");
-    }
-
     if (this.azureOpenAIApiKey) {
       if (!this.azureOpenAIApiInstanceName) {
         throw new Error("Azure OpenAI API instance name not found");
@@ -408,11 +413,13 @@ export class ChatOpenAI
                         choice.message.function_call.arguments +=
                           part.delta?.function_call?.arguments ?? "";
                       }
-                      // TODO this should pass part.index to the callback
-                      // when that's supported there
                       // eslint-disable-next-line no-void
                       void runManager?.handleLLMNewToken(
-                        part.delta?.content ?? ""
+                        part.delta?.content ?? "",
+                        {
+                          prompt: options.promptIndex ?? 0,
+                          completion: part.index,
+                        }
                       );
                       // TODO we don't currently have a callback method for
                       // sending the function call arguments

langchain/src/chat_models/tests/chatopenai.int.test.ts

@@ -14,6 +14,7 @@ import {
   SystemMessagePromptTemplate,
 } from "../../prompts/index.js";
 import { CallbackManager } from "../../callbacks/index.js";
+import { NewTokenIndices } from "../../callbacks/base.js";
 
 test("Test ChatOpenAI", async () => {
   const chat = new ChatOpenAI({ modelName: "gpt-3.5-turbo", maxTokens: 10 });
@@ -129,11 +130,43 @@ test("Test ChatOpenAI in streaming mode", async () => {
     ],
   });
   const message = new HumanChatMessage("Hello!");
-  const res = await model.call([message]);
-  console.log({ res });
+  const result = await model.call([message]);
+  console.log(result);
+
+  expect(nrNewTokens > 0).toBe(true);
+  expect(result.text).toBe(streamedCompletion);
+});
+
+test("Test ChatOpenAI in streaming mode with n > 1 and multiple prompts", async () => {
+  let nrNewTokens = 0;
+  const streamedCompletions = [
+    ["", ""],
+    ["", ""],
+  ];
+
+  const model = new ChatOpenAI({
+    modelName: "gpt-3.5-turbo",
+    streaming: true,
+    maxTokens: 10,
+    n: 2,
+    callbacks: [
+      {
+        async handleLLMNewToken(token: string, idx: NewTokenIndices) {
+          nrNewTokens += 1;
+          streamedCompletions[idx.prompt][idx.completion] += token;
+        },
+      },
+    ],
+  });
+  const message1 = new HumanChatMessage("Hello!");
+  const message2 = new HumanChatMessage("Bye!");
+  const result = await model.generate([[message1], [message2]]);
+  console.log(result.generations);
 
   expect(nrNewTokens > 0).toBe(true);
-  expect(res.text).toBe(streamedCompletion);
+  expect(result.generations.map((g) => g.map((gg) => gg.text))).toEqual(
+    streamedCompletions
+  );
 });
 
 test("Test ChatOpenAI prompt value", async () => {

langchain/src/llms/base.ts

@@ -332,8 +332,10 @@ export abstract class LLM extends BaseLLM {
     runManager?: CallbackManagerForLLMRun
   ): Promise<LLMResult> {
     const generations: Generation[][] = await Promise.all(
-      prompts.map((prompt) =>
-        this._call(prompt, options, runManager).then((text) => [{ text }])
+      prompts.map((prompt, promptIndex) =>
+        this._call(prompt, { ...options, promptIndex }, runManager).then(
+          (text) => [{ text }]
+        )
       )
     );
     return { generations };

langchain/src/llms/openai-chat.ts

@@ -20,7 +20,11 @@ import { CallbackManagerForLLMRun } from "../callbacks/manager.js";
 import { Generation, LLMResult } from "../schema/index.js";
 import { promptLayerTrackRequest } from "../util/prompt-layer.js";
 
-export { OpenAICallOptions, OpenAIChatInput, AzureOpenAIInput };
+export { OpenAIChatInput, AzureOpenAIInput };
+
+export interface OpenAIChatCallOptions extends OpenAICallOptions {
+  promptIndex?: number;
+}
 
 /**
  * Wrapper around OpenAI large language models that use the Chat endpoint.
@@ -48,10 +52,10 @@ export class OpenAIChat
   extends LLM
   implements OpenAIChatInput, AzureOpenAIInput
 {
-  declare CallOptions: OpenAICallOptions;
+  declare CallOptions: OpenAIChatCallOptions;
 
-  get callKeys(): (keyof OpenAICallOptions)[] {
-    return ["stop", "signal", "timeout", "options"];
+  get callKeys(): (keyof OpenAIChatCallOptions)[] {
+    return ["stop", "signal", "timeout", "options", "promptIndex"];
   }
 
   lc_serializable = true;
@@ -166,8 +170,10 @@ export class OpenAIChat
 
     this.streaming = fields?.streaming ?? false;
 
-    if (this.streaming && this.n > 1) {
-      throw new Error("Cannot stream results when n > 1");
+    if (this.n > 1) {
+      throw new Error(
+        "Cannot use n > 1 in OpenAIChat LLM. Use ChatOpenAI Chat Model instead."
+      );
     }
 
     if (this.azureOpenAIApiKey) {
@@ -329,7 +335,11 @@ export class OpenAIChat
                       choice.message.content += part.delta?.content ?? "";
                       // eslint-disable-next-line no-void
                       void runManager?.handleLLMNewToken(
-                        part.delta?.content ?? ""
+                        part.delta?.content ?? "",
+                        {
+                          prompt: options.promptIndex ?? 0,
+                          completion: part.index,
+                        }
                       );
                     }
                   }

langchain/src/llms/openai.ts

@@ -88,7 +88,7 @@ export class OpenAI extends BaseLLM implements OpenAIInput, AzureOpenAIInput {
 
   n = 1;
 
-  bestOf = 1;
+  bestOf?: number;
 
   logitBias?: Record<string, number>;
 
@@ -179,11 +179,7 @@ export class OpenAI extends BaseLLM implements OpenAIInput, AzureOpenAIInput {
 
     this.streaming = fields?.streaming ?? false;
 
-    if (this.streaming && this.n > 1) {
-      throw new Error("Cannot stream results when n > 1");
-    }
-
-    if (this.streaming && this.bestOf > 1) {
+    if (this.streaming && this.bestOf && this.bestOf > 1) {
       throw new Error("Cannot stream results when bestOf > 1");
     }
 
@@ -345,10 +341,11 @@ export class OpenAI extends BaseLLM implements OpenAIInput, AzureOpenAIInput {
                         choice.text = (choice.text ?? "") + (part.text ?? "");
                         choice.finish_reason = part.finish_reason;
                         choice.logprobs = part.logprobs;
-                        // TODO this should pass part.index to the callback
-                        // when that's supported there
                         // eslint-disable-next-line no-void
-                        void runManager?.handleLLMNewToken(part.text ?? "");
+                        void runManager?.handleLLMNewToken(part.text ?? "", {
+                          prompt: Math.floor(part.index / this.n),
+                          completion: part.index % this.n,
+                        });
                       }
                     }
 

langchain/src/llms/tests/openai.int.test.ts

@@ -4,6 +4,7 @@ import { OpenAIChat } from "../openai-chat.js";
 import { OpenAI } from "../openai.js";
 import { StringPromptValue } from "../../prompts/index.js";
 import { CallbackManager } from "../../callbacks/index.js";
+import { NewTokenIndices } from "../../callbacks/base.js";
 
 test("Test OpenAI", async () => {
   const model = new OpenAI({ maxTokens: 5, modelName: "text-ada-001" });
@@ -144,26 +145,63 @@ test("Test OpenAI in streaming mode", async () => {
 
 test("Test OpenAI in streaming mode with multiple prompts", async () => {
   let nrNewTokens = 0;
+  const completions = [
+    ["", ""],
+    ["", ""],
+  ];
 
   const model = new OpenAI({
     maxTokens: 5,
     modelName: "text-ada-001",
     streaming: true,
+    n: 2,
     callbacks: CallbackManager.fromHandlers({
-      async handleLLMNewToken(_token: string) {
+      async handleLLMNewToken(token: string, idx: NewTokenIndices) {
         nrNewTokens += 1;
+        completions[idx.prompt][idx.completion] += token;
       },
     }),
   });
   const res = await model.generate(["Print hello world", "print hello sea"]);
-  console.log({ res });
+  console.log(
+    res.generations,
+    res.generations.map((g) => g[0].generationInfo)
+  );
 
   expect(nrNewTokens > 0).toBe(true);
   expect(res.generations.length).toBe(2);
-  expect(res.generations.map((g) => typeof g[0].text === "string")).toEqual([
-    true,
-    true,
-  ]);
+  expect(res.generations.map((g) => g.map((gg) => gg.text))).toEqual(
+    completions
+  );
+});
+
+test("Test OpenAIChat in streaming mode with multiple prompts", async () => {
+  let nrNewTokens = 0;
+  const completions = [[""], [""]];
+
+  const model = new OpenAI({
+    maxTokens: 5,
+    modelName: "gpt-3.5-turbo",
+    streaming: true,
+    n: 1,
+    callbacks: CallbackManager.fromHandlers({
+      async handleLLMNewToken(token: string, idx: NewTokenIndices) {
+        nrNewTokens += 1;
+        completions[idx.prompt][idx.completion] += token;
+      },
+    }),
+  });
+  const res = await model.generate(["Print hello world", "print hello sea"]);
+  console.log(
+    res.generations,
+    res.generations.map((g) => g[0].generationInfo)
+  );
+
+  expect(nrNewTokens > 0).toBe(true);
+  expect(res.generations.length).toBe(2);
+  expect(res.generations.map((g) => g.map((gg) => gg.text))).toEqual(
+    completions
+  );
 });
 
 test("Test OpenAI prompt value", async () => {