feat(workers): Implement input/output mapping for Data Analyst subgraph
This commit is contained in:
39
backend/src/ea_chatbot/graph/workers/data_analyst/mapping.py
Normal file
39
backend/src/ea_chatbot/graph/workers/data_analyst/mapping.py
Normal file
@@ -0,0 +1,39 @@
|
||||
from typing import Dict, Any, List
|
||||
from langchain_core.messages import HumanMessage, AIMessage
|
||||
from ea_chatbot.graph.state import AgentState
|
||||
from ea_chatbot.graph.workers.data_analyst.state import WorkerState
|
||||
|
||||
def prepare_worker_input(state: AgentState) -> Dict[str, Any]:
|
||||
"""Prepare the initial state for the Data Analyst worker."""
|
||||
checklist = state.get("checklist", [])
|
||||
current_step = state.get("current_step", 0)
|
||||
|
||||
# Get the current task description
|
||||
task_desc = "Analyze data" # Default
|
||||
if 0 <= current_step < len(checklist):
|
||||
task_desc = checklist[current_step].get("task", task_desc)
|
||||
|
||||
return {
|
||||
"task": task_desc,
|
||||
"messages": [HumanMessage(content=task_desc)], # Start worker loop with the task
|
||||
"vfs_state": dict(state.get("vfs", {})),
|
||||
"iterations": 0,
|
||||
"plots": [],
|
||||
"code": None,
|
||||
"output": None,
|
||||
"error": None,
|
||||
"result": None
|
||||
}
|
||||
|
||||
def merge_worker_output(worker_state: WorkerState) -> Dict[str, Any]:
|
||||
"""Map worker results back to the global AgentState."""
|
||||
result = worker_state.get("result", "Analysis complete.")
|
||||
|
||||
# We bubble up the summary as an AI message
|
||||
updates = {
|
||||
"messages": [AIMessage(content=result)],
|
||||
"vfs": worker_state.get("vfs_state", {}),
|
||||
"plots": worker_state.get("plots", [])
|
||||
}
|
||||
|
||||
return updates
|
||||
Reference in New Issue
Block a user