feat(workers): Define WorkerState for the Data Analyst subgraph

This commit is contained in:
Yunxiao Xu
2026-02-23 04:55:50 -08:00
parent eeb2be409b
commit 5324cbe851
4 changed files with 45 additions and 0 deletions

View File

@@ -0,0 +1,28 @@
from typing import TypedDict, List, Dict, Any, Optional
from langchain_core.messages import BaseMessage
class WorkerState(TypedDict):
"""Internal state for the Data Analyst worker subgraph."""
# Internal worker conversation (not bubbled up to global unless summary)
messages: List[BaseMessage]
# The specific sub-task assigned by the Orchestrator
task: str
# Generated code and execution context
code: Optional[str]
output: Optional[str]
error: Optional[str]
# Number of internal retry attempts
iterations: int
# Temporary storage for analysis results
plots: List[Any]
# Isolated/Snapshot view of the VFS
vfs_state: Dict[str, Any]
# The final summary or result to return to the Orchestrator
result: Optional[str]

View File

@@ -0,0 +1,17 @@
from typing import get_type_hints, List
from langchain_core.messages import BaseMessage
from ea_chatbot.graph.workers.data_analyst.state import WorkerState
def test_data_analyst_worker_state():
"""Verify that DataAnalyst WorkerState has the required fields."""
hints = get_type_hints(WorkerState)
assert "messages" in hints
assert "task" in hints
assert hints["task"] == str
assert "code" in hints
assert "output" in hints
assert "error" in hints
assert "iterations" in hints
assert hints["iterations"] == int