Pass wrapped config object instead

langchain-ai · Feb 14, 2024 · 779ad09 · 779ad09
1 parent ef2cf6b
commit 779ad09
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 22 deletions.
diff --git a/js/package.json b/js/package.json
@@ -38,7 +38,7 @@
     "build:esm": "tsc --outDir dist/ && rm -rf dist/tests dist/**/tests",
     "build:cjs": "tsc --outDir dist-cjs/ -p tsconfig.cjs.json && node scripts/move-cjs-to-dist.js && rm -r dist-cjs",
     "test": "cross-env NODE_OPTIONS=--experimental-vm-modules jest --passWithNoTests --testPathIgnorePatterns='\\.int\\.test.[tj]s' --testTimeout 30000",
-    "test:integration": "cross-env NODE_OPTIONS=--experimental-vm-modules jest --testPathPattern=\\.int\\.test.ts --testTimeout 100000",
+    "test:integration": "cross-env NODE_OPTIONS=--experimental-vm-modules jest --testPathPattern=traceable_wrapper\\.int\\.test.ts --testTimeout 100000",
     "test:single": "NODE_OPTIONS=--experimental-vm-modules yarn run jest --config jest.config.cjs --testTimeout 100000",
     "lint": "eslint 'src/**/*.{ts,tsx}' --quiet --fix",
     "format": "prettier --write 'src/**/*.{ts,tsx}'",

diff --git a/js/src/client.ts b/js/src/client.ts
@@ -582,7 +582,6 @@ export class Client {
       runCreate,
     ]);
 
-    console.log(mergedRunCreateParams[0]);
     const response = await this.caller.call(fetch, `${this.apiUrl}/runs`, {
       method: "POST",
       headers,

diff --git a/js/src/run_helpers.ts b/js/src/run_helpers.ts
@@ -1,6 +1,8 @@
 import { RunTree, RunTreeConfig, isRunTree } from "./run_trees.js";
 import { KVMap } from "./schemas.js";
 
+type TraceableLastArg = RunTree | { config: Partial<RunTreeConfig> } | "root";
+
 /**
  * Higher-order function for creating or adding a run to a run tree.
  * The returned function will now expect the first argument to be a run tree
@@ -17,16 +19,20 @@ export function traceable<WrappedFunc extends (...args: any[]) => any>(
   type Output = Awaited<ReturnType<WrappedFunc>>;
 
   const traceableFunc: TraceableFunction<Inputs, Output> = async (
-    ...args: [...Inputs, RunTree | RunTreeConfig]
+    ...args: [...Inputs, TraceableLastArg]
   ): Promise<Output> => {
     let currentRunTree: RunTree;
 
-    const inputRunTree: RunTree | RunTreeConfig = args[args.length - 1];
+    const inputRunTree: TraceableLastArg = args[args.length - 1];
     const rawInputs = args.slice(0, args.length - 1) as Inputs;
 
-    if (!isRunTree(inputRunTree) && !isKVMap(inputRunTree)) {
+    if (
+      !isRunTree(inputRunTree) &&
+      inputRunTree !== "root" &&
+      !(isKVMap(inputRunTree) && "config" in inputRunTree)
+    ) {
       throw new Error(
-        "Last argument must be a RunTree or RunTreeConfig with a name"
+        "Last argument must be a RunTree instance or a config object with RunTree constructor arguments"
       );
     }
 
@@ -37,8 +43,13 @@ export function traceable<WrappedFunc extends (...args: any[]) => any>(
 
     if (isRunTree(inputRunTree)) {
       currentRunTree = await inputRunTree.createChild(ensuredConfig);
+    } else if (inputRunTree === "root") {
+      currentRunTree = new RunTree(ensuredConfig);
     } else {
-      currentRunTree = new RunTree({ ...ensuredConfig, ...inputRunTree });
+      currentRunTree = new RunTree({
+        ...ensuredConfig,
+        ...inputRunTree.config,
+      });
     }
 
     let inputs: KVMap;
@@ -91,7 +102,7 @@ export function traceable<WrappedFunc extends (...args: any[]) => any>(
 }
 
 export type TraceableFunction<Inputs extends unknown[], Output> = (
-  ...inputs: [...params: Inputs, runTree: RunTree | RunTreeConfig]
+  ...inputs: [...params: Inputs, runTree: TraceableLastArg]
 ) => Promise<Output>;
 
 export function isTraceableFunction(

diff --git a/js/src/tests/traceable_wrapper.int.test.ts b/js/src/tests/traceable_wrapper.int.test.ts
@@ -1,7 +1,6 @@
 import { Client } from "../client.js";
 import { traceable } from "../run_helpers.js";
 import { RunTree } from "../run_trees.js";
-import { v4 as uuidv4 } from "uuid";
 
 // async function deleteProject(langchainClient: Client, projectName: string) {
 //   try {
@@ -19,33 +18,43 @@ test.concurrent(
       callerOptions: { maxRetries: 0 },
     });
     const projectName = "__test_traceable_wrapper";
-    const runId = uuidv4();
-    const rootRunTree = new RunTree({
-      name: "test_run_tree",
-      run_type: "chain",
-      id: runId,
-      client: langchainClient,
-      project_name: projectName,
-    });
-
     const addValueTraceable = traceable(
       (a: string, b: number) => {
         return a + b;
       },
       { name: "testinger" }
     );
 
-    expect(await addValueTraceable("testing", 9, rootRunTree)).toBe("testing9");
+    expect(
+      await addValueTraceable("testing", 9, {
+        config: {
+          name: "test_run_tree_single",
+          run_type: "chain",
+          client: langchainClient,
+          project_name: projectName,
+        },
+      })
+    ).toBe("testing9");
 
     const entryTraceable = traceable(
       (complex: { value: string }, runTree: RunTree) =>
         addValueTraceable(complex.value, 1, runTree),
       { name: "nested_testinger" }
     );
 
-    expect(await entryTraceable({ value: "testing" }, rootRunTree)).toBe(
-      "testing1"
-    );
+    expect(
+      await entryTraceable(
+        { value: "testing" },
+        {
+          config: {
+            name: "test_run_tree_composed",
+            run_type: "chain",
+            client: langchainClient,
+            project_name: projectName,
+          },
+        }
+      )
+    ).toBe("testing1");
 
     // await deleteProject(langchainClient, projectName);
   },