langchain 0.0.93 → 0.0.95

package/dist/agents/agent.cjs

@@ -209,6 +209,9 @@ class Agent extends BaseSingleActionAgent {
             newInputs.stop = this._stop();
         }
         const output = await this.llmChain.predict(newInputs, callbackManager);
+        if (!this.outputParser) {
+            throw new Error("Output parser not set");
+        }
         return this.outputParser.parse(output, callbackManager);
     }
     /**

package/dist/agents/agent.d.ts

@@ -95,7 +95,7 @@ export interface AgentArgs {
  */
 export declare abstract class Agent extends BaseSingleActionAgent {
     llmChain: LLMChain;
-    outputParser: AgentActionOutputParser;
+    outputParser: AgentActionOutputParser | undefined;
     private _allowedTools?;
     get allowedTools(): string[] | undefined;
     get inputKeys(): string[];

package/dist/agents/agent.js

@@ -202,6 +202,9 @@ export class Agent extends BaseSingleActionAgent {
             newInputs.stop = this._stop();
         }
         const output = await this.llmChain.predict(newInputs, callbackManager);
+        if (!this.outputParser) {
+            throw new Error("Output parser not set");
+        }
         return this.outputParser.parse(output, callbackManager);
     }
     /**

package/dist/agents/index.cjs

@@ -1,6 +1,6 @@
 "use strict";
 Object.defineProperty(exports, "__esModule", { value: true });
-exports.StructuredChatOutputParserWithRetries = exports.StructuredChatOutputParser = exports.StructuredChatAgent = exports.AgentActionOutputParser = exports.ZeroShotAgentOutputParser = exports.ZeroShotAgent = exports.initializeAgentExecutorWithOptions = exports.initializeAgentExecutor = exports.AgentExecutor = exports.ChatConversationalAgentOutputParserWithRetries = exports.ChatConversationalAgentOutputParser = exports.ChatConversationalAgent = exports.ChatAgentOutputParser = exports.ChatAgent = exports.Toolkit = exports.createVectorStoreRouterAgent = exports.createVectorStoreAgent = exports.createSqlAgent = exports.createOpenApiAgent = exports.createJsonAgent = exports.ZapierToolKit = exports.VectorStoreToolkit = exports.VectorStoreRouterToolkit = exports.SqlToolkit = exports.RequestsToolkit = exports.OpenApiToolkit = exports.JsonToolkit = exports.LLMSingleActionAgent = exports.BaseSingleActionAgent = exports.Agent = void 0;
+exports.OpenAIAgent = exports.StructuredChatOutputParserWithRetries = exports.StructuredChatOutputParser = exports.StructuredChatAgent = exports.AgentActionOutputParser = exports.ZeroShotAgentOutputParser = exports.ZeroShotAgent = exports.initializeAgentExecutorWithOptions = exports.initializeAgentExecutor = exports.AgentExecutor = exports.ChatConversationalAgentOutputParserWithRetries = exports.ChatConversationalAgentOutputParser = exports.ChatConversationalAgent = exports.ChatAgentOutputParser = exports.ChatAgent = exports.Toolkit = exports.createVectorStoreRouterAgent = exports.createVectorStoreAgent = exports.createSqlAgent = exports.createOpenApiAgent = exports.createJsonAgent = exports.ZapierToolKit = exports.VectorStoreToolkit = exports.VectorStoreRouterToolkit = exports.SqlToolkit = exports.RequestsToolkit = exports.OpenApiToolkit = exports.JsonToolkit = exports.LLMSingleActionAgent = exports.BaseSingleActionAgent = exports.Agent = void 0;
 var agent_js_1 = require("./agent.cjs");
 Object.defineProperty(exports, "Agent", { enumerable: true, get: function () { return agent_js_1.Agent; } });
 Object.defineProperty(exports, "BaseSingleActionAgent", { enumerable: true, get: function () { return agent_js_1.BaseSingleActionAgent; } });
@@ -45,3 +45,5 @@ Object.defineProperty(exports, "StructuredChatAgent", { enumerable: true, get: f
 var outputParser_js_4 = require("./structured_chat/outputParser.cjs");
 Object.defineProperty(exports, "StructuredChatOutputParser", { enumerable: true, get: function () { return outputParser_js_4.StructuredChatOutputParser; } });
 Object.defineProperty(exports, "StructuredChatOutputParserWithRetries", { enumerable: true, get: function () { return outputParser_js_4.StructuredChatOutputParserWithRetries; } });
+var index_js_6 = require("./openai/index.cjs");
+Object.defineProperty(exports, "OpenAIAgent", { enumerable: true, get: function () { return index_js_6.OpenAIAgent; } });

package/dist/agents/index.d.ts

@@ -12,3 +12,4 @@ export { ZeroShotAgentOutputParser } from "./mrkl/outputParser.js";
 export { AgentActionOutputParser, AgentInput, SerializedAgent, SerializedAgentT, SerializedZeroShotAgent, StoppingMethod, } from "./types.js";
 export { StructuredChatAgent, StructuredChatAgentInput, StructuredChatCreatePromptArgs, } from "./structured_chat/index.js";
 export { StructuredChatOutputParser, StructuredChatOutputParserArgs, StructuredChatOutputParserWithRetries, } from "./structured_chat/outputParser.js";
+export { OpenAIAgent, OpenAIAgentInput, OpenAIAgentCreatePromptArgs, } from "./openai/index.js";

package/dist/agents/index.js

@@ -12,3 +12,4 @@ export { ZeroShotAgentOutputParser } from "./mrkl/outputParser.js";
 export { AgentActionOutputParser, } from "./types.js";
 export { StructuredChatAgent, } from "./structured_chat/index.js";
 export { StructuredChatOutputParser, StructuredChatOutputParserWithRetries, } from "./structured_chat/outputParser.js";
+export { OpenAIAgent, } from "./openai/index.js";

package/dist/agents/initialize.cjs

@@ -7,6 +7,7 @@ const index_js_2 = require("./chat_convo/index.cjs");
 const index_js_3 = require("./structured_chat/index.cjs");
 const executor_js_1 = require("./executor.cjs");
 const index_js_4 = require("./mrkl/index.cjs");
+const index_js_5 = require("./openai/index.cjs");
 /**
  * @deprecated use initializeAgentExecutorWithOptions instead
  */
