diff --git a/backend/src/ea_chatbot/graph/workers/__init__.py b/backend/src/ea_chatbot/graph/workers/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/backend/src/ea_chatbot/graph/workers/data_analyst/__init__.py b/backend/src/ea_chatbot/graph/workers/data_analyst/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/backend/src/ea_chatbot/graph/workers/data_analyst/state.py b/backend/src/ea_chatbot/graph/workers/data_analyst/state.py new file mode 100644 index 0000000..1c7995a --- /dev/null +++ b/backend/src/ea_chatbot/graph/workers/data_analyst/state.py @@ -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] diff --git a/backend/tests/test_data_analyst_state.py b/backend/tests/test_data_analyst_state.py new file mode 100644 index 0000000..793dcee --- /dev/null +++ b/backend/tests/test_data_analyst_state.py @@ -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