airflow.providers.common.ai.durable.replay_usage¶

Keeps durable replays out of the usage a run is counted and limited against.

Attributes¶

SUCCESSFUL_TOOL_RESULT_KINDS

Classes¶

ReplayUsageLedger

Nets durable replays out of the RunUsage a run is counted and limited against.

Functions¶

`subtract_request_usage`(run_usage, usage)	Subtract one response's (or one durable operation's) usage and cost from a run's usage, in place.
`fill_replayed_cost`(response)	Price a replayed response in place, the way pydantic-ai's graph is about to.
`is_successful_tool_payload`(payload)	Whether a recorded tool-call payload is a result pydantic-ai counts as a successful call.

Module Contents¶

airflow.providers.common.ai.durable.replay_usage.SUCCESSFUL_TOOL_RESULT_KINDS[source]¶

airflow.providers.common.ai.durable.replay_usage.subtract_request_usage(run_usage, usage)[source]¶

Subtract one response’s (or one durable operation’s) usage and cost from a run’s usage, in place.

airflow.providers.common.ai.durable.replay_usage.fill_replayed_cost(response)[source]¶

Price a replayed response in place, the way pydantic-ai’s graph is about to.

The journal records a live response before the graph prices it, so a replayed response comes back with usage.cost=None even for a priced model. The graph fills the cost right after the durable model request returns, and only when it is still unset, by the same lookup, so the cost this module subtracts is exactly the cost the graph then adds. Uses the public ModelResponse.cost(), which runs the same price lookup as the graph’s own best-effort pricing; a response it cannot price stays unpriced, as it would in the graph.

class airflow.providers.common.ai.durable.replay_usage.ReplayUsageLedger(*, run_usage, usage_limits)[source]¶

Nets durable replays out of the RunUsage a run is counted and limited against.

pydantic-ai counts a replayed step exactly like a live one: after a durable model request returns, the graph does requests += 1 and adds the response’s usage, and after a durable tool call returns successfully, the tool manager does tool_calls += 1. The ledger cancels each replay when it happens, so the run’s counts only move for live work. Nothing here is persisted: the result is the same whether the journal entry was written by the previous attempt, an attempt that failed halfway through its own replay, or an attempt before a clear (which resets the budget but keeps the durable journal).

Two of pydantic-ai’s limit checks run before the durable layer sees the step it guards: check_before_request before each model request, and the up-front check_before_tool_call projection over a whole step’s function-tool calls. For those, the ledger takes a credit ahead of time – when a replayed model response is followed in the journal by the next model request or by recorded tool results – and resolves each credit when the step runs: a replay keeps it, a live call gives it back and re-runs the check pydantic-ai made against the credited count. Credits that are never resolved (the run took a different path, or stopped) are given back by settle(), so they never reach the persisted total; until then, a credit for a stale journal entry the run never reaches has only loosened that one up-front check.

Parameters:

run_usage (pydantic_ai.usage.RunUsage) – The RunUsage passed to the run as usage=.
usage_limits (pydantic_ai.usage.UsageLimits | None) – The limits passed to the run; None means pydantic-ai’s defaults, as in Agent.run.

run_usage[source]¶

begin_model_request(messages)[source]¶

Note the start of a model request and return whether it continues a suspended response.

pydantic-ai re-issues a suspended response (Anthropic pause_turn, OpenAI background mode) by sending it back as the last message; the graph counts the whole chain as one request.

track_chain(response, model_request_parameters)[source]¶

Count the function-tool calls response issues, for crediting the steps after it.

credit_first_replay(durable_run)[source]¶

Credit the run’s first model request if the journal holds a response for it.

Called before the run starts, because pydantic-ai checks request_limit before the first request reaches the durable layer; without the credit, a retry whose seeded requests already equals request_limit could not start even when every step would replay for free. Steps recorded before the first model request (tool discovery, capability operations) are skipped over.

credit_successors(durable_run, position)[source]¶

Credit what the replayed, complete response at position is followed by in the journal.

The response’s function-tool calls took positions after it when they first ran, and the next model request a position after them. Tool calls with a recorded successful result are credited; the scan stops at the next model request, which is credited too. A position with nothing recorded is a call that raised or was not recorded, and the scan gives up once there are more of those than calls in the chain.

record_capability_replay(payload)[source]¶

Cancel the usage a replayed @durable_operation brings back with its result.

pydantic-ai adds the usage an operation accumulated (a summarization’s model call, say) to the run when the operation returns, including when it returns from the journal; that usage was already counted by the attempt that ran it.

credit_request()[source]¶

Credit the next model request, which the journal says will be a replay.

credit_tool_steps(steps, *, batch_calls)[source]¶

Credit the tool calls at these journal positions, which have recorded results.

Parameters:: batch_calls (int) – The number of function-tool calls this replayed response issued, i.e. the count pydantic-ai’s up-front check_before_tool_call projects forward before any of them run.

settle()[source]¶

Give back every unresolved credit.

Returns:: Whether a request credit was outstanding.
Return type:: bool

record_model_replay(response, *, continuation)[source]¶

Cancel what the graph is about to add for a replayed model response.

A continuation segment (Anthropic pause_turn, OpenAI background mode) is not a request of its own: the graph counts one request per chain and commits the merged response’s usage once. Merging sums the segments’ usage, except when a segment re-polls the same provider response id, where the merged usage is that segment’s cumulative snapshot – so the previous segment’s subtraction is given back before this one is taken.

Between this subtraction and the graph adding the response’s usage back, the graph awaits after_model_request; a non-ModelRetry exception raised there, or an AirflowTaskTimeout landing in that window, leaves this response’s tokens and cost under-counted.

record_live_model_request(*, had_request_credit)[source]¶

Re-run the pre-request check for a credited request that turned out to be live.

pydantic-ai checked this request against the credited count, which assumed it would replay for free.

record_tool_replay(step)[source]¶

Cancel the tool_calls += 1 the tool manager does after a replayed call returns.

record_live_tool_call(step)[source]¶

Give back a credit whose recorded result did not replay, and re-check the limit.

pydantic-ai’s up-front check projected the whole batch this credited call belongs to as if every credited call in it would replay for free. With this one turning out to be live, that assumption is wrong for one more call: re-run the check against pydantic-ai’s original projection plus one for each credited call in the batch that has turned out to be live so far (including this one). A credit that is still unresolved keeps counting as a replay – if it does replay, the original projection was already right for it; if it turns out live later, its own call to this method re-checks again. So the last credited call in a batch to turn live is checked against the sum of every live call in the batch, regardless of the order concurrent calls finish in, and the batch as a whole never exceeds the limit.

batch_calls (passed to credit_tool_steps()) is derived from the response’s tool-call parts matched against model_request_parameters’s function_tools, which also lists tools that need approval or are external – so in a non-resume batch, where those calls are deferred rather than executed, batch_calls can count more than pydantic-ai’s own projection does. That over-count flows into _tool_batch_projection and this method’s recheck, so it can only make the recheck stricter – a possible false block, never a missed one. pydantic-ai’s projection, in turn, also covers 'unknown'-kind calls that batch_calls does not; that under-count never reaches this method or record_tool_replay(), because nothing calls record_tool_call for calls of 'unknown' kind, so it can’t push the final count over the limit.

airflow.providers.common.ai.durable.replay_usage.is_successful_tool_payload(payload)[source]¶

Whether a recorded tool-call payload is a result pydantic-ai counts as a successful call.