@@ -78,6 +79,7 @@ async function initializeAgentExecutorWithOptions(tools, llm, options = {
                         returnMessages: true,
                         memoryKey: "chat_history",
                         inputKey: "input",
+                        outputKey: "output",
                     }),
                 ...rest,
             });
@@ -93,6 +95,22 @@ async function initializeAgentExecutorWithOptions(tools, llm, options = {
             });
             return executor;
         }
+        case "openai-functions": {
+            const { agentArgs, memory, ...rest } = options;
+            const executor = executor_js_1.AgentExecutor.fromAgentAndTools({
+                agent: index_js_5.OpenAIAgent.fromLLMAndTools(llm, tools, agentArgs),
+                tools,
+                memory: memory ??
+                    new buffer_memory_js_1.BufferMemory({
+                        returnMessages: true,
+                        memoryKey: "chat_history",
+                        inputKey: "input",
+                        outputKey: "output",
+                    }),
+                ...rest,
+            });
+            return executor;
+        }
         default: {
             throw new Error("Unknown agent type");
         }

package/dist/agents/initialize.d.ts

@@ -6,6 +6,7 @@ import { ChatConversationalAgent } from "./chat_convo/index.js";
 import { StructuredChatAgent } from "./structured_chat/index.js";
 import { AgentExecutor, AgentExecutorInput } from "./executor.js";
 import { ZeroShotAgent } from "./mrkl/index.js";
