10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185 | class Engine:
"""
The Protokol Execution Engine.
This class is completely synchronous and stateless. It evaluates the current RunPlan,
yields Contexts to the user, and advances the state machine based on user results.
"""
def __init__(
self,
on_step_start: Optional[Callable[[str, RunPlan], None]] = None,
on_step_complete: Optional[Callable[[str, RunPlan, Any], None]] = None,
on_flow_error: Optional[Callable[[str, RunPlan, Exception], None]] = None,
retry_strategy: Optional[RetryStrategy] = None,
):
# Observability hooks for integrating Datadog, Prometheus, etc.
self.on_step_start = on_step_start
self.on_step_complete = on_step_complete
self.on_flow_error = on_flow_error
self.retry_strategy = retry_strategy or SimpleRetryStrategy()
def _resolve_flow(self, root_flow: AbstractFlow, plan: RunPlan) -> AbstractFlow:
"""Resolve the currently active flow by walking down the nested flow call stack."""
current_flow = root_flow
for frame in plan.call_stack:
parent_id = frame["parent_step"]
current_flow = current_flow.get_step(parent_id)
return current_flow
def get_context(self, flow: AbstractFlow, plan: RunPlan) -> Union[StepContext, Dict[str, StepContext]]:
"""Yields the StepContext(s) for the current step in the RunPlan."""
if not plan.current_step:
raise ValueError("RunPlan.current_step is not set.")
active_flow = self._resolve_flow(flow, plan)
# Handle Parallel Execution Contexts
if isinstance(plan.current_step, list):
contexts = {}
for s_id in plan.current_step:
if self.on_step_start:
self.on_step_start(s_id, plan)
contexts[s_id] = active_flow.get_step(s_id).get_context(plan)
return contexts
if self.on_step_start:
self.on_step_start(plan.current_step, plan)
step = active_flow.get_step(plan.current_step)
# Handle Nested Flow: Push the parent step to the stack and descend into the sub-flow
if isinstance(step, AbstractFlow):
plan.call_stack.append({"parent_step": plan.current_step})
plan.current_step = step.start_step_id
return self.get_context(flow, plan)
return step.get_context(plan)
def advance(self, flow: AbstractFlow, plan: RunPlan, user_result: Any):
"""Processes the LLM result returned by the user and advances the state machine pointers."""
try:
active_flow = self._resolve_flow(flow, plan)
# --- Parallel Step Advancement ---
if isinstance(plan.current_step, list):
if not isinstance(user_result, dict):
raise ValueError("For parallel execution, user_result must be a dict.")
next_steps: List[str] = []
for s_id in plan.current_step:
step = active_flow.get_step(s_id)
step_state = plan.get_step_state(s_id)
try:
output = step.process(plan, user_result.get(s_id))
except Exception as e:
# Built-in Retry Logic: If processing fails, increment attempt and skip advancing.
step_state.attempt += 1
step_state.last_error = str(e)
if self.retry_strategy.should_retry(step_state.attempt, e):
logger.warning(f"Step {s_id} failed, retrying: {e}")
next_steps.append(s_id) # Keep step in the next parallel batch
continue
else:
raise e
attempt_before_reset = step_state.attempt
if self.on_step_complete:
self.on_step_complete(s_id, plan, output)
# Calculate next step(s)
ns = step.next(plan, output)
# Write to immutable audit log
plan.add_trace_entry(
step=s_id,
attempt=attempt_before_reset,
result=user_result.get(s_id),
next_step=ns,
parallel=True,
)
plan.reset_step_state(s_id)
if ns:
next_steps.extend(ns if isinstance(ns, list) else [ns])
self._handle_next_steps(flow, plan, next_steps)
# --- Standard Step Advancement ---
else:
current_step_id = plan.current_step
step = active_flow.get_step(current_step_id)
step_state = plan.get_step_state(current_step_id)
try:
output = step.process(plan, user_result)
except Exception as e:
# Built-in Retry Logic
step_state.attempt += 1
step_state.last_error = str(e)
if self.retry_strategy.should_retry(step_state.attempt, e):
logger.warning(f"Step {current_step_id} failed, retrying: {e}")
return # Abort advancement, exact same context will be yielded on next loop
else:
raise e
attempt_before_reset = step_state.attempt
if self.on_step_complete:
self.on_step_complete(current_step_id, plan, output)
ns = step.next(plan, output)
# Write to immutable audit log
plan.add_trace_entry(
step=current_step_id,
attempt=attempt_before_reset,
result=user_result,
next_step=ns,
parallel=False,
)
plan.reset_step_state(current_step_id)
self._handle_next_steps(flow, plan, ns)
except Exception as e:
if self.on_flow_error:
curr = str(plan.current_step)
self.on_flow_error(curr, plan, e)
plan.failed = True
plan.error = str(e)
plan.is_terminal = True
def _handle_next_steps(self, root_flow: AbstractFlow, plan: RunPlan, ns: Union[str, List[str], None]):
"""Helper to mutate the plan's current_step pointer and handle nested flow stack pops."""
if not ns:
# The current flow has ended. Check if we are inside a nested sub-flow.
if plan.call_stack:
# Pop the stack and return control to the parent flow
parent_info = plan.call_stack.pop()
parent_step_id = parent_info["parent_step"]
# Re-resolve the active flow to the parent's level
parent_flow = self._resolve_flow(root_flow, plan)
parent_step = parent_flow.get_step(parent_step_id)
# Evaluate the parent step's next() method to continue routing
ns_after_pop = parent_step.next(plan, {})
self._handle_next_steps(root_flow, plan, ns_after_pop)
else:
# Stack is empty, the root flow is complete
plan.is_terminal = True
elif isinstance(ns, list):
# Parallel execution - preserve order and deduplicate while keeping deterministic ordering
plan.current_step = _unique_ordered(ns)
else:
plan.current_step = ns
|