kaxil commented on code in PR #76790:
URL: https://github.com/apache/airflow/pull/76790#discussion_r4532050737
##########
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:
Good call, done in c02d75d708b. `_as_airflow_tool` now hangs a `restore(run,
count)` hook on the tool's function and the adapter calls it through
`retry_count_restorer`, so a replay never goes through the function. A tool
without the hook (your Dag-edit case) just replays the recorded result as text;
added a test for exactly that, which fails on the old code with `KeyError:
'key'`. It also let me drop `ToolCallScope.replayed` and the `replayed=` arg on
`tool_call_scope`.
##########
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:
Added in c02d75d708b:
`test_a_replayed_correctable_failure_counts_against_the_tools_retry_limit` now
asserts the exact message from that logger, so the line the guide quotes can't
drift.
--
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]