+import { OpenAIAgent } from "./openai/index.js";
 type AgentType = "zero-shot-react-description" | "chat-zero-shot-react-description" | "chat-conversational-react-description";
 /**
  * @deprecated use initializeAgentExecutorWithOptions instead
@@ -29,10 +30,13 @@ export type InitializeAgentExecutorOptions = ({
 /**
  * @interface
  */
-export type InitializeAgentExecutorOptionsStructured = {
+export type InitializeAgentExecutorOptionsStructured = ({
     agentType: "structured-chat-zero-shot-react-description";
     agentArgs?: Parameters<typeof StructuredChatAgent.fromLLMAndTools>[2];
-} & Omit<AgentExecutorInput, "agent" | "tools">;
+} & Omit<AgentExecutorInput, "agent" | "tools">) | ({
+    agentType: "openai-functions";
+    agentArgs?: Parameters<typeof OpenAIAgent.fromLLMAndTools>[2];
+} & Omit<AgentExecutorInput, "agent" | "tools">);
 /**
  * Initialize an agent executor with options
  * @param tools Array of tools to use in the agent

package/dist/agents/initialize.js

@@ -4,6 +4,7 @@ import { ChatConversationalAgent } from "./chat_convo/index.js";
 import { StructuredChatAgent } from "./structured_chat/index.js";
 import { AgentExecutor } from "./executor.js";
 import { ZeroShotAgent } from "./mrkl/index.js";
+import { OpenAIAgent } from "./openai/index.js";
 /**
  * @deprecated use initializeAgentExecutorWithOptions instead
  */
@@ -74,6 +75,7 @@ export async function initializeAgentExecutorWithOptions(tools, llm, options = {
                         returnMessages: true,
                         memoryKey: "chat_history",
                         inputKey: "input",
+                        outputKey: "output",
                     }),
                 ...rest,
             });
@@ -89,6 +91,22 @@ export async function initializeAgentExecutorWithOptions(tools, llm, options = {
             });
             return executor;
         }
+        case "openai-functions": {
+            const { agentArgs, memory, ...rest } = options;
+            const executor = AgentExecutor.fromAgentAndTools({
+                agent: OpenAIAgent.fromLLMAndTools(llm, tools, agentArgs),
+                tools,
+                memory: memory ??
+                    new BufferMemory({
+                        returnMessages: true,
+                        memoryKey: "chat_history",
+                        inputKey: "input",
+                        outputKey: "output",
+                    }),
+                ...rest,
+            });
+            return executor;
+        }
         default: {
             throw new Error("Unknown agent type");
         }

package/dist/agents/openai/index.cjs

@@ -0,0 +1,119 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.OpenAIAgent = void 0;
+const index_js_1 = require("../../schema/index.cjs");
+const agent_js_1 = require("../agent.cjs");
+const prompt_js_1 = require("./prompt.cjs");
+const chat_js_1 = require("../../prompts/chat.cjs");
+const llm_chain_js_1 = require("../../chains/llm_chain.cjs");
+function parseOutput(message) {
+    if (message.additional_kwargs.function_call) {
+        // eslint-disable-next-line prefer-destructuring
+        const function_call = message.additional_kwargs.function_call;
+        return {
+            tool: function_call.name,
+            toolInput: function_call.arguments
+                ? JSON.parse(function_call.arguments)
+                : {},
+            log: message.text,
+        };
+    }
+    else {
+        return { returnValues: { output: message.text }, log: message.text };
+    }
+}
+class OpenAIAgent extends agent_js_1.Agent {
+    _agentType() {
+        return "openai-functions";
+    }
+    observationPrefix() {
+        return "Observation: ";
+    }
+    llmPrefix() {
+        return "Thought:";
+    }
+    _stop() {
+        return ["Observation:"];
+    }
+    constructor(input) {
+        super({ ...input, outputParser: undefined });
+        Object.defineProperty(this, "lc_namespace", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: ["langchain", "agents", "openai"]
+        });
+        Object.defineProperty(this, "tools", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        this.tools = input.tools;
+    }
+    static createPrompt(_tools, fields) {
+        const { prefix = prompt_js_1.PREFIX } = fields || {};
+        return chat_js_1.ChatPromptTemplate.fromPromptMessages([
+            chat_js_1.SystemMessagePromptTemplate.fromTemplate(prefix),
+            new chat_js_1.MessagesPlaceholder("chat_history"),
+            chat_js_1.HumanMessagePromptTemplate.fromTemplate("{input}"),
+            new chat_js_1.MessagesPlaceholder("agent_scratchpad"),
+        ]);
+    }
+    static fromLLMAndTools(llm, tools, args) {
+        OpenAIAgent.validateTools(tools);
+        if (llm._modelType() !== "base_chat_model" || llm._llmType() !== "openai") {
+            throw new Error("OpenAIAgent requires an OpenAI chat model");
+        }
+        const prompt = OpenAIAgent.createPrompt(tools, args);
+        const chain = new llm_chain_js_1.LLMChain({
+            prompt,
+            llm,
+            callbacks: args?.callbacks,
+        });
+        return new OpenAIAgent({
+            llmChain: chain,
+            allowedTools: tools.map((t) => t.name),
+            tools,
+        });
+    }
+    async constructScratchPad(steps) {
+        return steps.flatMap(({ action, observation }) => [
+            new index_js_1.AIChatMessage("", {
+                function_call: {
+                    name: action.tool,
+                    arguments: JSON.stringify(action.toolInput),
+                },
+            }),
+            new index_js_1.FunctionChatMessage(observation, action.tool),
+        ]);
+    }
+    async plan(steps, inputs, callbackManager) {
+        // Add scratchpad and stop to inputs
+        const thoughts = await this.constructScratchPad(steps);
+        const newInputs = {
+            ...inputs,
+            agent_scratchpad: thoughts,
+        };
+        if (this._stop().length !== 0) {
+            newInputs.stop = this._stop();
+        }
+        // Split inputs between prompt and llm
+        const llm = this.llmChain.llm;
+        const valuesForPrompt = { ...newInputs };
+        const valuesForLLM = {
+            tools: this.tools,
+        };
+        for (const key of this.llmChain.llm.callKeys) {
+            if (key in inputs) {
+                valuesForLLM[key] = inputs[key];
+                delete valuesForPrompt[key];
+            }
+        }
+        const promptValue = await this.llmChain.prompt.formatPromptValue(valuesForPrompt);
+        const message = await llm.predictMessages(promptValue.toChatMessages(), valuesForLLM, callbackManager);
+        console.log("message", message);
+        return parseOutput(message);
+    }
+}
+exports.OpenAIAgent = OpenAIAgent;

package/dist/agents/openai/index.d.ts

@@ -0,0 +1,26 @@
+import { CallbackManager } from "../../callbacks/manager.js";
+import { BasePromptTemplate } from "../../prompts/base.js";
+import { AgentAction, AgentFinish, AgentStep, BaseChatMessage, ChainValues } from "../../schema/index.js";
+import { StructuredTool } from "../../tools/base.js";
+import { Agent, AgentArgs } from "../agent.js";
+import { AgentInput } from "../types.js";
+import { BaseLanguageModel } from "../../base_language/index.js";
+export interface OpenAIAgentInput extends AgentInput {
+    tools: StructuredTool[];
+}
+export interface OpenAIAgentCreatePromptArgs {
+    prefix?: string;
+}
+export declare class OpenAIAgent extends Agent {
+    lc_namespace: string[];
+    _agentType(): "openai-functions";
+    observationPrefix(): string;
+    llmPrefix(): string;
+    _stop(): string[];
+    tools: StructuredTool[];
+    constructor(input: Omit<OpenAIAgentInput, "outputParser">);
+    static createPrompt(_tools: StructuredTool[], fields?: OpenAIAgentCreatePromptArgs): BasePromptTemplate;
+    static fromLLMAndTools(llm: BaseLanguageModel, tools: StructuredTool[], args?: OpenAIAgentCreatePromptArgs & Pick<AgentArgs, "callbacks">): OpenAIAgent;
+    constructScratchPad(steps: AgentStep[]): Promise<string | BaseChatMessage[]>;
+    plan(steps: Array<AgentStep>, inputs: ChainValues, callbackManager?: CallbackManager): Promise<AgentAction | AgentFinish>;
+}

