Add test, revert prebuilt

langchain-ai · Aug 28, 2024 · 532597d · 532597d
1 parent 645e5b2
commit 532597d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 15 deletions.
diff --git a/libs/langgraph/src/prebuilt/react_agent_executor.ts b/libs/langgraph/src/prebuilt/react_agent_executor.ts
@@ -21,8 +21,9 @@ import {
 import { ChatPromptTemplate } from "@langchain/core/prompts";
 import { BaseCheckpointSaver } from "@langchain/langgraph-checkpoint";
 import {
+  END,
   messagesStateReducer,
-  type START,
+  START,
   StateGraph,
 } from "../graph/index.js";
 import { MessagesAnnotation } from "../graph/messages_annotation.js";
@@ -109,9 +110,9 @@ export function createReactAgent(
       isAIMessage(lastMessage) &&
       (!lastMessage.tool_calls || lastMessage.tool_calls.length === 0)
     ) {
-      return "__end__";
+      return END;
     } else {
-      return "tools";
+      return "continue";
     }
   };
 
@@ -129,8 +130,11 @@ export function createReactAgent(
       RunnableLambda.from(callModel).withConfig({ runName: "agent" })
     )
     .addNode("tools", new ToolNode<AgentState>(toolClasses))
-    .addEdge("__start__", "agent")
-    .addConditionalEdges("agent", shouldContinue, ["tools", "__end__"])
+    .addEdge(START, "agent")
+    .addConditionalEdges("agent", shouldContinue, {
+      continue: "tools",
+      [END]: END,
+    })
     .addEdge("tools", "agent");
 
   return workflow.compile({

diff --git a/libs/langgraph/src/tests/pregel.test.ts b/libs/langgraph/src/tests/pregel.test.ts
@@ -2036,9 +2036,9 @@ describe("StateGraph", () => {
     data: typeof AgentAnnotation.State
   ): Promise<string> => {
     if (data.agentOutcome && "returnValues" in data.agentOutcome) {
-      return "exit";
+      return "__end__";
     }
-    return "continue";
+    return "tools";
   };
 
   it("can invoke", async () => {
@@ -2088,10 +2088,7 @@ describe("StateGraph", () => {
       .addNode("tools", executeTools)
       .addEdge(START, "agent")
       .addEdge("agent", "passthrough")
-      .addConditionalEdges("passthrough", shouldContinue, {
-        continue: "tools",
-        exit: END,
-      })
+      .addConditionalEdges("passthrough", shouldContinue, ["tools", "__end__"])
       .addEdge("tools", "agent")
       .compile();
 
@@ -2182,10 +2179,7 @@ describe("StateGraph", () => {
       .addNode("agent", agent)
       .addNode("tools", executeTools)
       .addEdge(START, "agent")
-      .addConditionalEdges("agent", shouldContinue, {
-        continue: "tools",
-        exit: END,
-      })
+      .addConditionalEdges("agent", shouldContinue, ["tools", "__end__"])
       .addEdge("tools", "agent")
       .compile();