amoghrajesh commented on code in PR #76790:
URL: https://github.com/apache/airflow/pull/76790#discussion_r4525253807
##########
providers/common/ai/src/airflow/providers/common/ai/tools/_from_toolset.py:
##########
@@ -206,16 +279,36 @@ def _as_airflow_tool(
) -> AirflowTool:
budget = _RetryBudget(tool.max_retries)
- def correctable(started: float, turn: Hashable | None, message: str) ->
ToolResult:
- if budget.failed(started, turn):
+ def correctable(started: float, scope: ToolCallScope, message: str) ->
ToolResult:
+ if budget.failed(started, scope.turn):
+ scope.limit_reached = True
raise ToolCallError(f"{name} kept failing after
{budget.max_retries} correction(s): {message}")
+ scope.counted = budget.take()
return ToolResult(content=message, is_error=True)
- async def call(arguments: dict[str, Any], started: float, turn: Hashable |
None) -> ToolResult:
+ def replay(count: RetryCount) -> ToolResult:
+ # The result the earlier attempt recorded reaches the model from the
journal, not from here.
+ restored = budget.failures
+ budget.restore(count)
+ if budget.failures != restored:
+ log.info(
Review Comment:
The log.info "an earlier attempt left N of M correction(s) used" line is
documented in `openai_agents.rst`, but no test covers it. Maybe add a
structured caplog check to one replay test.
##########
providers/common/ai/src/airflow/providers/common/ai/tools/openai_agents.py:
##########
@@ -163,25 +200,29 @@ async def call_airflow_tool(ctx: ToolContext[Any],
arguments: str) -> str:
return serialize_for_llm({"error": f"The arguments are not valid
JSON: {e}"})
if not isinstance(parsed, dict):
return serialize_for_llm({"error": "The arguments must be a JSON
object."})
- # The SDK does not identify the model turn, so calls that run at the
same time count once.
- try:
- with calling_framework("openai_agents"),
tool_call_scope(run=_SameRun(ctx.usage)):
- result = await tool.call(parsed)
- except ToolCallError as e:
- # Already masked, and logged to the task log with its traceback by
AirflowTool.call.
- raise _FailedToolCall(str(e)) from None
+ with _calling_tool(ctx, tool) as scope:
+ result = await tool.call(parsed)
+ setattr(ctx, _RETRY_COUNT, scope.counted)
if result.is_error:
return serialize_for_llm({"error": result.content})
return serialize_for_llm(result.content)
- return FunctionTool(
+ async def replay_retry(ctx: ToolContext[Any], count: RetryCount) -> None:
+ # Not through AirflowTool.call: nothing reaches the model, and its
task log lines would
+ # show a call that never ran. Only a toolset tool records a count, and
it does not run.
+ with _calling_tool(ctx, tool, replayed=count):
+ await tool.function({})
Review Comment:
`replay_retry` restores the count by calling `tool.function({})` and relies
on a `ToolCallScope.replayed` contextvar to skip the real call. If a retry tool
with the same name is a custom wrapper (user changed the Dag between tries,
same schema, so the fingerprint matches), the wrapper runs with `{}`. A cleaner
fix is to have `_as_airflow_tool` expose its `replay(count)` directly as a
hook, so replay never goes through the tool function at all.
--
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]