package/dist/agents/openai/index.js

@@ -0,0 +1,115 @@
+import { AIChatMessage, FunctionChatMessage, } from "../../schema/index.js";
+import { Agent } from "../agent.js";
+import { PREFIX } from "./prompt.js";
+import { ChatPromptTemplate, HumanMessagePromptTemplate, MessagesPlaceholder, SystemMessagePromptTemplate, } from "../../prompts/chat.js";
+import { LLMChain } from "../../chains/llm_chain.js";
+function parseOutput(message) {
+    if (message.additional_kwargs.function_call) {
+        // eslint-disable-next-line prefer-destructuring
+        const function_call = message.additional_kwargs.function_call;
+        return {
+            tool: function_call.name,
+            toolInput: function_call.arguments
+                ? JSON.parse(function_call.arguments)
+                : {},
+            log: message.text,
+        };
+    }
+    else {
+        return { returnValues: { output: message.text }, log: message.text };
+    }
+}
+export class OpenAIAgent extends Agent {
+    _agentType() {
+        return "openai-functions";
+    }
+    observationPrefix() {
+        return "Observation: ";
+    }
+    llmPrefix() {
+        return "Thought:";
+    }
+    _stop() {
+        return ["Observation:"];
+    }
+    constructor(input) {
+        super({ ...input, outputParser: undefined });
+        Object.defineProperty(this, "lc_namespace", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: ["langchain", "agents", "openai"]
+        });
+        Object.defineProperty(this, "tools", {
+            enumerable: true,
+            configurable: true,
+            writable: true,
+            value: void 0
+        });
+        this.tools = input.tools;
+    }
+    static createPrompt(_tools, fields) {
+        const { prefix = PREFIX } = fields || {};
+        return ChatPromptTemplate.fromPromptMessages([
+            SystemMessagePromptTemplate.fromTemplate(prefix),
+            new MessagesPlaceholder("chat_history"),
+            HumanMessagePromptTemplate.fromTemplate("{input}"),
+            new MessagesPlaceholder("agent_scratchpad"),
+        ]);
+    }
+    static fromLLMAndTools(llm, tools, args) {
+        OpenAIAgent.validateTools(tools);
+        if (llm._modelType() !== "base_chat_model" || llm._llmType() !== "openai") {
+            throw new Error("OpenAIAgent requires an OpenAI chat model");
+        }
+        const prompt = OpenAIAgent.createPrompt(tools, args);
+        const chain = new LLMChain({
+            prompt,
+            llm,
+            callbacks: args?.callbacks,
+        });
+        return new OpenAIAgent({
+            llmChain: chain,
+            allowedTools: tools.map((t) => t.name),
+            tools,
+        });
+    }
+    async constructScratchPad(steps) {
+        return steps.flatMap(({ action, observation }) => [
+            new AIChatMessage("", {
+                function_call: {
+                    name: action.tool,
+                    arguments: JSON.stringify(action.toolInput),
+                },
+            }),
+            new FunctionChatMessage(observation, action.tool),
+        ]);
+    }
+    async plan(steps, inputs, callbackManager) {
+        // Add scratchpad and stop to inputs
+        const thoughts = await this.constructScratchPad(steps);
+        const newInputs = {
+            ...inputs,
+            agent_scratchpad: thoughts,
+        };
+        if (this._stop().length !== 0) {
+            newInputs.stop = this._stop();
+        }
+        // Split inputs between prompt and llm
+        const llm = this.llmChain.llm;
+        const valuesForPrompt = { ...newInputs };
+        const valuesForLLM = {
+            tools: this.tools,
+        };
+        for (const key of this.llmChain.llm.callKeys) {
+            if (key in inputs) {
+                valuesForLLM[key] = inputs[key];
+                delete valuesForPrompt[key];
+            }
+        }
+        const promptValue = await this.llmChain.prompt.formatPromptValue(valuesForPrompt);
+        const message = await llm.predictMessages(promptValue.toChatMessages(), valuesForLLM, callbackManager);
+        console.log("message", message);
+        return parseOutput(message);
+    }
+}

