feat: Setup initialize node (#1)

* feat: Setup initialize node

* cr
This commit is contained in:
Brace Sproul 2025-05-21 17:53:43 -07:00 • committed by GitHub
parent e2415aee3f
commit 0830c9c3e8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 291 additions and 139 deletions

21
e2b.js.Dockerfile Normal file
View file

@ -0,0 +1,21 @@
# Always start from the base E2B image
FROM e2bdev/code-interpreter:latest
# Ensure git is installed
RUN apt-get update && apt-get install -y git curl
# Set the working directory
WORKDIR /app
# Download and install nvm, then Node.js, then Yarn
ENV NVM_DIR /root/.nvm
ENV NODE_VERSION 22
RUN curl -o- https://raw.githubusercontent.com/nvm-sh/nvm/v0.40.3/install.sh | bash \
&& . "$NVM_DIR/nvm.sh" \
&& nvm install "$NODE_VERSION" \
&& nvm alias default "$NODE_VERSION" \
&& nvm use default \
&& corepack enable yarn
# Add NVM's bin to the PATH for subsequent commands
ENV PATH $NVM_DIR/versions/node/v$NODE_VERSION/bin:$PATH

20
e2b.python.Dockerfile Normal file
View file

@ -0,0 +1,20 @@
# Always start from the base E2B image
FROM e2bdev/code-interpreter:latest
# Ensure git is installed
RUN apt-get update && apt-get install -y git curl
# Set the working directory
WORKDIR /app
# Ensure Python and pip are installed
RUN apt-get install -y python3 python3-pip python3-venv
# Upgrade pip
RUN python3 -m pip install --upgrade pip
# Install common Python development tools
RUN python3 -m pip install pytest black isort mypy
# Set Python 3 as the default python
RUN ln -sf /usr/bin/python3 /usr/bin/python

16
e2b.toml Normal file
View file

@ -0,0 +1,16 @@
# This is a config for E2B sandbox template.
# You can use template ID (eh0860emqx28qyxmbctu) to create a sandbox:
# Python SDK
# from e2b import Sandbox, AsyncSandbox
# sandbox = Sandbox("eh0860emqx28qyxmbctu") # Sync sandbox
# sandbox = await AsyncSandbox.create("eh0860emqx28qyxmbctu") # Async sandbox
# JS SDK
# import { Sandbox } from 'e2b'
# const sandbox = await Sandbox.create('eh0860emqx28qyxmbctu')
team_id = "af4eaf9c-3d40-415c-bb98-4831c47745cf"
start_cmd = "/root/.jupyter/start-up.sh"
dockerfile = "e2b.js.Dockerfile"
template_id = "eh0860emqx28qyxmbctu"

View file

@ -20,11 +20,13 @@
"test:single": "NODE_OPTIONS=--experimental-vm-modules yarn run jest --config jest.config.js --testTimeout 100000"
},
"dependencies": {
"@e2b/code-interpreter": "^1.2.0-beta.4",
"@langchain/anthropic": "^0.3.20",
"@langchain/core": "^0.3.56",
"@langchain/google-genai": "^0.2.9",
"@langchain/langgraph": "^0.2.73",
"@langchain/openai": "^0.5.10",
"diff": "^8.0.1",
"langchain": "^0.3.26",
"uuid": "^11.0.5",
"zod": "^3.23.8"
@ -32,20 +34,20 @@
"devDependencies": {
"@eslint/eslintrc": "^3.1.0",
"@eslint/js": "^9.19.0",
"@jest/globals": "^29.7.0",
"@langchain/langgraph-cli": "latest",
"@tsconfig/recommended": "^1.0.8",
"@types/jest": "^29.5.0",
"@types/node": "^22.13.5",
"dotenv": "^16.4.7",
"eslint": "^9.19.0",
"ts-jest": "^29.1.0",
"eslint-config-prettier": "^8.8.0",
"@types/jest": "^29.5.0",
"eslint-plugin-import": "^2.27.5",
"jest": "^29.7.0",
"eslint-plugin-no-instanceof": "^1.0.1",
"@jest/globals": "^29.7.0",
"eslint-plugin-prettier": "^4.2.1",
"jest": "^29.7.0",
"prettier": "^3.5.2",
"ts-jest": "^29.1.0",
"typescript": "~5.7.2",
"typescript-eslint": "^8.22.0"
},

4
src/constants.ts Normal file
View file

@ -0,0 +1,4 @@
// Always extend timeouts by 5 min after each command execution.
export const TIMEOUT_EXTENSION_OPT = {
timeoutMs: 300_000, // 5 minutes
};

View file

@ -6,7 +6,6 @@ import {
generateAction,
takeAction,
rewritePlan,
summarizeSession,
} from "./nodes/index.js";
import { isAIMessage, isToolMessage } from "@langchain/core/messages";
@ -31,8 +30,6 @@ function routeAfterPlan(state: GraphState): "generate-action" | "rewrite-plan" {
/**
* After taking action, ensure there is an AI message with tool calls.
* If there is, route to the take-action node. Otherwise, end the graph.
* We do not want to route to the summarize session node here, as the generation
* likely means the LLM is asking a question, or failed to generate an action.
*/
function takeActionOrEnd(state: GraphState): typeof END | "take-action" {
const { messages } = state;
@ -45,31 +42,12 @@ function takeActionOrEnd(state: GraphState): typeof END | "take-action" {
return END;
}
/**
* After taking action, check if there are any uncompleted plan items.
* If there are, route to the generate-action node. Otherwise, route to the summarize-session node.
* The summarize session node is the final node, and will end the graph.
*/
function generateActionOrSummarize(
state: GraphState,
): "generate-action" | "summarize-session" {
const { plan } = state;
const hasUncompletedPlanItems = !!plan.filter((item) => !item.completed)
.length;
if (hasUncompletedPlanItems) {
return "generate-action";
}
return "summarize-session";
}
const workflow = new StateGraph(GraphAnnotation, GraphConfiguration)
.addNode("initialize", initialize)
.addNode("generate-plan", generatePlan)
.addNode("rewrite-plan", rewritePlan)
.addNode("generate-action", generateAction)
.addNode("take-action", takeAction)
.addNode("summarize-session", summarizeSession)
.addEdge(START, "initialize")
.addEdge("initialize", "generate-plan")
.addConditionalEdges("generate-plan", routeAfterPlan, [
@ -82,11 +60,7 @@ const workflow = new StateGraph(GraphAnnotation, GraphConfiguration)
])
.addEdge("generate-plan", "generate-action")
.addConditionalEdges("generate-action", takeActionOrEnd, ["take-action", END])
.addConditionalEdges("take-action", generateActionOrSummarize, [
"generate-action",
"summarize-session",
])
.addEdge("summarize-session", END);
.addEdge("take-action", "generate-action");
export const graph = workflow.compile();
graph.name = "LangGraph ReAct MCP";

View file

@ -1,6 +1,5 @@
export * from "./generate-plan.js";
export * from "./initialize.js";
export * from "./generate-action.js";
export * from "./generate-message.js";
export * from "./take-action.js";
export * from "./rewrite-plan.js";
export * from "./summarize-session.js";

View file

@ -1,12 +1,86 @@
import { GraphState, GraphConfig, GraphUpdate } from "../types.js";
import { Sandbox } from "@e2b/code-interpreter";
import {
GraphState,
GraphConfig,
GraphUpdate,
TargetRepository,
} from "../types.js";
import { TIMEOUT_EXTENSION_OPT } from "../constants.js";
const JS_SANDBOX_TEMPLATE_ID = "eh0860emqx28qyxmbctu";
async function cloneRepo(sandbox: Sandbox, targetRepository: TargetRepository) {
if (!process.env.GITHUB_PAT) {
throw new Error("GITHUB_PAT environment variable not set.");
}
const gitCloneCommand = ["git", "clone"];
const repoUrlWithToken = `https://${process.env.GITHUB_PAT}@github.com/${targetRepository.owner}/${targetRepository.repo}.git`;
if (targetRepository.branch) {
gitCloneCommand.push("-b", targetRepository.branch, repoUrlWithToken);
} else {
gitCloneCommand.push(repoUrlWithToken);
}
return await sandbox.commands.run(
gitCloneCommand.join(" "),
TIMEOUT_EXTENSION_OPT,
);
}
/**
* Initializes the session. This ensures there's an active VM session, and that
* the proper credentials are provided for taking actions on GitHub.
* It also clones the repository the user has specified to be used, and an optional
* branch.
*/
export function initialize(
state: GraphState,
export async function initialize(
_state: GraphState,
config: GraphConfig,
): Promise<GraphUpdate> {
throw new Error("Not implemented");
if (!config.configurable) {
throw new Error("Configuration object not found.");
}
const { sandbox_session_id, target_repository, sandbox_language } =
config.configurable;
if (sandbox_session_id) {
// Session already initialized. Return early.
return {};
}
if (!sandbox_language || !target_repository) {
throw new Error(
"Missing required configuration. Please provide a sandbox language and git repository URL.",
);
}
if (sandbox_language === "js") {
const sandbox = await Sandbox.create(JS_SANDBOX_TEMPLATE_ID);
config.configurable.sandbox_session_id = sandbox.sandboxId;
const res = await cloneRepo(sandbox, target_repository);
if (res.error) {
let errorMessage = "Failed to clone repository.";
if (res.error) {
// Basic redaction attempt for the token in the URL if present in the error
errorMessage +=
": " + res.error.replace(/:\/\/[^@]+@/, "://<REDACTED_TOKEN>@");
} else if (typeof res.error === "string") {
errorMessage +=
": " + res.error.replace(/:\/\/[^@]+@/, "://<REDACTED_TOKEN>@");
} else {
errorMessage += " An unknown error occurred.";
}
throw new Error(errorMessage);
}
return {};
}
if (sandbox_language === "python") {
throw new Error("Python sandbox not implemented yet.");
}
throw new Error("Unsupported sandbox language: " + sandbox_language);
}

View file

@ -1,12 +0,0 @@
import { GraphState, GraphConfig, GraphUpdate } from "../types.js";
/**
* Initializes the session. This ensures there's an active VM session, and that
* the proper credentials are provided for taking actions on GitHub.
*/
export function summarizeSession(
state: GraphState,
config: GraphConfig,
): Promise<GraphUpdate> {
throw new Error("Not implemented");
}

21
src/tools/diff.ts Normal file
View file

@ -0,0 +1,21 @@
import { tool } from "@langchain/core/tools";
import { z } from "zod";
import { applyPatch } from "diff";
const applyDiffToolSchema = z.object({
diff: z.string().describe("The diff to apply"),
file_path: z.string().describe("The file path to apply the diff to"),
});
export const applyDiffTool = tool(
(input) => {
const { diff, file_path } = input;
const res = applyPatch(file_path, diff);
return `Successfully applied diff to \`${file_path}\``;
},
{
name: "apply_diff",
description: "Applies a diff to a file.",
schema: applyDiffToolSchema,
},
);

43
src/tools/shell.ts Normal file
View file

@ -0,0 +1,43 @@
import { tool } from "@langchain/core/tools";
import { z } from "zod";
import { Sandbox } from "@e2b/code-interpreter";
import { GraphConfig } from "../types.js";
const shellToolSchema = z.object({
command: z.array(z.string()).describe("The command to run"),
workdir: z
.string()
.optional()
.describe("The working directory for the command."),
timeout: z
.number()
.optional()
.describe(
"The maximum time to wait for the command to complete in milliseconds.",
),
});
export const shellTool = tool(
async (input, config: GraphConfig) => {
const sessionId = config.configurable?.sandbox_session_id;
if (!sessionId) {
return "FAILED TO RUN COMMAND: No sandbox session ID provided";
}
const sandbox = await Sandbox.connect(sessionId);
const { command, workdir, timeout } = input;
const result = await sandbox.commands.run(command.join(" "), {
timeoutMs: timeout,
cwd: workdir,
});
return `Command exited with code ${result.exitCode}:\n
${result.error ? `Error: ${result.error}\n` : ""}
Stdout:\n${result.stdout}\n\nStderr:\n${result.stderr}`;
},
{
name: "shell",
description: "Runs a shell command, and returns its output.",
schema: shellToolSchema,
},
);

View file

@ -12,6 +12,12 @@ type PlanItem = {
completed: boolean;
};
export type TargetRepository = {
owner: string;
repo: string;
branch?: string;
};
export const GraphAnnotation = Annotation.Root({
messages: MessagesAnnotation.spec.messages,
plan: Annotation<PlanItem[]>({
@ -36,107 +42,30 @@ export const MCPConfig = z.object({
export const GraphConfiguration = z.object({
/**
* The model ID to use for the reflection generation.
* Should be in the format `provider/model_name`.
* Defaults to `anthropic/claude-3-7-sonnet-latest`.
* The session ID of the Sandbox to use.
*/
modelName: z
sandbox_session_id: z
.string()
.optional()
.langgraph.metadata({
x_lg_ui_config: {
type: "select",
default: "anthropic/claude-3-7-sonnet-latest",
description: "The model to use in all generations",
options: [
{
label: "Claude 3.7 Sonnet",
value: "anthropic/claude-3-7-sonnet-latest",
},
{
label: "Claude 3.5 Sonnet",
value: "anthropic/claude-3-5-sonnet-latest",
},
{
label: "GPT 4o",
value: "openai/gpt-4o",
},
{
label: "GPT 4.1",
value: "openai/gpt-4.1",
},
{
label: "o3",
value: "openai/o3",
},
{
label: "o3 mini",
value: "openai/o3-mini",
},
{
label: "o4",
value: "openai/o4",
},
],
type: "hidden",
},
}),
/**
* The temperature to use for the reflection generation.
* Defaults to `0.7`.
* The URL of the repository to clone.
*/
temperature: z
.number()
.optional()
.langgraph.metadata({
x_lg_ui_config: {
type: "slider",
default: 0.7,
min: 0,
max: 2,
step: 0.1,
description: "Controls randomness (0 = deterministic, 2 = creative)",
},
}),
target_repository: z
.object({
owner: z.string(),
repo: z.string(),
branch: z.string().optional(),
})
.langgraph.metadata({}),
/**
* The maximum number of tokens to generate.
* Defaults to `1000`.
* The language of the sandbox to use.
*/
maxTokens: z
.number()
.optional()
.langgraph.metadata({
x_lg_ui_config: {
type: "number",
default: 4000,
min: 1,
description: "The maximum number of tokens to generate",
},
}),
systemPrompt: z
.string()
.optional()
.langgraph.metadata({
x_lg_ui_config: {
type: "textarea",
placeholder: "Enter a system prompt...",
description: "The system prompt to use in all generations",
},
}),
/**
* MCP configuration for tool selection
*/
mcpConfig: z
.lazy(() => MCPConfig)
.optional()
.langgraph.metadata({
x_lg_ui_config: {
type: "mcp",
// Add custom tools here.
// default: {
// tools: ["Math_Divide", "Math_Mod"]
// }
},
}),
sandbox_language: z.enum(["js", "python"]).optional().langgraph.metadata({}),
});
export type GraphConfig = LangGraphRunnableConfig<

View file

@ -291,6 +291,11 @@
resolved "https://registry.yarnpkg.com/@bcoe/v8-coverage/-/v8-coverage-0.2.3.tgz#75a2e8b51cb758a7553d6804a5932d7aace75c39"
integrity sha512-0hYQ8SB4Db5zvZB4axdMHGwEaQjkZzFjQiN9LVYvIFB2nSUHW9tYpxWriPrWDASIxiaXax83REcLxuSdnGPZtw==
"@bufbuild/protobuf@^2.2.2":
version "2.4.0"
resolved "https://registry.yarnpkg.com/@bufbuild/protobuf/-/protobuf-2.4.0.tgz#ef6bfd2fc0374a95a4904cdc5fd8e02cc999cf73"
integrity sha512-RN9M76x7N11QRihKovEglEjjVCQEA9PRBVnDgk9xw8JHLrcUrp4FpAVSPSH91cNbcTft3u2vpLN4GMbiKY9PJw==
"@cfworker/json-schema@^4.0.2":
version "4.1.1"
resolved "https://registry.yarnpkg.com/@cfworker/json-schema/-/json-schema-4.1.1.tgz#4a2a3947ee9fa7b7c24be981422831b8674c3be6"
@ -306,6 +311,16 @@
resolved "https://registry.yarnpkg.com/@commander-js/extra-typings/-/extra-typings-13.1.0.tgz#026e29b04401c92fc4307223fbaadf1ff3e5551e"
integrity sha512-q5P52BYb1hwVWE6dtID7VvuJWrlfbCv4klj7BjUUOqMz4jbSZD4C9fJ9lRjL2jnBGTg+gDDlaXN51rkWcLk4fg==
"@connectrpc/connect-web@2.0.0-rc.3":
version "2.0.0-rc.3"
resolved "https://registry.yarnpkg.com/@connectrpc/connect-web/-/connect-web-2.0.0-rc.3.tgz#4dd9c18b8105cb7395077808c4967a9eafdffe76"
integrity sha512-w88P8Lsn5CCsA7MFRl2e6oLY4J/5toiNtJns/YJrlyQaWOy3RO8pDgkz+iIkG98RPMhj2thuBvsd3Cn4DKKCkw==
"@connectrpc/connect@2.0.0-rc.3":
version "2.0.0-rc.3"
resolved "https://registry.yarnpkg.com/@connectrpc/connect/-/connect-2.0.0-rc.3.tgz#f75cd1bd132dd983ff4f16432aaf5309ee673662"
integrity sha512-ARBt64yEyKbanyRETTjcjJuHr2YXorzQo0etyS5+P6oSeW8xEuzajA9g+zDnMcj1hlX2dQE93foIWQGfpru7gQ==
"@dabh/diagnostics@^2.0.2":
version "2.0.3"
resolved "https://registry.yarnpkg.com/@dabh/diagnostics/-/diagnostics-2.0.3.tgz#7f7e97ee9a725dffc7808d93668cc984e1dc477a"
@ -315,6 +330,13 @@
enabled "2.0.x"
kuler "^2.0.0"
"@e2b/code-interpreter@^1.2.0-beta.4":
version "1.2.0-beta.4"
resolved "https://registry.yarnpkg.com/@e2b/code-interpreter/-/code-interpreter-1.2.0-beta.4.tgz#03e5a829e1c90d660f5179356c9048ffed71294f"
integrity sha512-09Rciui9D/d7Fl7nzlNMYKFzloakauYNyXGMZxlD8UUY5oLvbuoT1zMDEKgKhfAbBkO8cRNPAZ1jRGLKYJlrUw==
dependencies:
e2b "1.2.0-beta.4"
"@emnapi/core@^1.4.0", "@emnapi/core@^1.4.3":
version "1.4.3"
resolved "https://registry.yarnpkg.com/@emnapi/core/-/core-1.4.3.tgz#9ac52d2d5aea958f67e52c40a065f51de59b77d6"
@ -1901,6 +1923,11 @@ commander@^13.0.0:
resolved "https://registry.yarnpkg.com/commander/-/commander-13.1.0.tgz#776167db68c78f38dcce1f9b8d7b8b9a488abf46"
integrity sha512-/rFeCpNJQbhSZjGVwO9RFV3xPqbnERS8MmIQzCtD/zl6gpJuV/bMLuN92oG3F7d8oDEHHRrujSXNUr8fpjntKw==
compare-versions@^6.1.0:
version "6.1.1"
resolved "https://registry.yarnpkg.com/compare-versions/-/compare-versions-6.1.1.tgz#7af3cc1099ba37d244b3145a9af5201b629148a9"
integrity sha512-4hm4VPpIecmlg59CHXnRDnqGplJFrbLG4aFEl5vl6cK1u76ws3LLvX7ikFnTDl5vo39sjWD6AaDPYodJp/NNHg==
concat-map@0.0.1:
version "0.0.1"
resolved "https://registry.yarnpkg.com/concat-map/-/concat-map-0.0.1.tgz#d8a96bd77fd68df7793a73036a3ba0d5405d477b"
@ -2111,6 +2138,11 @@ diff-sequences@^29.6.3:
resolved "https://registry.yarnpkg.com/diff-sequences/-/diff-sequences-29.6.3.tgz#4deaf894d11407c51efc8418012f9e70b84ea921"
integrity sha512-EjePK1srD3P08o2j4f0ExnylqRs5B9tJjcp9t1krH2qRi8CCdsYfwe9JgSLurFBWwq4uOlipzfk5fHNvwFKr8Q==
diff@^8.0.1:
version "8.0.1"
resolved "https://registry.yarnpkg.com/diff/-/diff-8.0.1.tgz#1562531dbca460fe5be98a026412996ccab1abd5"
integrity sha512-rEaM3KmVm78zE3dFZaop3aCQa2MTm+T4kcigUFLVU/KbOYdiY6JnL2g2puOYnct3QFw9pjZadaCbCZ1O8ArMlQ==
doctrine@^2.1.0:
version "2.1.0"
resolved "https://registry.yarnpkg.com/doctrine/-/doctrine-2.1.0.tgz#5cd01fc101621b42c4cd7f5d1a66243716d3f39d"
@ -2132,6 +2164,18 @@ dunder-proto@^1.0.0, dunder-proto@^1.0.1:
es-errors "^1.3.0"
gopd "^1.2.0"
e2b@1.2.0-beta.4:
version "1.2.0-beta.4"
resolved "https://registry.yarnpkg.com/e2b/-/e2b-1.2.0-beta.4.tgz#b2b2666724f96f99fddc489bf5cacc5662ec791d"
integrity sha512-LIXU1oWV9MnsSOtGPJh5LjZ17CQ/K252Ko/JJ+URn7EoFYShtRz+LZD4YsSWVOQUiI/+8EybHRsZtjhxw8zPXw==
dependencies:
"@bufbuild/protobuf" "^2.2.2"
"@connectrpc/connect" "2.0.0-rc.3"
"@connectrpc/connect-web" "2.0.0-rc.3"
compare-versions "^6.1.0"
openapi-fetch "^0.9.7"
platform "^1.3.6"
ee-first@1.1.1:
version "1.1.1"
resolved "https://registry.yarnpkg.com/ee-first/-/ee-first-1.1.1.tgz#590c61156b0ae2f4f0255732a158b266bc56b21d"
@ -4394,11 +4438,23 @@ openai@^4.96.0:
formdata-node "^4.3.2"
node-fetch "^2.6.7"
openapi-fetch@^0.9.7:
version "0.9.8"
resolved "https://registry.yarnpkg.com/openapi-fetch/-/openapi-fetch-0.9.8.tgz#177cd9d1c9a4f42ac6aec7adf075ac9638e1fbf5"
integrity sha512-zM6elH0EZStD/gSiNlcPrzXcVQ/pZo3BDvC6CDwRDUt1dDzxlshpmQnpD6cZaJ39THaSmwVCxxRrPKNM1hHrDg==
dependencies:
openapi-typescript-helpers "^0.0.8"
openapi-types@^12.1.3:
version "12.1.3"
resolved "https://registry.yarnpkg.com/openapi-types/-/openapi-types-12.1.3.tgz#471995eb26c4b97b7bd356aacf7b91b73e777dd3"
integrity sha512-N4YtSYJqghVu4iek2ZUvcN/0aqH1kRDuNqzcycDxhOUpg7GdvLa2F3DgS6yBNhInhv2r/6I0Flkn7CqL8+nIcw==
openapi-typescript-helpers@^0.0.8:
version "0.0.8"
resolved "https://registry.yarnpkg.com/openapi-typescript-helpers/-/openapi-typescript-helpers-0.0.8.tgz#460f395362cc16e4a5de56264b7b1c5a03746e35"
integrity sha512-1eNjQtbfNi5Z/kFhagDIaIRj6qqDzhjNJKz8cmMW0CVdGwT6e1GLbAfgI0d28VTJa1A8jz82jm/4dG8qNoNS8g==
optionator@^0.9.3:
version "0.9.4"
resolved "https://registry.yarnpkg.com/optionator/-/optionator-0.9.4.tgz#7ea1c1a5d91d764fb282139c88fe11e182a3a734"
@ -4570,6 +4626,11 @@ pkg-dir@^4.2.0:
dependencies:
find-up "^4.0.0"
platform@^1.3.6:
version "1.3.6"
resolved "https://registry.yarnpkg.com/platform/-/platform-1.3.6.tgz#48b4ce983164b209c2d45a107adb31f473a6e7a7"
integrity sha512-fnWVljUchTro6RiCFvCXBbNhJc2NijN7oIQxbwsyL0buWJPG85v81ehlHI9fXrJsMNgTofEoWIQeClKpgxFLrg==
possible-typed-array-names@^1.0.0:
version "1.1.0"
resolved "https://registry.yarnpkg.com/possible-typed-array-names/-/possible-typed-array-names-1.1.0.tgz#93e3582bc0e5426586d9d07b79ee40fc841de4ae"