feat(executor): Integrate VFS helper for in-memory artifact tracking

This commit is contained in:
Yunxiao Xu
2026-02-23 04:42:18 -08:00
parent 92d9288f38
commit eeb2be409b
2 changed files with 45 additions and 2 deletions

View File

@@ -8,6 +8,7 @@ from matplotlib.figure import Figure
from ea_chatbot.graph.state import AgentState
from ea_chatbot.utils.db_client import DBClient
from ea_chatbot.utils.vfs import VFSHelper
from ea_chatbot.utils.logging import get_logger
from ea_chatbot.config import Settings
@@ -37,12 +38,17 @@ def executor_node(state: AgentState) -> dict:
db_client = DBClient(settings=db_settings)
# Initialize the Virtual File System (VFS) helper
vfs_state = dict(state.get("vfs", {}))
vfs_helper = VFSHelper(vfs_state)
# Initialize local variables for execution
# 'db' is the DBClient instance, 'plots' is for matplotlib figures
local_vars = {
'db': db_client,
'plots': [],
'pd': pd
'pd': pd,
'vfs': vfs_helper
}
stdout_buffer = io.StringIO()
@@ -98,5 +104,6 @@ def executor_node(state: AgentState) -> dict:
"code_output": code_output,
"error": error,
"plots": plots,
"dfs": dfs
"dfs": dfs,
"vfs": vfs_state
}

View File

@@ -0,0 +1,36 @@
from ea_chatbot.graph.nodes.executor import executor_node
from ea_chatbot.graph.state import AgentState
def test_executor_with_vfs():
"""Verify that the executor node provides VFS access to the code."""
state = AgentState(
messages=[],
question="test",
analysis={},
next_action="test",
iterations=0,
checklist=[],
current_step=0,
vfs={},
plots=[],
dfs={}
)
# Code that uses the 'vfs' helper
code = """
vfs.write("output.txt", "Execution Result", metadata={"type": "text"})
print("VFS Write Complete")
"""
state["code"] = code
result = executor_node(state)
# Check if the execution was successful
assert result["error"] is None
assert "VFS Write Complete" in result["code_output"]
# Verify that the VFS state was updated
# Note: executor_node returns a dict of updates, which should include the updated 'vfs'
assert "vfs" in result
assert "output.txt" in result["vfs"]
assert result["vfs"]["output.txt"]["content"] == "Execution Result"