uranusjr commented on code in PR #76015:
URL: https://github.com/apache/airflow/pull/76015#discussion_r4450662502
##########
providers/common/ai/src/airflow/providers/common/ai/durable/openai_agents.py:
##########
@@ -345,24 +376,50 @@ async def get_response(
)
if step.replayed:
try:
- return ModelResponse(
+ replayed = ModelResponse(
output=_OUTPUT_ITEMS.validate_python(step.payload["output"]),
usage=Usage(),
response_id=step.payload.get("response_id"),
)
except (ValidationError, KeyError, TypeError) as error:
self.durable_run.reject(step, reason=f"the recorded model
response cannot be read: {error}")
+ else:
+ self._reserve_tool_steps(replayed, tools)
+ return replayed
async def call() -> ModelResponse:
return await self.wrapped.get_response(
system_instructions, input, model_settings, tools,
output_schema, handoffs, tracing, **kwargs
)
try:
- return await step.run(call, to_record=_record_response)
+ response = await step.run(call, to_record=_record_response)
except Exception:
self.failed_step = step
raise
+ self._reserve_tool_steps(response, tools)
+ return response
+
+ def _reserve_tool_steps(self, response: ModelResponse, tools: list[Tool])
-> None:
+ """
+ Claim a position for each of the response's tool calls, in the order
the model made them.
+
+ A call whose tool never runs (held for approval, stopped by a
guardrail) leaves its position
+ empty, which a retry runs live without disturbing the positions after
it.
+ """
+ durable_names = {tool.name for tool in tools if getattr(tool,
_DURABLE_MARKER, False)}
+ reserved = _RESERVED_TOOL_STEPS.setdefault(self.durable_run, {})
+ for item in response.output:
+ if (
+ isinstance(item, ResponseFunctionToolCall)
+ and item.name in durable_names
+ and item.call_id not in reserved
Review Comment:
A reservation that is never taken lingers for the rest of the run. For
example, a tool input guardrail rejects the call, so the tool never reaches its
wrapper. If a later response reuses the call id, the later call gets no
reservation. Its tool then finds the stale step.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
