From 6e970fb73e18976440ea7be9b0bede961c8c26d8 Mon Sep 17 00:00:00 2001
From: Brace Sproul <braceasproul@gmail.com>
Date: Mon, 15 Jan 2024 17:41:19 -0800
Subject: [PATCH] core[patch]: Allow runnable func to have async (#4033)

* core[patch]: Allow runnable func to have async

* chore: lint files

* cr

* undo cr

* pass config through

* tests

* chore: lint files
---
 langchain-core/src/runnables/passthrough.ts   |  8 ++++--
 .../tests/runnable_passthrough.test.ts        | 28 +++++++++++++++++++
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/langchain-core/src/runnables/passthrough.ts b/langchain-core/src/runnables/passthrough.ts
index 3074eff4b838..3e8f3cdeb9a3 100644
--- a/langchain-core/src/runnables/passthrough.ts
+++ b/langchain-core/src/runnables/passthrough.ts
@@ -10,7 +10,9 @@ import type { RunnableConfig } from "./config.js";
 // eslint-disable-next-line @typescript-eslint/no-explicit-any
 type RunnablePassthroughFunc<RunInput = any> =
   | ((input: RunInput) => void)
-  | ((input: RunInput, config?: RunnableConfig) => void);
+  | ((input: RunInput, config?: RunnableConfig) => void)
+  | ((input: RunInput) => Promise<void>)
+  | ((input: RunInput, config?: RunnableConfig) => Promise<void>);
 
 /**
  * A runnable to passthrough inputs unchanged or with additional keys.
@@ -64,7 +66,7 @@ export class RunnablePassthrough<RunInput> extends Runnable<
     options?: Partial<RunnableConfig>
   ): Promise<RunInput> {
     if (this.func) {
-      this.func(input);
+      await this.func(input, options);
     }
 
     return this._callWithConfig(
@@ -103,7 +105,7 @@ export class RunnablePassthrough<RunInput> extends Runnable<
     }
 
     if (this.func && finalOutput !== undefined) {
-      this.func(finalOutput);
+      await this.func(finalOutput, options);
     }
   }
 
diff --git a/langchain-core/src/runnables/tests/runnable_passthrough.test.ts b/langchain-core/src/runnables/tests/runnable_passthrough.test.ts
index 18f6e8635fd0..551ff3769b9c 100644
--- a/langchain-core/src/runnables/tests/runnable_passthrough.test.ts
+++ b/langchain-core/src/runnables/tests/runnable_passthrough.test.ts
@@ -3,6 +3,7 @@ import { FakeChatModel } from "../../utils/testing/index.js";
 import { RunnablePassthrough } from "../passthrough.js";
 import { JsonOutputParser } from "../../output_parsers/json.js";
 import { RunnableSequence } from "../base.js";
+import { RunnableConfig } from "../config.js";
 
 test("RunnablePassthrough can call .assign and pass prev result through", async () => {
   const promptTemplate = PromptTemplate.fromTemplate("{input}");
@@ -81,3 +82,30 @@ test("RunnablePassthrough can transform a function as constructor args", async (
   expect(results).toEqual([1]);
   expect(wasCalled).toEqual(true);
 });
+
+test("RunnablePassthrough can invoke a function and pass through config", async () => {
+  let wasCalled = false;
+  let addOneResult: number = 0;
+  const addOne = (input: number, config?: RunnableConfig) => {
+    wasCalled = true;
+    if (
+      !config?.configurable?.number ??
+      Number.isNaN(config?.configurable?.number)
+    ) {
+      throw new Error("configurable.number is NaN");
+    }
+    console.log(config.configurable.number);
+    addOneResult = input + config.configurable.number;
+  };
+  const passthrough = new RunnablePassthrough<number>({
+    func: addOne,
+  });
+  const result = await passthrough.invoke(1, {
+    configurable: {
+      number: 1,
+    },
+  });
+  expect(result).toEqual(1);
+  expect(addOneResult).toEqual(2);
+  expect(wasCalled).toEqual(true);
+});