Lee-W commented on code in PR #76790:
URL: https://github.com/apache/airflow/pull/76790#discussion_r4534473184
##########
providers/common/ai/src/airflow/providers/common/ai/tools/_from_toolset.py:
##########
@@ -39,31 +41,80 @@
from airflow.providers.common.ai.utils.coroutines import run_coroutine_sync
if TYPE_CHECKING:
+ from collections.abc import Awaitable, Callable
+
from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool
+log = logging.getLogger(__name__)
+
+
+class RetryCount(TypedDict):
+ """The corrections a tool had used when a call finished, and when it
finished."""
+
+ failures: int
+ # Wall-clock time, not monotonic: a retry compares it with times recorded
by earlier tries,
+ # which may have run on other workers.
+ finished: float
+
-# Set by a framework adapter around each call: which agent run and which model
turn the call
-# belongs to, where the framework says. The retry budget counts by them.
-_call_scope: ContextVar[tuple[Hashable | None, Hashable | None]] = ContextVar(
- "common_ai_tool_call_scope", default=(None, None)
-)
+@dataclass(slots=True)
+class ToolCallScope:
+ """Which agent run and model turn a tool call belongs to, and what it did
to the tool's retry budget."""
+
+ run: Hashable | None = None
+ turn: Hashable | None = None
+ # The tool function the adapter calls. Only its own count is recorded, not
that of a toolset
+ # tool a custom function calls in turn.
+ function: Callable[[dict[str, Any]], Awaitable[ToolResult]] | None = None
+ # Set by a call that failed or succeeded; None when it left the count
alone, such as a ToolFailed.
+ counted: RetryCount | None = None
+ # Set when the call failed once more than the budget allows, which is the
model's doing.
+ limit_reached: bool = False
+
+
+# Set by a framework adapter around each call. The retry budget counts by its
run and turn.
+_call_scope: ContextVar[ToolCallScope | None] =
ContextVar("common_ai_tool_call_scope", default=None)
@contextmanager
-def tool_call_scope(*, run: Hashable | None, turn: Hashable | None = None) ->
Iterator[None]:
+def tool_call_scope(
+ *,
+ run: Hashable | None,
+ turn: Hashable | None = None,
+ function: Callable[[dict[str, Any]], Awaitable[ToolResult]] | None = None,
+) -> Iterator[ToolCallScope]:
"""
Tell the tools which agent run and which model turn the calls inside the
block belong to.
Each tool's retry budget starts over when the run changes, and counts at
most one failure
per turn, as pydantic-ai does. Without a turn, calls that run at the same
time count as one.
+
+ Given the ``function`` of the tool the block calls, the block yields the
count that call left
+ in the tool's retry budget, so durable execution can record it with the
call's result and
+ restore it through ``retry_count_restorer`` when a retry replays that
result.
"""
- token = _call_scope.set((run, turn))
+ scope = ToolCallScope(run=run, turn=turn, function=function)
+ token = _call_scope.set(scope)
try:
- yield
+ yield scope
finally:
_call_scope.reset(token)
+# Set on the function of each tool this module builds: restores a count an
earlier attempt recorded.
+_RESTORE_RETRY_COUNT = "_airflow_restore_retry_count"
+
+
+def retry_count_restorer(tool: AirflowTool) -> Callable[[Hashable,
RetryCount], Awaitable[None]] | None:
Review Comment:
```suggestion
def get_retry_count_restorer(tool: AirflowTool) -> Callable[[Hashable,
RetryCount], Awaitable[None]] | None:
```
--
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]
