mirror of
https://github.com/Sea-Haven-Industries/open-swe.git
synced 2026-09-30 09:13:14 +00:00
feat: Pause sandboxes before ending the graph
This commit is contained in:
parent
424124c286
commit
68ab03647c
2 changed files with 34 additions and 2 deletions
19
src/index.ts
19
src/index.ts
|
|
@ -1,5 +1,10 @@
|
|||
import { END, START, StateGraph } from "@langchain/langgraph";
|
||||
import { GraphAnnotation, GraphConfiguration, GraphState } from "./types.js";
|
||||
import {
|
||||
GraphAnnotation,
|
||||
GraphConfig,
|
||||
GraphConfiguration,
|
||||
GraphState,
|
||||
} from "./types.js";
|
||||
import {
|
||||
generatePlan,
|
||||
initialize,
|
||||
|
|
@ -9,6 +14,7 @@ import {
|
|||
interruptPlan,
|
||||
} from "./nodes/index.js";
|
||||
import { isAIMessage } from "@langchain/core/messages";
|
||||
import { pauseSandbox } from "./utils/sandbox.js";
|
||||
|
||||
/**
|
||||
* @param {GraphState} state - The current graph state.
|
||||
|
|
@ -33,7 +39,10 @@ function routeAfterPlan(state: GraphState): "interrupt-plan" | typeof END {
|
|||
* @param {GraphState} state - The current graph state.
|
||||
* @returns {typeof END | "take-action"} The next node to execute, or END if the process should stop.
|
||||
*/
|
||||
function takeActionOrEnd(state: GraphState): typeof END | "take-action" {
|
||||
async function takeActionOrEnd(
|
||||
state: GraphState,
|
||||
config: GraphConfig,
|
||||
): Promise<typeof END | "take-action"> {
|
||||
const { messages } = state;
|
||||
const lastMessage = messages[messages.length - 1];
|
||||
|
||||
|
|
@ -41,6 +50,12 @@ function takeActionOrEnd(state: GraphState): typeof END | "take-action" {
|
|||
if (isAIMessage(lastMessage) && lastMessage.tool_calls?.length) {
|
||||
return "take-action";
|
||||
}
|
||||
|
||||
// First, pause the sandbox before ending the graph.
|
||||
if (config.configurable?.sandbox_session_id) {
|
||||
await pauseSandbox(config.configurable.sandbox_session_id);
|
||||
}
|
||||
|
||||
return END;
|
||||
}
|
||||
|
||||
|
|
|
|||
17
src/utils/sandbox.ts
Normal file
17
src/utils/sandbox.ts
Normal file
|
|
@ -0,0 +1,17 @@
|
|||
import { Sandbox } from "@e2b/code-interpreter";
|
||||
|
||||
/**
|
||||
* Pauses the sandbox. Either pass an existing sandbox client, or a sandbox session ID.
|
||||
* If no sandbox client is provided, the sandbox will be connected to.
|
||||
|
||||
* @param sandboxSessionId The ID of the sandbox to pause.
|
||||
* @param sandbox The sandbox client to pause. If not provided, the sandbox will be connected to.
|
||||
* @returns The sandbox session ID.
|
||||
*/
|
||||
export async function pauseSandbox(
|
||||
sandboxSessionId: string,
|
||||
sandbox?: Sandbox,
|
||||
): Promise<string> {
|
||||
const sandboxClient = sandbox ?? (await Sandbox.connect(sandboxSessionId));
|
||||
return await sandboxClient.pause();
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue