Add batch size arg (langchain-ai#3310)

malandis · Nov 19, 2023 · 0a2b8a6 · 0a2b8a6
1 parent ca2bc81
commit 0a2b8a6
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 10 deletions.
diff --git a/langchain/src/chat_models/base.ts b/langchain/src/chat_models/base.ts
@@ -1,15 +1,15 @@
 import {
   AIMessage,
   BaseMessage,
+  BaseMessageChunk,
+  BaseMessageLike,
   BasePromptValue,
   ChatGeneration,
+  ChatGenerationChunk,
   ChatResult,
   HumanMessage,
-  BaseMessageChunk,
   LLMResult,
   RUN_KEY,
-  ChatGenerationChunk,
-  BaseMessageLike,
   coerceMessageLikeToMessage,
 } from "../schema/index.js";
 import {
@@ -23,6 +23,7 @@ import {
   CallbackManagerForLLMRun,
   Callbacks,
 } from "../callbacks/manager.js";
+
 import { RunnableConfig } from "../schema/runnable/config.js";
 
 /**
@@ -164,6 +165,7 @@ export abstract class BaseChatModel<
       const extra = {
         options: callOptions,
         invocation_params: this?.invocationParams(callOptions),
+        batch_size: 1,
       };
       const runManagers = await callbackManager_?.handleChatModelStart(
         this.toJSON(),
@@ -231,6 +233,7 @@ export abstract class BaseChatModel<
     const extra = {
       options: parsedOptions,
       invocation_params: this?.invocationParams(parsedOptions),
+      batch_size: 1,
     };
     const runManagers = await callbackManager_?.handleChatModelStart(
       this.toJSON(),

diff --git a/langchain/src/llms/base.ts b/langchain/src/llms/base.ts
@@ -7,20 +7,21 @@ import {
   LLMResult,
   RUN_KEY,
 } from "../schema/index.js";
-import {
-  BaseLanguageModel,
-  BaseLanguageModelCallOptions,
-  BaseLanguageModelInput,
-  BaseLanguageModelParams,
-} from "../base_language/index.js";
 import {
   BaseCallbackConfig,
   CallbackManager,
   CallbackManagerForLLMRun,
   Callbacks,
 } from "../callbacks/manager.js";
-import { getBufferString } from "../memory/base.js";
+import {
+  BaseLanguageModel,
+  BaseLanguageModelCallOptions,
+  BaseLanguageModelInput,
+  BaseLanguageModelParams,
+} from "../base_language/index.js";
+
 import { RunnableConfig } from "../schema/runnable/config.js";
+import { getBufferString } from "../memory/base.js";
 
 export type SerializedLLM = {
   _model: string;
@@ -120,6 +121,7 @@ export abstract class BaseLLM<
       const extra = {
         options: callOptions,
         invocation_params: this?.invocationParams(callOptions),
+        batch_size: 1,
       };
       const runManagers = await callbackManager_?.handleLLMStart(
         this.toJSON(),
@@ -247,6 +249,7 @@ export abstract class BaseLLM<
     const extra = {
       options: parsedOptions,
       invocation_params: this?.invocationParams(parsedOptions),
+      batch_size: prompts.length,
     };
     const runManagers = await callbackManager_?.handleLLMStart(
       this.toJSON(),