[BREAKING] Python: Refactor SharedState to State with sync methods and superstep caching (#3667)

* Refactor SharedState to State with sync methods and superstep caching

* Fixes

* Address PR feedback

* Remove dead links

* Fix lab test import
This commit is contained in:
Evan Mattson
2026-02-05 10:42:52 +09:00
committed by GitHub
Unverified
parent 4e25917644
commit 10afb86213
48 changed files with 1971 additions and 1724 deletions
@@ -36,7 +36,7 @@ Purpose:
Demonstrate how to use a multi-selection edge group to fan out from one executor to multiple possible targets.
Show how to:
- Implement a selection function that chooses one or more downstream branches based on analysis.
- Share state across branches so different executors can read the same email content.
- Share workflow state across branches so different executors can read the same email content.
- Validate agent outputs with Pydantic models for robust structured data exchange.
- Merge results from multiple branches (e.g., a summary) back into a typed state.
- Apply conditional persistence logic (short vs long emails).
@@ -44,7 +44,7 @@ Show how to:
Prerequisites:
- Familiarity with WorkflowBuilder, executors, edges, and events.
- Understanding of multi-selection edge groups and how their selection function maps to target ids.
- Experience with shared state in workflows for persisting and reusing objects.
- Experience with workflow state for persisting and reusing objects.
"""
@@ -87,8 +87,8 @@ class DatabaseEvent(WorkflowEvent): ...
@executor(id="store_email")
async def store_email(email_text: str, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
new_email = Email(email_id=str(uuid4()), email_content=email_text)
await ctx.set_shared_state(f"{EMAIL_STATE_PREFIX}{new_email.email_id}", new_email)
await ctx.set_shared_state(CURRENT_EMAIL_ID_KEY, new_email.email_id)
ctx.set_state(f"{EMAIL_STATE_PREFIX}{new_email.email_id}", new_email)
ctx.set_state(CURRENT_EMAIL_ID_KEY, new_email.email_id)
await ctx.send_message(
AgentExecutorRequest(messages=[ChatMessage("user", text=new_email.email_content)], should_respond=True)
@@ -98,8 +98,8 @@ async def store_email(email_text: str, ctx: WorkflowContext[AgentExecutorRequest
@executor(id="to_analysis_result")
async def to_analysis_result(response: AgentExecutorResponse, ctx: WorkflowContext[AnalysisResult]) -> None:
parsed = AnalysisResultAgent.model_validate_json(response.agent_response.text)
email_id: str = await ctx.get_shared_state(CURRENT_EMAIL_ID_KEY)
email: Email = await ctx.get_shared_state(f"{EMAIL_STATE_PREFIX}{email_id}")
email_id: str = ctx.get_state(CURRENT_EMAIL_ID_KEY)
email: Email = ctx.get_state(f"{EMAIL_STATE_PREFIX}{email_id}")
await ctx.send_message(
AnalysisResult(
spam_decision=parsed.spam_decision,
@@ -116,7 +116,7 @@ async def submit_to_email_assistant(analysis: AnalysisResult, ctx: WorkflowConte
if analysis.spam_decision != "NotSpam":
raise RuntimeError("This executor should only handle NotSpam messages.")
email: Email = await ctx.get_shared_state(f"{EMAIL_STATE_PREFIX}{analysis.email_id}")
email: Email = ctx.get_state(f"{EMAIL_STATE_PREFIX}{analysis.email_id}")
await ctx.send_message(
AgentExecutorRequest(messages=[ChatMessage("user", text=email.email_content)], should_respond=True)
)
@@ -131,7 +131,7 @@ async def finalize_and_send(response: AgentExecutorResponse, ctx: WorkflowContex
@executor(id="summarize_email")
async def summarize_email(analysis: AnalysisResult, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
# Only called for long NotSpam emails by selection_func
email: Email = await ctx.get_shared_state(f"{EMAIL_STATE_PREFIX}{analysis.email_id}")
email: Email = ctx.get_state(f"{EMAIL_STATE_PREFIX}{analysis.email_id}")
await ctx.send_message(
AgentExecutorRequest(messages=[ChatMessage("user", text=email.email_content)], should_respond=True)
)
@@ -140,8 +140,8 @@ async def summarize_email(analysis: AnalysisResult, ctx: WorkflowContext[AgentEx
@executor(id="merge_summary")
async def merge_summary(response: AgentExecutorResponse, ctx: WorkflowContext[AnalysisResult]) -> None:
summary = EmailSummaryModel.model_validate_json(response.agent_response.text)
email_id: str = await ctx.get_shared_state(CURRENT_EMAIL_ID_KEY)
email: Email = await ctx.get_shared_state(f"{EMAIL_STATE_PREFIX}{email_id}")
email_id: str = ctx.get_state(CURRENT_EMAIL_ID_KEY)
email: Email = ctx.get_state(f"{EMAIL_STATE_PREFIX}{email_id}")
# Build an AnalysisResult mirroring to_analysis_result but with summary
await ctx.send_message(
AnalysisResult(
@@ -165,7 +165,7 @@ async def handle_spam(analysis: AnalysisResult, ctx: WorkflowContext[Never, str]
@executor(id="handle_uncertain")
async def handle_uncertain(analysis: AnalysisResult, ctx: WorkflowContext[Never, str]) -> None:
if analysis.spam_decision == "Uncertain":
email: Email | None = await ctx.get_shared_state(f"{EMAIL_STATE_PREFIX}{analysis.email_id}")
email: Email | None = ctx.get_state(f"{EMAIL_STATE_PREFIX}{analysis.email_id}")
await ctx.yield_output(
f"Email marked as uncertain: {analysis.reason}. Email content: {getattr(email, 'email_content', '')}"
)
@@ -25,13 +25,13 @@ from typing_extensions import Never
"""
Sample: Switch-Case Edge Group with an explicit Uncertain branch.
The workflow stores a single email in shared state, asks a spam detection agent for a three way decision,
The workflow stores a single email in workflow state, asks a spam detection agent for a three way decision,
then routes with a switch-case group: NotSpam to the drafting assistant, Spam to a spam handler, and
Default to an Uncertain handler.
Purpose:
Demonstrate deterministic one of N routing with switch-case edges. Show how to:
- Persist input once in shared state, then pass around a small typed pointer that carries the email id.
- Persist input once in workflow state, then pass around a small typed pointer that carries the email id.
- Validate agent JSON with Pydantic models for robust parsing.
- Keep executor responsibilities narrow. Transform model output to a typed DetectionResult, then route based
on that type.
@@ -74,7 +74,7 @@ class DetectionResult:
@dataclass
class Email:
# In memory record of the email content stored in shared state.
# In memory record of the email content stored in workflow state.
email_id: str
email_content: str
@@ -93,8 +93,8 @@ def get_case(expected_decision: str):
async def store_email(email_text: str, ctx: WorkflowContext[AgentExecutorRequest]) -> None:
# Persist the raw email once. Store under a unique key and set the current pointer for convenience.
new_email = Email(email_id=str(uuid4()), email_content=email_text)
await ctx.set_shared_state(f"{EMAIL_STATE_PREFIX}{new_email.email_id}", new_email)
await ctx.set_shared_state(CURRENT_EMAIL_ID_KEY, new_email.email_id)
ctx.set_state(f"{EMAIL_STATE_PREFIX}{new_email.email_id}", new_email)
ctx.set_state(CURRENT_EMAIL_ID_KEY, new_email.email_id)
# Kick off the detector by forwarding the email as a user message to the spam_detection_agent.
await ctx.send_message(
@@ -106,7 +106,7 @@ async def store_email(email_text: str, ctx: WorkflowContext[AgentExecutorRequest
async def to_detection_result(response: AgentExecutorResponse, ctx: WorkflowContext[DetectionResult]) -> None:
# Parse the detector JSON into a typed model. Attach the current email id for downstream lookups.
parsed = DetectionResultAgent.model_validate_json(response.agent_response.text)
email_id: str = await ctx.get_shared_state(CURRENT_EMAIL_ID_KEY)
email_id: str = ctx.get_state(CURRENT_EMAIL_ID_KEY)
await ctx.send_message(DetectionResult(spam_decision=parsed.spam_decision, reason=parsed.reason, email_id=email_id))
@@ -116,8 +116,8 @@ async def submit_to_email_assistant(detection: DetectionResult, ctx: WorkflowCon
if detection.spam_decision != "NotSpam":
raise RuntimeError("This executor should only handle NotSpam messages.")
# Load the original content from shared state using the id carried in DetectionResult.
email: Email = await ctx.get_shared_state(f"{EMAIL_STATE_PREFIX}{detection.email_id}")
# Load the original content from workflow state using the id carried in DetectionResult.
email: Email = ctx.get_state(f"{EMAIL_STATE_PREFIX}{detection.email_id}")
await ctx.send_message(
AgentExecutorRequest(messages=[ChatMessage("user", text=email.email_content)], should_respond=True)
)
@@ -143,7 +143,7 @@ async def handle_spam(detection: DetectionResult, ctx: WorkflowContext[Never, st
async def handle_uncertain(detection: DetectionResult, ctx: WorkflowContext[Never, str]) -> None:
# Uncertain path terminal. Surface the original content to aid human review.
if detection.spam_decision == "Uncertain":
email: Email | None = await ctx.get_shared_state(f"{EMAIL_STATE_PREFIX}{detection.email_id}")
email: Email | None = ctx.get_state(f"{EMAIL_STATE_PREFIX}{detection.email_id}")
await ctx.yield_output(
f"Email marked as uncertain: {detection.reason}. Email content: {getattr(email, 'email_content', '')}"
)