Lee-W commented on code in PR #74404: URL: https://github.com/apache/airflow/pull/74404#discussion_r4262327505
########## providers/common/ai/src/airflow/providers/common/ai/durable/strands_storage.py: ########## @@ -0,0 +1,137 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +A `Strands Agents <https://strandsagents.com/>`__ ``Storage`` backed by the task state store. + +Available on Airflow >= 3.3, where the task state store keeps key/value state for a +task instance across its tries within a Dag run; ``NEVER_EXPIRE`` does not exist on older +Airflow versions. +""" + +from __future__ import annotations + +import asyncio +import base64 +import bisect +import builtins +import threading +from typing import TYPE_CHECKING + +from airflow.providers.common.compat.sdk import AirflowOptionalProviderFeatureException + +try: + from strands.storage import Storage +except ImportError as e: + raise AirflowOptionalProviderFeatureException(e) + +from airflow.providers.common.ai.durable.base import DURABLE_KEY_PREFIX +from airflow.sdk.execution_time.context import NEVER_EXPIRE + +if TYPE_CHECKING: + from pydantic import JsonValue + + from airflow.sdk.execution_time.context import TaskStateStoreAccessor + +__all__ = ["TaskStateStoreStorage"] + +_KEY_PREFIX = f"{DURABLE_KEY_PREFIX}strands/" +# The task state store cannot list keys by prefix, so the keys this storage holds are kept +# in a sorted list under one more key. +_INDEX_KEY = f"{DURABLE_KEY_PREFIX}strands_index" +# Every instance over the same task instance shares that list, so updates to it are +# serialized across the process, not per instance. +_index_lock = threading.Lock() + + +class TaskStateStoreStorage(Storage): + """ + A Strands ``Storage`` that keeps its keys in the running task instance's task state store. + + The keys share the task instance's key namespace with anything user code writes through + ``context["task_state_store"]``, under a reserved prefix that keeps them apart. They are + written with ``NEVER_EXPIRE``, so a retry finds them whatever ``retry_delay`` or the + retention config is. Keys left behind by a task that never succeeds are removed with + the Dag run. + + A value that is valid UTF-8, such as a Strands session snapshot, is stored as text, and + anything else as base64. + + :param accessor: The task state store of the running task instance, + ``context["task_state_store"]``. + """ + + def __init__(self, accessor: TaskStateStoreAccessor) -> None: + self._store = accessor + + async def write(self, key: str, data: bytes) -> None: + await asyncio.to_thread(self._write, key, data) + + async def read(self, key: str) -> bytes | None: + return _decode(await asyncio.to_thread(self._store.get, _KEY_PREFIX + key)) + + async def delete(self, key: str) -> None: + await asyncio.to_thread(self._delete, key) + + async def list(self, query: str = "") -> builtins.list[str]: + keys = await asyncio.to_thread(self._load_index) + return [key for key in keys if key.startswith(query)] + + def _write(self, key: str, data: bytes) -> None: Review Comment: We might want handle a failing write as well so that a backend error or a non-JSON-serializable object in `agent.state` won't fail the task. `durable/task_state_store.py` on the pydantic-ai side deliberately swallows failed writes (warning, then skip), and only the final `delete_session` here is best-effort. -- 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]
