YunxiaoXu/ea-chatbot-lg
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor(graph): Use piped Runnables for worker nodes to enable subgraph event streaming

Browse Source
This commit is contained in:
Yunxiao Xu
2026-02-23 17:48:23 -08:00
parent 02d93120e0
commit b8d8651924
3 changed files with 40 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
backend/src/ea_chatbot/graph/workflow.py
View File

@@ -16,17 +16,9 @@ from ea_chatbot.graph.nodes.summarize_conversation import summarize_conversation
_DATA_ANALYST_WORKER = create_data_analyst_worker()
_RESEARCHER_WORKER = create_researcher_worker()
def data_analyst_worker_node(state: AgentState) -> dict:
"""Wrapper node for the Data Analyst subgraph with state mapping."""
worker_input = prepare_worker_input(state)
worker_result = _DATA_ANALYST_WORKER.invoke(worker_input)
return merge_worker_output(worker_result)
def researcher_worker_node(state: AgentState) -> dict:
"""Wrapper node for the Researcher subgraph with state mapping."""
worker_input = prepare_researcher_input(state)
worker_result = _RESEARCHER_WORKER.invoke(worker_input)
return merge_researcher_output(worker_result)
# Define worker nodes as piped runnables to enable subgraph event propagation
data_analyst_worker_runnable = prepare_worker_input | _DATA_ANALYST_WORKER | merge_worker_output
researcher_worker_runnable = prepare_researcher_input | _RESEARCHER_WORKER | merge_researcher_output
def main_router(state: AgentState) -> str:
"""Route from query analyzer based on initial assessment."""
@@ -52,8 +44,8 @@ def create_workflow(
query_analyzer=query_analyzer_node,
planner=planner_node,
delegate=delegate_node,
data_analyst_worker=data_analyst_worker_node,
researcher_worker=researcher_worker_node,
data_analyst_worker=data_analyst_worker_runnable,
researcher_worker=researcher_worker_runnable,
reflector=reflector_node,
synthesizer=synthesizer_node,
clarification=clarification_node,