kaxil opened a new pull request, #64199: URL: https://github.com/apache/airflow/pull/64199
Adds `durable=True` parameter to `AgentOperator` and `@task.agent` that enables step-level caching of model responses and tool results via ObjectStorage. On task retry, cached steps are replayed instead of re-executing expensive LLM calls and tool operations. ## Design - **CachingModel(WrapperModel)**: Intercepts `model.request()` calls. Returns cached `ModelResponse` on retry, otherwise calls the real model and caches. - **CachingToolset(WrapperToolset)**: Intercepts `tool.call_tool()` calls with the same pattern. - **DurableStepCounter**: Shared monotonic counter between CachingModel and CachingToolset for deterministic step indexing. - **DurableStorage**: Persists all step caches in a single JSON file on ObjectStorage. Base path configured via `[common.ai] durable_cache_path`. ## Why `ObjectStorage` instead of `XCom`? Airflow 3.x clears XCom for a task on retry. By storing caches on `ObjectStorage` (file://, s3://, gs://, etc.), they survive retry clearing. The cache file is deleted on successful task completion. Once [AIP-103 Task State Management](https://cwiki.apache.org/confluence/x/oZM8G) is out, we will move to it. ## Why step-level caching instead of message history checkpointing? Step-level caching provides finer granularity and lower storage overhead than checkpointing the full message history. Each model call and tool call is individually cached, so a retry only re-executes from the exact step that failed. ## Usage ```python AgentOperator( task_id="analyst", prompt="Analyze the data", llm_conn_id="my_llm", durable=True, retries=3, toolsets=[SQLToolset(db_conn_id="postgres_default")], ) ``` Requires `[common.ai] durable_cache_path` config (e.g. `file:///tmp/airflow_durable_cache`). Bumps pydantic-ai-slim to >=1.34.0 for WrapperModel/WrapperToolset base classes. -- 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]