package/dist/agents/openai/prompt.cjs

@@ -0,0 +1,5 @@
+"use strict";
+Object.defineProperty(exports, "__esModule", { value: true });
+exports.SUFFIX = exports.PREFIX = void 0;
+exports.PREFIX = `You are a helpful AI assistant.`;
+exports.SUFFIX = ``;

package/dist/agents/openai/prompt.d.ts

@@ -0,0 +1,2 @@
+export declare const PREFIX = "You are a helpful AI assistant.";
+export declare const SUFFIX = "";

package/dist/agents/openai/prompt.js

@@ -0,0 +1,2 @@
+export const PREFIX = `You are a helpful AI assistant.`;
+export const SUFFIX = ``;

package/dist/agents/types.d.ts

@@ -4,7 +4,7 @@ import { AgentAction, AgentFinish } from "../schema/index.js";
 import { BaseOutputParser } from "../schema/output_parser.js";
 export interface AgentInput {
     llmChain: LLMChain;
-    outputParser: AgentActionOutputParser;
+    outputParser: AgentActionOutputParser | undefined;
     allowedTools?: string[];
 }
 export declare abstract class AgentActionOutputParser extends BaseOutputParser<AgentAction | AgentFinish> {

package/dist/chains/base.cjs

@@ -39,11 +39,12 @@ class BaseChain extends index_js_2.BaseLangChain {
     async run(
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     input, callbacks) {
-        const isKeylessInput = this.inputKeys.length <= 1;
+        const inputKeys = this.inputKeys.filter((k) => !this.memory?.memoryKeys.includes(k) ?? true);
+        const isKeylessInput = inputKeys.length <= 1;
         if (!isKeylessInput) {
             throw new Error(`Chain ${this._chainType()} expects multiple inputs, cannot use 'run' `);
         }
-        const values = this.inputKeys.length ? { [this.inputKeys[0]]: input } : {};
+        const values = inputKeys.length ? { [inputKeys[0]]: input } : {};
         const returnValues = await this.call(values, callbacks);
         const keys = Object.keys(returnValues);
         if (keys.length === 1) {

package/dist/chains/base.js

@@ -36,11 +36,12 @@ export class BaseChain extends BaseLangChain {
     async run(
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     input, callbacks) {
-        const isKeylessInput = this.inputKeys.length <= 1;
+        const inputKeys = this.inputKeys.filter((k) => !this.memory?.memoryKeys.includes(k) ?? true);
+        const isKeylessInput = inputKeys.length <= 1;
         if (!isKeylessInput) {
             throw new Error(`Chain ${this._chainType()} expects multiple inputs, cannot use 'run' `);
         }
-        const values = this.inputKeys.length ? { [this.inputKeys[0]]: input } : {};
+        const values = inputKeys.length ? { [inputKeys[0]]: input } : {};
         const returnValues = await this.call(values, callbacks);
         const keys = Object.keys(returnValues);
         if (keys.length === 1) {

package/dist/chains/combine_docs_chain.cjs

@@ -42,17 +42,21 @@ class StuffDocumentsChain extends base_js_1.BaseChain {
         this.inputKey = fields.inputKey ?? this.inputKey;
     }
     /** @ignore */
-    async _call(values, runManager) {
+    _prepInputs(values) {
         if (!(this.inputKey in values)) {
             throw new Error(`Document key ${this.inputKey} not found.`);
         }
         const { [this.inputKey]: docs, ...rest } = values;
         const texts = docs.map(({ pageContent }) => pageContent);
         const text = texts.join("\n\n");
-        const result = await this.llmChain.call({
+        return {
             ...rest,
             [this.documentVariableName]: text,
-        }, runManager?.getChild("combine_documents"));
+        };
+    }
+    /** @ignore */
+    async _call(values, runManager) {
+        const result = await this.llmChain.call(this._prepInputs(values), runManager?.getChild("combine_documents"));
         return result;
     }
     _chainType() {
@@ -163,11 +167,11 @@ class MapReduceDocumentsChain extends base_js_1.BaseChain {
             const canSkipMapStep = i !== 0 || !this.ensureMapStep;
             if (canSkipMapStep) {
                 // Calculate the total tokens required in the input
-                const promises = inputs.map(async (i) => {
-                    const prompt = await this.llmChain.prompt.format(i);
-                    return this.llmChain.llm.getNumTokens(prompt);
-                });
-                const length = await Promise.all(promises).then((results) => results.reduce((a, b) => a + b, 0));
+                const formatted = await this.combineDocumentChain.llmChain.prompt.format(this.combineDocumentChain._prepInputs({
+                    [this.combineDocumentChain.inputKey]: currentDocs,
+                    ...rest,
+                }));
+                const length = await this.combineDocumentChain.llmChain.llm.getNumTokens(formatted);
                 const withinTokenLimit = length < this.maxTokens;
                 // If we can skip the map step, and we're within the token limit, we don't
                 // need to run the map step, so just break out of the loop.
@@ -193,7 +197,10 @@ class MapReduceDocumentsChain extends base_js_1.BaseChain {
         }
         // Now, with the final result of all the inputs from the `llmChain`, we can
         // run the `combineDocumentChain` over them.
-        const newInputs = { input_documents: currentDocs, ...rest };
+        const newInputs = {
+            [this.combineDocumentChain.inputKey]: currentDocs,
+            ...rest,
+        };
         const result = await this.combineDocumentChain.call(newInputs, runManager?.getChild("combine_documents"));
         // Return the intermediate steps results if the flag is set
         if (this.returnIntermediateSteps) {
@@ -213,7 +220,7 @@ class MapReduceDocumentsChain extends base_js_1.BaseChain {
         }
         return new MapReduceDocumentsChain({
             llmChain: await llm_chain_js_1.LLMChain.deserialize(data.llm_chain),
-            combineDocumentChain: await base_js_1.BaseChain.deserialize(data.combine_document_chain),
+            combineDocumentChain: await StuffDocumentsChain.deserialize(data.combine_document_chain),
         });
     }
     serialize() {

package/dist/chains/combine_docs_chain.d.ts

@@ -25,6 +25,8 @@ export declare class StuffDocumentsChain extends BaseChain implements StuffDocum
     get outputKeys(): string[];
     constructor(fields: StuffDocumentsChainInput);
     /** @ignore */
+    _prepInputs(values: ChainValues): ChainValues;
+    /** @ignore */
     _call(values: ChainValues, runManager?: CallbackManagerForChainRun): Promise<ChainValues>;
     _chainType(): "stuff_documents_chain";
     static deserialize(data: SerializedStuffDocumentsChain): Promise<StuffDocumentsChain>;
@@ -38,7 +40,7 @@ export interface MapReduceDocumentsChainInput extends StuffDocumentsChainInput {
     /** Ensures that the map step is taken regardless of max tokens */
     ensureMapStep?: boolean;
     /** Chain to use to combine results of applying llm_chain to documents. */
-    combineDocumentChain: BaseChain;
+    combineDocumentChain: StuffDocumentsChain;
     /** Return the results of the map steps in the output. */
     returnIntermediateSteps?: boolean;
 }
@@ -57,7 +59,7 @@ export declare class MapReduceDocumentsChain extends BaseChain implements MapRed
     maxTokens: number;
     maxIterations: number;
     ensureMapStep: boolean;
-    combineDocumentChain: BaseChain;
+    combineDocumentChain: StuffDocumentsChain;
     constructor(fields: MapReduceDocumentsChainInput);
     /** @ignore */
     _call(values: ChainValues, runManager?: CallbackManagerForChainRun): Promise<ChainValues>;

package/dist/chains/combine_docs_chain.js

@@ -39,17 +39,21 @@ export class StuffDocumentsChain extends BaseChain {
         this.inputKey = fields.inputKey ?? this.inputKey;
     }
     /** @ignore */
-    async _call(values, runManager) {
+    _prepInputs(values) {
         if (!(this.inputKey in values)) {
             throw new Error(`Document key ${this.inputKey} not found.`);
         }
         const { [this.inputKey]: docs, ...rest } = values;
         const texts = docs.map(({ pageContent }) => pageContent);
         const text = texts.join("\n\n");
-        const result = await this.llmChain.call({
+        return {
             ...rest,
             [this.documentVariableName]: text,
-        }, runManager?.getChild("combine_documents"));
+        };
+    }
+    /** @ignore */
+    async _call(values, runManager) {
+        const result = await this.llmChain.call(this._prepInputs(values), runManager?.getChild("combine_documents"));
         return result;
     }
     _chainType() {
@@ -159,11 +163,11 @@ export class MapReduceDocumentsChain extends BaseChain {
             const canSkipMapStep = i !== 0 || !this.ensureMapStep;
             if (canSkipMapStep) {
                 // Calculate the total tokens required in the input
-                const promises = inputs.map(async (i) => {
-                    const prompt = await this.llmChain.prompt.format(i);
-                    return this.llmChain.llm.getNumTokens(prompt);
-                });
-                const length = await Promise.all(promises).then((results) => results.reduce((a, b) => a + b, 0));
+                const formatted = await this.combineDocumentChain.llmChain.prompt.format(this.combineDocumentChain._prepInputs({
+                    [this.combineDocumentChain.inputKey]: currentDocs,
+                    ...rest,
+                }));
+                const length = await this.combineDocumentChain.llmChain.llm.getNumTokens(formatted);
                 const withinTokenLimit = length < this.maxTokens;
                 // If we can skip the map step, and we're within the token limit, we don't
                 // need to run the map step, so just break out of the loop.
@@ -189,7 +193,10 @@ export class MapReduceDocumentsChain extends BaseChain {
         }
         // Now, with the final result of all the inputs from the `llmChain`, we can
         // run the `combineDocumentChain` over them.
-        const newInputs = { input_documents: currentDocs, ...rest };
+        const newInputs = {
+            [this.combineDocumentChain.inputKey]: currentDocs,
+            ...rest,
+        };
         const result = await this.combineDocumentChain.call(newInputs, runManager?.getChild("combine_documents"));
         // Return the intermediate steps results if the flag is set
         if (this.returnIntermediateSteps) {
@@ -209,7 +216,7 @@ export class MapReduceDocumentsChain extends BaseChain {
         }
         return new MapReduceDocumentsChain({
             llmChain: await LLMChain.deserialize(data.llm_chain),
-            combineDocumentChain: await BaseChain.deserialize(data.combine_document_chain),
+            combineDocumentChain: await StuffDocumentsChain.deserialize(data.combine_document_chain),
         });
     }
     serialize() {

package/dist/chains/serde.d.ts

@@ -45,7 +45,7 @@ export type SerializedChatVectorDBQAChain = {
 export type SerializedMapReduceDocumentsChain = {
     _type: "map_reduce_documents_chain";
     llm_chain?: SerializedLLMChain;
-    combine_document_chain?: SerializedBaseChain;
+    combine_document_chain?: SerializedStuffDocumentsChain;
 };
 export type SerializedRefineDocumentsChain = {
     _type: "refine_documents_chain";