feat(langgraph): Clean up traces (langchain-ai#627)

benjamincburns · Oct 24, 2024 · 70453f2 · 70453f2
1 parent 74ee400
commit 70453f2
Show file tree

Hide file tree

Showing 14 changed files with 64 additions and 180 deletions.
diff --git a/examples/package.json b/examples/package.json
@@ -8,7 +8,7 @@
   "devDependencies": {
     "@langchain/anthropic": "^0.3.0",
     "@langchain/community": "^0.3.0",
-    "@langchain/core": "^0.3.10",
+    "@langchain/core": "^0.3.14",
     "@langchain/groq": "^0.1.1",
     "@langchain/langgraph": "workspace:*",
     "@langchain/mistralai": "^0.1.0",

diff --git a/libs/langgraph/package.json b/libs/langgraph/package.json
@@ -43,7 +43,7 @@
     "@jest/globals": "^29.5.0",
     "@langchain/anthropic": "^0.3.0",
     "@langchain/community": "^0.3.0",
-    "@langchain/core": "^0.3.10",
+    "@langchain/core": "^0.3.14",
     "@langchain/langgraph-checkpoint-postgres": "workspace:*",
     "@langchain/langgraph-checkpoint-sqlite": "workspace:*",
     "@langchain/openai": "^0.3.0",

diff --git a/libs/langgraph/src/graph/graph.ts b/libs/langgraph/src/graph/graph.ts
@@ -80,6 +80,7 @@ export class Branch<
   ) {
     return ChannelWrite.registerWriter(
       new RunnableCallable({
+        trace: false,
         func: async (input: IO, config: CallOptions) => {
           try {
             return await this._route(input, config, writer, reader);

diff --git a/libs/langgraph/src/graph/state.ts b/libs/langgraph/src/graph/state.ts
@@ -316,9 +316,20 @@ export class StateGraph<
       this._addSchema(options.input.spec);
     }
 
-    const runnable = _coerceToRunnable(action) as unknown as Runnable<S, U>;
+    let runnable;
+    if (Runnable.isRunnable(action)) {
+      runnable = action;
+    } else if (typeof action === "function") {
+      runnable = new RunnableCallable({
+        func: action,
+        name: key,
+        trace: false,
+      });
+    } else {
+      runnable = _coerceToRunnable(action);
+    }
     const nodeSpec: StateGraphNodeSpec<S, U> = {
-      runnable,
+      runnable: runnable as unknown as Runnable<S, U>,
       retryPolicy: options?.retryPolicy,
       metadata: options?.metadata,
       input: options?.input?.spec ?? this._schemaDefinition,

diff --git a/libs/langgraph/src/prebuilt/react_agent_executor.ts b/libs/langgraph/src/prebuilt/react_agent_executor.ts
@@ -163,10 +163,7 @@ export function createReactAgent(
   const workflow = new StateGraph<AgentState>({
     channels: schema,
   })
-    .addNode(
-      "agent",
-      RunnableLambda.from(callModel).withConfig({ runName: "agent" })
-    )
+    .addNode("agent", callModel)
     .addNode("tools", new ToolNode<AgentState>(toolClasses))
     .addEdge(START, "agent")
     .addConditionalEdges("agent", shouldContinue, {

diff --git a/libs/langgraph/src/prebuilt/tool_node.ts b/libs/langgraph/src/prebuilt/tool_node.ts
@@ -139,6 +139,8 @@ export class ToolNode<T = any> extends RunnableCallable<T, T> {
 
   handleToolErrors = true;
 
+  trace = false;
+
   constructor(
     tools: (StructuredToolInterface | RunnableToolLike)[],
     options?: ToolNodeOptions

diff --git a/libs/langgraph/src/pregel/index.ts b/libs/langgraph/src/pregel/index.ts
@@ -732,8 +732,10 @@ export class Pregel<
       name: asNode,
       input: values,
       proc:
-        // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        writers.length > 1 ? RunnableSequence.from(writers as any) : writers[0],
+        writers.length > 1
+          ? // eslint-disable-next-line @typescript-eslint/no-explicit-any
+            RunnableSequence.from(writers as any, { omitSequenceTags: true })
+          : writers[0],
       writes: [],
       triggers: [INTERRUPT],
       id: uuid5(INTERRUPT, checkpoint.id),

diff --git a/libs/langgraph/src/pregel/read.ts b/libs/langgraph/src/pregel/read.ts
@@ -194,12 +194,14 @@ export class PregelNode<
         first: writers[0],
         middle: writers.slice(1, writers.length - 1),
         last: writers[writers.length - 1],
+        omitSequenceTags: true,
       });
     } else if (writers.length > 0) {
       return new RunnableSequence({
         first: this.bound,
         middle: writers.slice(0, writers.length - 1),
         last: writers[writers.length - 1],
+        omitSequenceTags: true,
       });
     } else {
       return this.bound;
@@ -237,7 +239,7 @@ export class PregelNode<
         channels: this.channels,
         triggers: this.triggers,
         mapper: this.mapper,
-        writers: [...this.writers, coerceable as ChannelWrite],
+        writers: [...this.writers, coerceable],
         bound: this.bound as unknown as PregelNode<
           RunInput,
           Exclude<NewRunOutput, Error>

diff --git a/libs/langgraph/src/pregel/write.ts b/libs/langgraph/src/pregel/write.ts
@@ -121,7 +121,7 @@ export class ChannelWrite<
     write(filtered);
   }
 
-  static isWriter(runnable: RunnableLike): boolean {
+  static isWriter(runnable: RunnableLike): runnable is ChannelWrite {
     return (
       // eslint-disable-next-line no-instanceof/no-instanceof
       runnable instanceof ChannelWrite ||

diff --git a/libs/langgraph/src/tests/prebuilt.test.ts b/libs/langgraph/src/tests/prebuilt.test.ts
@@ -319,7 +319,7 @@ describe("ToolNode", () => {
       })
     );
     await wrapper.invoke({}, { callbacks: callbackManager });
-    expect(runnableStartCount).toEqual(3);
+    expect(runnableStartCount).toEqual(2);
   });
 
   it("Should work in a state graph", async () => {

diff --git a/libs/langgraph/src/tests/pregel.test.ts b/libs/langgraph/src/tests/pregel.test.ts
@@ -31,6 +31,7 @@ import {
   ToolMessage,
 } from "@langchain/core/messages";
 import { ToolCall } from "@langchain/core/messages/tool";
+import { awaitAllCallbacks } from "@langchain/core/callbacks/promises";
 import {
   BaseCheckpointSaver,
   BaseStore,
@@ -2351,13 +2352,14 @@ export function runPregelTests(
           callbacks: [
             {
               handleChainEnd(outputs) {
+                // The final time this is called should be the final output from graph
                 callbackOutputs = outputs;
               },
             },
           ],
         }
       );
-      await new Promise((resolve) => setTimeout(resolve, 200));
+      await awaitAllCallbacks();
       expect(result).toEqual({
         input: "what is the weather in sf?",
         agentOutcome: {
@@ -4174,8 +4176,8 @@ export function runPregelTests(
             id: "action",
             type: "runnable",
             data: {
-              id: ["langchain_core", "runnables", "RunnableLambda"],
-              name: "RunnableLambda",
+              id: ["langgraph", "RunnableCallable"],
+              name: "action",
             },
           },
           {
@@ -4330,8 +4332,8 @@ export function runPregelTests(
             id: "action",
             type: "runnable",
             data: {
-              id: ["langchain_core", "runnables", "RunnableLambda"],
-              name: "RunnableLambda",
+              id: ["langgraph", "RunnableCallable"],
+              name: "action",
             },
           },
         ]),
@@ -5013,39 +5015,39 @@ export function runPregelTests(
     expect(tool.getGraph().toJSON()).toMatchObject({
       nodes: expect.arrayContaining([
         expect.objectContaining({ id: "__start__", type: "schema" }),
-        expect.objectContaining({ id: "__end__", type: "schema" }),
         {
           id: "prepare",
           type: "runnable",
           data: {
-            id: ["langchain_core", "runnables", "RunnableLambda"],
-            name: "RunnableLambda",
+            id: ["langgraph", "RunnableCallable"],
+            name: "prepare",
           },
         },
         {
           id: "tool_two_slow",
           type: "runnable",
           data: {
-            id: ["langchain_core", "runnables", "RunnableLambda"],
-            name: "RunnableLambda",
+            id: ["langgraph", "RunnableCallable"],
+            name: "tool_two_slow",
           },
         },
         {
           id: "tool_two_fast",
           type: "runnable",
           data: {
-            id: ["langchain_core", "runnables", "RunnableLambda"],
-            name: "RunnableLambda",
+            id: ["langgraph", "RunnableCallable"],
+            name: "tool_two_fast",
           },
         },
         {
           id: "finish",
           type: "runnable",
           data: {
-            id: ["langchain_core", "runnables", "RunnableLambda"],
-            name: "RunnableLambda",
+            id: ["langgraph", "RunnableCallable"],
+            name: "finish",
           },
         },
+        expect.objectContaining({ id: "__end__", type: "schema" }),
       ]),
       edges: expect.arrayContaining([
         { source: "__start__", target: "prepare", conditional: false },