run-llama/workflows-py
1
0
Fork 0
mirror of https://github.com/run-llama/workflows-py.git synced 2026-08-24 20:01:34 -04:00
Code Issues 18 Packages Projects Releases 10 Wiki Activity

Key state storage by run id and namespace (#724)

Browse Source
This commit is contained in:
Adrian Lyjak
2026-07-07 18:48:04 -04:00
committed by GitHub
parent 126d40a66a
commit 26050d5803
30 changed files with 665 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.changeset/per-namespace-state-records.md
+7
View File
@@ -0,0 +1,7 @@
---
"llama-index-workflows": minor
"llama-agents-dbos": minor
"llama-agents-server": minor
---
State storage is now keyed by run id and namespace; existing rows migrate to the root namespace.
packages/llama-agents-dbos/src/llama_agents/dbos/runtime.py
+23 -11
View File
@@ -144,20 +144,24 @@ class DBOSWorkflowStore(AbstractWorkflowStore):
state_type: type[Any] | None = None,
serialized_state: dict[str, Any] | None = None,
serializer: BaseSerializer | None = None,
namespace: tuple[str, ...] = (),
) -> StateStore[Any]:
# Delegate the whole template method so memoization lives in the
# inner store's single cache (the proxy's own cache stays unused).
return self._resolve().create_state_store(
run_id, state_type, serialized_state, serializer
run_id, state_type, serialized_state, serializer, namespace
)
def _build_state_store(
self,
run_id: str,
namespace: tuple[str, ...],
state_type: type[Any] | None,
serializer: BaseSerializer | None,
) -> StateStoreFacade[Any]:
return self._resolve()._build_state_store(run_id, state_type, serializer)
return self._resolve()._build_state_store(
run_id, namespace, state_type, serializer
)
async def query(self, query: HandlerQuery) -> list[PersistentHandler]:
return await self._resolve().query(query)
@@ -1055,7 +1059,7 @@ class InternalDBOSAdapter(InternalRunAdapter):
self._db_path = db_path
self._closed = False
self._shutdown_event = asyncio.Event()
self._state_store: StateStore[Any] | None = None
self._state_stores: dict[tuple[str, ...], StateStore[Any]] = {}
# Journal for deterministic task ordering - lazily initialized
self._journal: TaskJournal | None = None
self._orphan_purge_done = False
@@ -1125,24 +1129,29 @@ class InternalDBOSAdapter(InternalRunAdapter):
self._resolved_pool = await self._pool_provider.get()
return self._resolved_pool
def _get_or_create_state_store(self) -> StateStore[Any]:
"""Get or lazily create the state store.
def _get_or_create_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any]:
"""Get or lazily create the per-namespace state store.
For PostgreSQL, the pool must be resolved first via _resolve_pool().
Call _ensure_resources() before accessing the state store.
"""
if self._state_store is None:
store = self._state_stores.get(namespace)
if store is None:
if self._resolved_pool is not None:
self._state_store = PostgresStateStore(
store = PostgresStateStore(
pool=self._resolved_pool,
run_id=self._run_id,
namespace=namespace,
state_type=cast(type[Any], self._state_type),
schema=self._schema,
)
elif self._db_path is not None:
self._state_store = SqliteStateStore(
store = SqliteStateStore(
db_path=self._db_path,
run_id=self._run_id,
namespace=namespace,
state_type=cast(type[Any], self._state_type),
)
else:
@@ -1150,10 +1159,13 @@ class InternalDBOSAdapter(InternalRunAdapter):
"No pool or db_path configured for state store. "
"Ensure the runtime pool is initialized before accessing state."
)
return self._state_store
self._state_stores[namespace] = store
return store
def get_state_store(self) -> StateStore[Any] | None:
return self._get_or_create_state_store()
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any] | None:
return self._get_or_create_state_store(namespace)
def is_replaying(self) -> bool:
if (
packages/llama-agents-dbos/tests/test_dbos_idle_release.py
+7 -3
View File
@@ -80,7 +80,9 @@ class StubInternalAdapter(InternalRunAdapter):
async def close(self) -> None:
self.closed = True
def get_state_store(self) -> StateStore[Any] | None:
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any] | None:
return None
@@ -107,7 +109,9 @@ class StubExternalAdapter(ExternalRunAdapter):
async def get_result(self) -> StopEvent:
return self._result
def get_state_store(self) -> StateStore[Any] | None:
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any] | None:
return None
@@ -820,7 +824,7 @@ async def test_do_resume_carries_over_serialized_state(
# Seed the state store with actual state data
state_store = InMemoryStateStore(MyState(counter=42))
store.state_stores["run-1"] = state_store
store.state_stores[("run-1", ())] = state_store
claim = await _claim_resume(lifecycle)
await decorator._do_resume("run-1", resume_claim=claim)
packages/llama-agents-integration-tests/src/llama_agents_integration_tests/fake_agent_data.py
+2
View File
@@ -210,6 +210,7 @@ def create_agent_data_state_store(
monkeypatch: pytest.MonkeyPatch,
run_id: str,
state_type: type[Any] | None = None,
namespace: tuple[str, ...] = (),
) -> AgentDataStateStore[Any]:
"""Create an AgentDataStateStore with httpx patched to use the fake backend."""
client = AgentDataClient(
@@ -222,6 +223,7 @@ def create_agent_data_state_store(
store = AgentDataStateStore(
client=client,
run_id=run_id,
namespace=namespace,
state_type=state_type,
)
return store
packages/llama-agents-integration-tests/tests/test_state_store_matrix.py
+4 -2
View File
@@ -116,12 +116,14 @@ async def _create_postgres_pool(dsn: str) -> asyncpg.Pool:
await conn.execute("CREATE SCHEMA IF NOT EXISTS dbos")
await conn.execute("""
CREATE TABLE IF NOT EXISTS dbos.workflow_state (
run_id VARCHAR(255) PRIMARY KEY,
run_id VARCHAR(255) NOT NULL,
namespace VARCHAR(255) NOT NULL DEFAULT '',
state_json TEXT NOT NULL,
state_type VARCHAR(255),
state_module VARCHAR(255),
created_at TIMESTAMPTZ,
updated_at TIMESTAMPTZ
updated_at TIMESTAMPTZ,
PRIMARY KEY (run_id, namespace)
)
""")
return pool
packages/llama-agents-server/src/llama_agents/server/_runtime/server_runtime.py
+17 -8
View File
@@ -73,22 +73,31 @@ class _ServerInternalRunAdapter(BaseInternalRunAdapterDecorator):
self._runtime = runtime
self._store = runtime._store
self._state_type = state_type
self._state_store: StateStore[Any] | None = None
self._state_stores: dict[tuple[str, ...], StateStore[Any]] = {}
self._write_lock: asyncio.Lock | None = None
@override
def get_state_store(self) -> StateStore[Any]:
if self._state_store is not None:
return self._state_store
initial = self._runtime._initial_state.pop(self.run_id, None)
def get_state_store(self, namespace: tuple[str, ...] = ()) -> StateStore[Any]:
cached = self._state_stores.get(namespace)
if cached is not None:
return cached
initial = (
None if namespace else self._runtime._initial_state.pop(self.run_id, None)
)
if initial is not None:
serialized_state, serializer = initial
store = self._store.create_state_store(
self.run_id, self._state_type, serialized_state, serializer
self.run_id,
self._state_type,
serialized_state,
serializer,
namespace=namespace,
)
else:
store = self._store.create_state_store(self.run_id, self._state_type)
self._state_store = store
store = self._store.create_state_store(
self.run_id, self._state_type, namespace=namespace
)
self._state_stores[namespace] = store
return store
@override
packages/llama-agents-server/src/llama_agents/server/_store/abstract_workflow_store.py
+22 -13
View File
@@ -116,9 +116,9 @@ class AbstractWorkflowStore(ABC):
# Weak-valued by default so facades die with their last consumer.
# Backends needing a different lifecycle (strong refs + explicit
# eviction) assign a plain dict in their __init__.
self._state_store_cache: MutableMapping[str, StateStoreFacade[Any]] = (
weakref.WeakValueDictionary()
)
self._state_store_cache: MutableMapping[
tuple[str, tuple[str, ...]], StateStoreFacade[Any]
] = weakref.WeakValueDictionary()
async def start(self) -> None:
"""Initialize backend resources. Default is a no-op."""
@@ -129,18 +129,22 @@ class AbstractWorkflowStore(ABC):
state_type: type[Any] | None = None,
serialized_state: dict[str, Any] | None = None,
serializer: BaseSerializer | None = None,
namespace: tuple[str, ...] = (),
) -> StateStore[Any]:
"""Return the per-run state store, creating and caching it on first use.
"""Return the per-(run, namespace) state store, caching on first use.
One facade per run per process, so its write lock is a real guarantee.
If *serialized_state* is provided, it is staged as a seed on the
(possibly already handed-out) facade: validation is eager, the I/O to
materialize it stays lazy (first async state access or handoff).
Namespace is a key dimension alongside ``run_id``: the default ``()``
is the root namespace and reproduces the single-record behavior. One
facade per (run, namespace) per process, so its write lock is a real
guarantee. If *serialized_state* is provided, it is staged as a seed on
the (possibly already handed-out) facade: validation is eager, the I/O
to materialize it stays lazy (first async state access or handoff).
"""
store = self._state_store_cache.get(run_id)
cache_key = (run_id, namespace)
store = self._state_store_cache.get(cache_key)
if store is None:
store = self._build_state_store(run_id, state_type, serializer)
self._state_store_cache[run_id] = store
store = self._build_state_store(run_id, namespace, state_type, serializer)
self._state_store_cache[cache_key] = store
elif state_type is not None and store.state_type is DictState:
# An earlier type-less caller (e.g. handler continuation) must
# not shadow the workflow's concrete state type.
@@ -149,14 +153,20 @@ class AbstractWorkflowStore(ABC):
store.add_seed(serialized_state, serializer or JsonSerializer())
return store
def _evict_run_state_stores(self, run_id: str) -> None:
"""Drop every cached namespace facade for *run_id* (all namespaces)."""
for key in [k for k in self._state_store_cache if k[0] == run_id]:
self._state_store_cache.pop(key, None)
@abstractmethod
def _build_state_store(
self,
run_id: str,
namespace: tuple[str, ...],
state_type: type[Any] | None,
serializer: BaseSerializer | None,
) -> StateStoreFacade[Any]:
"""Construct the backend facade for a run. No caching, no seeding."""
"""Construct the backend facade for a (run, namespace). No caching."""
@abstractmethod
async def query(self, query: HandlerQuery) -> list[PersistentHandler]: ...
@@ -236,7 +246,6 @@ class AbstractWorkflowStore(ABC):
@staticmethod
def _is_terminal_event(event: StoredEvent) -> bool:
"""Check if a stored event is terminal (StopEvent or subclass, etc.)."""
types = (event.event.types or []) + [event.event.type]
return StopEvent.__name__ in types
packages/llama-agents-server/src/llama_agents/server/_store/agent_data_state_store.py
+118 -32
View File
@@ -23,12 +23,19 @@ from .agent_data_client import AgentDataClient
MODEL_T = TypeVar("MODEL_T", bound=BaseModel, default=DictState) # type: ignore[reportGeneralTypeIssues]
_FIELD_RUN_ID = "run_id"
_FIELD_NAMESPACE = "namespace"
_STATE_PAGE_SIZE = 100
class _AgentDataStateRecord(BaseModel):
"""Validates the shape persisted in the Agent Data API."""
"""Validates the shape persisted in the Agent Data API.
Root items carry no ``namespace`` field (matching pre-namespace items);
the ``""`` default reads both back as the root namespace.
"""
run_id: str
namespace: str = ""
data: str
state_type: str | None = None
state_module: str | None = None
@@ -55,10 +62,14 @@ class _AgentDataStateStorage:
*,
client: AgentDataClient,
run_id: str,
namespace: tuple[str, ...] = (),
collection: str = "workflow_state",
) -> None:
self._client = client
self._run_id = run_id
self._namespace = namespace
# Persisted key: () -> "" (today's single root item), ("child",) -> "child".
self._namespace_key = "/".join(namespace)
self._collection = collection
self._item_id: str | None = None
@@ -71,16 +82,37 @@ class _AgentDataStateStorage:
# HTTP-backed: no per-call connections, the storage scopes itself.
yield self
async def _load_record(self) -> _AgentDataStateRecord | None:
async def _matching_item(self) -> tuple[str, _AgentDataStateRecord] | None:
"""Find this namespace's item for the run.
Filtering is by ``run_id`` *and* ``namespace`` server-side, so the query
returns this namespace's single row directly rather than scanning the
run's first page and filtering in Python.
Root rows carry no ``namespace`` field (same shape as pre-namespace
rows), and the backend's ``eq: null`` matches a missing field, so the
root lookup is a single query covering legacy and new rows alike.
"""
namespace_filter = None if self._namespace_key == "" else self._namespace_key
items = await self._client.search(
self._collection,
{_FIELD_RUN_ID: {"eq": self._run_id}},
page_size=1,
{
_FIELD_RUN_ID: {"eq": self._run_id},
_FIELD_NAMESPACE: {"eq": namespace_filter},
},
)
if not items:
for item in items:
record = _AgentDataStateRecord.model_validate(item["data"])
if record.namespace == self._namespace_key:
return item["id"], record
return None
async def _load_record(self) -> _AgentDataStateRecord | None:
match = await self._matching_item()
if match is None:
return None
self._item_id = items[0]["id"]
return _AgentDataStateRecord.model_validate(items[0]["data"])
self._item_id, record = match
return record
async def load(self) -> StateRecord | None:
record = await self._load_record()
@@ -91,26 +123,26 @@ class _AgentDataStateStorage:
async def save(self, record: StateRecord) -> None:
stored = _AgentDataStateRecord(
run_id=self._run_id,
namespace=self._namespace_key,
data=record.data,
state_type=record.state_type,
state_module=record.state_module,
)
payload = stored.model_dump()
if self._namespace_key == "":
# Root rows keep the pre-namespace shape (no namespace field) so
# the root lookup stays a single eq-null query for all rows.
del payload[_FIELD_NAMESPACE]
if self._item_id is not None:
await self._client.update_item(self._item_id, payload)
return
match = await self._matching_item()
if match is not None:
self._item_id = match[0]
await self._client.update_item(self._item_id, payload)
else:
items = await self._client.search(
self._collection,
{_FIELD_RUN_ID: {"eq": self._run_id}},
page_size=1,
)
if items:
item_id = items[0]["id"]
self._item_id = item_id
await self._client.update_item(item_id, payload)
else:
result = await self._client.create(self._collection, payload)
self._item_id = result["id"]
result = await self._client.create(self._collection, payload)
self._item_id = result["id"]
def to_handle(self) -> dict[str, Any]:
payload = AgentDataSerializedState(
@@ -125,21 +157,69 @@ class _AgentDataStateStorage:
return None
return AgentDataSerializedState.model_validate(payload)
async def copy_from_handle(self, handle: AgentDataSerializedState) -> None:
"""Copy the source target's record into this one (no-op if absent).
async def _all_run_items(
self, collection: str, run_id: str
) -> AsyncIterator[dict[str, Any]]:
"""Yield every state item for a run, paginating in full.
Goes through ``save`` so ``_item_id`` stays consistent with the
copied row.
A run may hold more namespace rows than a single search page (one per
child invocation), so iterate with a keyset cursor over ``namespace``
(unique per run) rather than reading only the first page.
The root row has no ``namespace`` field, so comparison filters can
never match it — fetch it with its own eq-null query, then paginate
the namespaced rows (every child key sorts after ``""``).
"""
source = _AgentDataStateStorage(
client=self._client,
run_id=handle.run_id,
collection=handle.collection,
root_items = await self._client.search(
collection,
{
_FIELD_RUN_ID: {"eq": run_id},
_FIELD_NAMESPACE: {"eq": None},
},
)
record = await source.load()
if record is None:
return
await self.save(record)
for item in root_items:
yield item
cursor = ""
while True:
filters: dict[str, Any] = {
_FIELD_RUN_ID: {"eq": run_id},
_FIELD_NAMESPACE: {"gt": cursor},
}
page = await self._client.search(
collection,
filters,
page_size=_STATE_PAGE_SIZE,
order_by=_FIELD_NAMESPACE,
)
for item in page:
yield item
cursor = _AgentDataStateRecord.model_validate(item["data"]).namespace
if len(page) < _STATE_PAGE_SIZE:
return
async def copy_from_handle(self, handle: AgentDataSerializedState) -> None:
"""Copy every namespace item of the source run into this run.
The source's items are enumerated by ``run_id`` (paginated in full) and
each is saved under the same namespace in this run. Saves go through
per-namespace storages so each ``_item_id`` stays consistent.
"""
async for item in self._all_run_items(handle.collection, handle.run_id):
source = _AgentDataStateRecord.model_validate(item["data"])
namespace = tuple(source.namespace.split("/")) if source.namespace else ()
dest = _AgentDataStateStorage(
client=self._client,
run_id=self._run_id,
namespace=namespace,
collection=self._collection,
)
await dest.save(
StateRecord(
data=source.data,
state_type=source.state_type,
state_module=source.state_module,
)
)
class AgentDataStateStore(StateStoreFacade[MODEL_T], Generic[MODEL_T]):
@@ -156,12 +236,18 @@ class AgentDataStateStore(StateStoreFacade[MODEL_T], Generic[MODEL_T]):
*,
client: AgentDataClient,
run_id: str,
namespace: tuple[str, ...] = (),
state_type: type[MODEL_T] | None = None,
collection: str = "workflow_state",
serializer: BaseSerializer | None = None,
) -> None:
super().__init__(
_AgentDataStateStorage(client=client, run_id=run_id, collection=collection),
_AgentDataStateStorage(
client=client,
run_id=run_id,
namespace=namespace,
collection=collection,
),
state_type,
serializer,
)
packages/llama-agents-server/src/llama_agents/server/_store/agent_data_store.py
+3 -1
View File
@@ -169,7 +169,7 @@ class AgentDataStore(AbstractWorkflowStore):
# Clean up sequence counters and cached state store
self._event_sequences.pop(run_id, None)
self._tick_sequences.pop(run_id, None)
self._state_store_cache.pop(run_id, None)
self._evict_run_state_stores(run_id)
# ------------------------------------------------------------------
# Sequence helpers
@@ -491,12 +491,14 @@ class AgentDataStore(AbstractWorkflowStore):
def _build_state_store(
self,
run_id: str,
namespace: tuple[str, ...],
state_type: type[Any] | None,
serializer: BaseSerializer | None,
) -> AgentDataStateStore[Any]:
return AgentDataStateStore(
client=self._client,
run_id=run_id,
namespace=namespace,
state_type=state_type,
collection=f"{self._collection}_state",
serializer=serializer,
packages/llama-agents-server/src/llama_agents/server/_store/memory_workflow_store.py
+3 -2
View File
@@ -70,7 +70,7 @@ class MemoryWorkflowStore(AbstractWorkflowStore):
self.ticks: dict[str, list[StoredTick]] = {}
# Strong refs: facades live until eviction. Public alias kept for
# tests/plugins that inject stores; the ABC template reads the cache.
self.state_stores: dict[str, StateStoreFacade[Any]] = {}
self.state_stores: dict[tuple[str, tuple[str, ...]], StateStoreFacade[Any]] = {}
self._state_store_cache = self.state_stores
self._conditions: weakref.WeakValueDictionary[str, asyncio.Condition] = (
weakref.WeakValueDictionary()
@@ -81,6 +81,7 @@ class MemoryWorkflowStore(AbstractWorkflowStore):
def _build_state_store(
self,
run_id: str,
namespace: tuple[str, ...],
state_type: type[Any] | None,
serializer: BaseSerializer | None,
) -> InMemoryStateStore[Any]:
@@ -134,7 +135,7 @@ class MemoryWorkflowStore(AbstractWorkflowStore):
if run_id is not None:
self.events.pop(run_id, None)
self.ticks.pop(run_id, None)
self.state_stores.pop(run_id, None)
self._evict_run_state_stores(run_id)
def _get_or_create_condition(self, run_id: str) -> asyncio.Condition:
cond = self._conditions.get(run_id)
packages/llama-agents-server/src/llama_agents/server/_store/postgres/migrations/0002_namespace_state.sql
+8
View File
@@ -0,0 +1,8 @@
-- migration: 2
-- Per-namespace state records: namespace becomes a key dimension alongside
-- run_id. Existing single rows are root-namespace rows (namespace = '').
ALTER TABLE workflow_state ADD COLUMN IF NOT EXISTS namespace VARCHAR(255) NOT NULL DEFAULT '';
ALTER TABLE workflow_state DROP CONSTRAINT IF EXISTS workflow_state_pkey;
ALTER TABLE workflow_state ADD CONSTRAINT workflow_state_pkey PRIMARY KEY (run_id, namespace);
packages/llama-agents-server/src/llama_agents/server/_store/postgres_state_store.py
+25 -10
View File
@@ -40,11 +40,15 @@ class _PostgresStateStorage:
self,
pool: asyncpg.Pool,
run_id: str,
namespace: tuple[str, ...] = (),
schema: str | None = None,
connection: asyncpg.Connection | asyncpg.pool.PoolConnectionProxy | None = None,
) -> None:
self._pool = pool
self._run_id = run_id
self._namespace = namespace
# Persisted key: () -> "" (today's single root row), ("child",) -> "child".
self._namespace_key = "/".join(namespace)
self._schema = schema
self._shared_conn = connection
@@ -81,15 +85,17 @@ class _PostgresStateStorage:
return
async with self._pool.acquire() as conn:
yield _PostgresStateStorage(
self._pool, self._run_id, self._schema, connection=conn
self._pool, self._run_id, self._namespace, self._schema, connection=conn
)
async def load(self) -> StateRecord | None:
"""Load raw state from the database."""
async with self._acquire() as conn:
row = await conn.fetchrow(
f"SELECT state_json FROM {self._table_ref} WHERE run_id = $1",
f"SELECT state_json FROM {self._table_ref} "
"WHERE run_id = $1 AND namespace = $2",
self._run_id,
self._namespace_key,
)
if row is None:
return None
@@ -101,15 +107,16 @@ class _PostgresStateStorage:
async with self._acquire() as conn:
await conn.execute(
f"""
INSERT INTO {self._table_ref} (run_id, state_json, state_type, state_module, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, $6)
ON CONFLICT(run_id) DO UPDATE SET
INSERT INTO {self._table_ref} (run_id, namespace, state_json, state_type, state_module, created_at, updated_at)
VALUES ($1, $2, $3, $4, $5, $6, $7)
ON CONFLICT(run_id, namespace) DO UPDATE SET
state_json = EXCLUDED.state_json,
state_type = EXCLUDED.state_type,
state_module = EXCLUDED.state_module,
updated_at = EXCLUDED.updated_at
""",
self._run_id,
self._namespace_key,
record.data,
record.state_type,
record.state_module,
@@ -129,15 +136,20 @@ class _PostgresStateStorage:
return PostgresSerializedState.model_validate(payload)
async def copy_from_handle(self, handle: PostgresSerializedState) -> None:
"""Copy state from another run's row using SQL INSERT...SELECT."""
"""Copy every namespace row from another run using INSERT...SELECT.
The ``run_id`` filter spans all namespaces, so a single statement
copies the source run's root and every child namespace; ``namespace``
is carried through unchanged.
"""
now = _utc_now()
async with self._acquire() as conn:
await conn.execute(
f"""
INSERT INTO {self._table_ref} (run_id, state_json, state_type, state_module, created_at, updated_at)
SELECT $1, state_json, state_type, state_module, $2, $3
INSERT INTO {self._table_ref} (run_id, namespace, state_json, state_type, state_module, created_at, updated_at)
SELECT $1, namespace, state_json, state_type, state_module, $2, $3
FROM {self._table_ref} WHERE run_id = $4
ON CONFLICT(run_id) DO UPDATE SET
ON CONFLICT(run_id, namespace) DO UPDATE SET
state_json = EXCLUDED.state_json,
state_type = EXCLUDED.state_type,
state_module = EXCLUDED.state_module,
@@ -157,12 +169,15 @@ class PostgresStateStore(StateStoreFacade[MODEL_T], Generic[MODEL_T]):
self,
pool: asyncpg.Pool,
run_id: str,
namespace: tuple[str, ...] = (),
state_type: type[MODEL_T] | None = None,
serializer: BaseSerializer | None = None,
schema: str | None = None,
) -> None:
super().__init__(
_PostgresStateStorage(pool, run_id, schema), state_type, serializer
_PostgresStateStorage(pool, run_id, namespace, schema),
state_type,
serializer,
)
@classmethod
packages/llama-agents-server/src/llama_agents/server/_store/postgres_workflow_store.py
+2
View File
@@ -298,6 +298,7 @@ class PostgresWorkflowStore(AbstractWorkflowStore):
def _build_state_store(
self,
run_id: str,
namespace: tuple[str, ...],
state_type: type[Any] | None,
serializer: BaseSerializer | None,
) -> PostgresStateStore[Any]:
@@ -308,6 +309,7 @@ class PostgresWorkflowStore(AbstractWorkflowStore):
return PostgresStateStore(
pool=self._pool,
run_id=run_id,
namespace=namespace,
state_type=state_type,
serializer=serializer,
schema=self._schema,
packages/llama-agents-server/src/llama_agents/server/_store/sqlite/migrations/0005_namespace_state.sql
+23
View File
@@ -0,0 +1,23 @@
-- migration: 5
-- Per-namespace state records: namespace becomes a key dimension alongside
-- run_id. SQLite cannot ALTER a primary key, so rebuild the table with the
-- composite PK and migrate existing rows as root-namespace rows (namespace '').
CREATE TABLE workflow_state_new (
run_id TEXT NOT NULL,
namespace TEXT NOT NULL DEFAULT '',
state_json TEXT NOT NULL DEFAULT '{}',
state_type TEXT NOT NULL DEFAULT 'DictState',
state_module TEXT NOT NULL DEFAULT 'workflows.context.state_store',
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL,
PRIMARY KEY (run_id, namespace)
);
INSERT INTO workflow_state_new (run_id, namespace, state_json, state_type, state_module, created_at, updated_at)
SELECT run_id, '', state_json, state_type, state_module, created_at, updated_at
FROM workflow_state;
DROP TABLE workflow_state;
ALTER TABLE workflow_state_new RENAME TO workflow_state;
packages/llama-agents-server/src/llama_agents/server/_store/sqlite/sqlite_state_store.py
+26 -10
View File
@@ -40,10 +40,14 @@ class _SqliteStateStorage:
self,
db_path: str,
run_id: str,
namespace: tuple[str, ...] = (),
connection: sqlite3.Connection | None = None,
) -> None:
self._db_path = db_path
self._run_id = run_id
self._namespace = namespace
# Persisted key: () -> "" (today's single root row), ("child",) -> "child".
self._namespace_key = "/".join(namespace)
self._shared_conn = connection
@property
@@ -73,7 +77,9 @@ class _SqliteStateStorage:
return
conn = sqlite3.connect(self._db_path, timeout=30.0)
try:
yield _SqliteStateStorage(self._db_path, self._run_id, connection=conn)
yield _SqliteStateStorage(
self._db_path, self._run_id, self._namespace, connection=conn
)
finally:
conn.close()
@@ -82,8 +88,9 @@ class _SqliteStateStorage:
with self._connect() as conn:
cursor = conn.cursor()
cursor.execute(
"SELECT state_json FROM workflow_state WHERE run_id = ?",
(self._run_id,),
"SELECT state_json FROM workflow_state "
"WHERE run_id = ? AND namespace = ?",
(self._run_id, self._namespace_key),
)
row = cursor.fetchone()
if row is None:
@@ -96,9 +103,9 @@ class _SqliteStateStorage:
now = _utc_now().isoformat()
conn.execute(
"""
INSERT INTO workflow_state (run_id, state_json, state_type, state_module, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?)
ON CONFLICT(run_id) DO UPDATE SET
INSERT INTO workflow_state (run_id, namespace, state_json, state_type, state_module, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?)
ON CONFLICT(run_id, namespace) DO UPDATE SET
state_json = excluded.state_json,
state_type = excluded.state_type,
state_module = excluded.state_module,
@@ -106,6 +113,7 @@ class _SqliteStateStorage:
""",
(
self._run_id,
self._namespace_key,
record.data,
record.state_type,
record.state_module,
@@ -125,13 +133,18 @@ class _SqliteStateStorage:
return SqliteSerializedState.model_validate(payload)
async def copy_from_handle(self, handle: SqliteSerializedState) -> None:
"""Copy state from another run's row using SQL INSERT...SELECT."""
"""Copy every namespace row from another run using INSERT...SELECT.
The ``run_id`` filter spans all namespaces, so a single statement
copies the source run's root and every child namespace; ``namespace``
is carried through unchanged.
"""
with self._connect() as conn:
now = _utc_now().isoformat()
conn.execute(
"""
INSERT OR REPLACE INTO workflow_state (run_id, state_json, state_type, state_module, created_at, updated_at)
SELECT ?, state_json, state_type, state_module, ?, ?
INSERT OR REPLACE INTO workflow_state (run_id, namespace, state_json, state_type, state_module, created_at, updated_at)
SELECT ?, namespace, state_json, state_type, state_module, ?, ?
FROM workflow_state WHERE run_id = ?
""",
(self._run_id, now, now, handle.run_id),
@@ -146,13 +159,16 @@ class SqliteStateStore(StateStoreFacade[MODEL_T], Generic[MODEL_T]):
self,
db_path: str,
run_id: str,
namespace: tuple[str, ...] = (),
state_type: type[MODEL_T] | None = None,
serializer: BaseSerializer | None = None,
connection: sqlite3.Connection | None = None,
) -> None:
self._db_path = db_path
super().__init__(
_SqliteStateStorage(db_path, run_id, connection), state_type, serializer
_SqliteStateStorage(db_path, run_id, namespace, connection),
state_type,
serializer,
)
@classmethod
packages/llama-agents-server/src/llama_agents/server/_store/sqlite/sqlite_workflow_store.py
+2
View File
@@ -93,12 +93,14 @@ class SqliteWorkflowStore(AbstractWorkflowStore):
def _build_state_store(
self,
run_id: str,
namespace: tuple[str, ...],
state_type: type[Any] | None,
serializer: BaseSerializer | None,
) -> SqliteStateStore[Any]:
return SqliteStateStore(
db_path=self.db_path,
run_id=run_id,
namespace=namespace,
state_type=state_type,
serializer=serializer,
connection=self._persistent_conn,
packages/llama-agents-server/tests/server/test_agent_data_store.py
+134
View File
@@ -4,6 +4,7 @@
from __future__ import annotations
import asyncio
import json
from datetime import datetime, timezone
from typing import Any
@@ -1028,3 +1029,136 @@ async def test_persist_error_does_not_block_in_memory_delivery(
await store.append_event("run-1", make_envelope(event=StopEvent(data="done")))
await asyncio.wait_for(task, timeout=2.0)
assert len(collected) == 3
# ---------------------------------------------------------------------------
# Per-namespace records
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_namespace_round_trip_and_isolation(
backend: FakeAgentDataBackend, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Root and child namespaces persist as independent items under one run."""
root = create_agent_data_state_store(backend, monkeypatch, "run-ns")
child = create_agent_data_state_store(
backend, monkeypatch, "run-ns", namespace=("child",)
)
await root.set("k", "root-val")
await child.set("k", "child-val")
root2 = create_agent_data_state_store(backend, monkeypatch, "run-ns")
child2 = create_agent_data_state_store(
backend, monkeypatch, "run-ns", namespace=("child",)
)
assert await root2.get("k") == "root-val"
assert await child2.get("k") == "child-val"
@pytest.mark.asyncio
async def test_single_namespace_lookup_beyond_first_search_page(
backend: FakeAgentDataBackend, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A namespace lookup filters by run_id AND namespace server-side, so it
finds a row even when the run holds more than one search page of rows."""
n = 150
for i in range(n):
s = create_agent_data_state_store(
backend, monkeypatch, "run-many", namespace=(f"child#{i}",)
)
await s.set("k", f"v{i}")
# A fresh store (no cache) for the last-written namespace must find its row.
late = create_agent_data_state_store(
backend, monkeypatch, "run-many", namespace=("child#149",)
)
assert await late.get("k") == "v149"
@pytest.mark.asyncio
async def test_copy_from_handle_copies_every_namespace_across_pages(
store: AgentDataStore,
) -> None:
"""Copying a run via its durable handle reproduces every namespace row,
paginating past the first search page."""
n = 150
root = store.create_state_store("run-src")
await root.set("k", "root")
for i in range(n):
child = store.create_state_store("run-src", namespace=(f"child#{i}",))
await child.set("k", f"v{i}")
handle = root.to_dict(JsonSerializer())
target_root = AgentDataStateStore.from_dict(
handle, JsonSerializer(), client=store._client, run_id="run-target"
)
await target_root.ensure_seeded()
assert await store.create_state_store("run-target").get("k") == "root"
for i in range(n):
tchild = store.create_state_store("run-target", namespace=(f"child#{i}",))
assert await tchild.get("k") == f"v{i}"
@pytest.mark.asyncio
async def test_root_lookup_is_single_query_matching_pre_namespace_rows(
backend: FakeAgentDataBackend, monkeypatch: pytest.MonkeyPatch
) -> None:
"""A row without a namespace field IS the root shape: the root lookup is a
single eq-null query that matches pre-namespace and new rows alike."""
serializer = JsonSerializer()
backend.create(
"test-deploy",
"workflow_state",
{
"run_id": "legacy-run",
"data": json.dumps({"_data": {"k": serializer.serialize("legacy")}}),
"state_type": "DictState",
"state_module": "workflows.context.state_store",
},
)
search_filters: list[dict[str, Any] | None] = []
original_search = backend.search
def counting_search(
deployment_name: str,
collection: str,
filters: dict[str, Any] | None = None,
page_size: int = 100,
order_by: str | None = None,
) -> list[dict[str, Any]]:
search_filters.append(filters)
return original_search(
deployment_name, collection, filters, page_size, order_by
)
monkeypatch.setattr(backend, "search", counting_search)
root = create_agent_data_state_store(backend, monkeypatch, "legacy-run")
assert await root.get("k") == "legacy"
assert search_filters == [
{"run_id": {"eq": "legacy-run"}, "namespace": {"eq": None}}
]
@pytest.mark.asyncio
async def test_root_write_omits_namespace_field(
backend: FakeAgentDataBackend, monkeypatch: pytest.MonkeyPatch
) -> None:
"""Root rows persist without a namespace field; child rows carry theirs."""
root = create_agent_data_state_store(backend, monkeypatch, "run-shape")
child = create_agent_data_state_store(
backend, monkeypatch, "run-shape", namespace=("child",)
)
await root.set("k", "r")
await child.set("k", "c")
rows = [item["data"] for item in backend.search("test-deploy", "workflow_state")]
assert len(rows) == 2
root_rows = [row for row in rows if "namespace" not in row]
child_rows = [row for row in rows if row.get("namespace") == "child"]
assert len(root_rows) == 1
assert len(child_rows) == 1
packages/llama-agents-server/tests/server/test_event_interceptor.py
+3 -1
View File
@@ -51,7 +51,9 @@ class _RecordingInternalAdapter(InternalRunAdapter):
async def close(self) -> None:
self.closed = True
def get_state_store(self) -> StateStore[Any] | None:
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any] | None:
return None
packages/llama-agents-server/tests/server/test_memory_workflow_store.py
+19 -2
View File
@@ -478,7 +478,7 @@ async def test_max_completed_cleans_up_events_ticks_and_state() -> None:
)
assert "run-old" in store.events
assert "run-old" in store.ticks
assert "run-old" in store.state_stores
assert ("run-old", ()) in store.state_stores
await _insert(
store,
@@ -490,12 +490,29 @@ async def test_max_completed_cleans_up_events_ticks_and_state() -> None:
assert "run-old" not in store.events
assert "run-old" not in store.ticks
assert "run-old" not in store.state_stores
assert ("run-old", ()) not in store.state_stores
remaining = await store.query(HandlerQuery())
assert len(remaining) == 1
assert remaining[0].handler_id == "h-new"
async def test_evict_run_state_stores_drops_every_namespace() -> None:
"""Run-scoped eviction removes all of a run's namespace facades at once."""
store = MemoryWorkflowStore()
store.create_state_store("run-x")
store.create_state_store("run-x", namespace=("child",))
store.create_state_store("run-other")
assert ("run-x", ()) in store.state_stores
assert ("run-x", ("child",)) in store.state_stores
store._evict_run_state_stores("run-x")
assert ("run-x", ()) not in store.state_stores
assert ("run-x", ("child",)) not in store.state_stores
# Other runs are untouched.
assert ("run-other", ()) in store.state_stores
@pytest.mark.asyncio
async def test_max_completed_eviction_via_update_handler_status() -> None:
"""Eviction triggers when status changes to terminal via update_handler_status."""
packages/llama-agents-server/tests/server/test_postgres_state_store.py
+58 -8
View File
@@ -4,6 +4,7 @@
from __future__ import annotations
import json
from typing import AsyncGenerator
import asyncpg
@@ -35,12 +36,14 @@ async def pool(postgres_dsn: str) -> AsyncGenerator[asyncpg.Pool, None]:
await conn.execute(f"CREATE SCHEMA IF NOT EXISTS {SCHEMA}")
await conn.execute(f"""
CREATE TABLE IF NOT EXISTS {SCHEMA}.workflow_state (
run_id VARCHAR(255) PRIMARY KEY,
run_id VARCHAR(255) NOT NULL,
namespace VARCHAR(255) NOT NULL DEFAULT '',
state_json TEXT NOT NULL,
state_type VARCHAR(255),
state_module VARCHAR(255),
created_at TIMESTAMPTZ,
updated_at TIMESTAMPTZ
updated_at TIMESTAMPTZ,
PRIMARY KEY (run_id, namespace)
)
""")
await conn.execute(f"DELETE FROM {SCHEMA}.workflow_state")
@@ -324,16 +327,18 @@ class FakeConnection:
def __init__(self, rows: dict[str, str]) -> None:
self._rows = rows
async def fetchrow(self, query: str, run_id: str) -> dict[str, str] | None:
state_json = self._rows.get(run_id)
async def fetchrow(
self, query: str, run_id: str, namespace: str
) -> dict[str, str] | None:
state_json = self._rows.get(f"{run_id}\x00{namespace}")
if state_json is None:
return None
return {"state_json": state_json}
async def execute(self, query: str, *args: object) -> None:
# Save upsert: (run_id, state_json, state_type, state_module, now, now)
run_id, state_json = str(args[0]), str(args[1])
self._rows[run_id] = state_json
# Save upsert: (run_id, namespace, state_json, state_type, state_module, now, now)
run_id, namespace, state_json = str(args[0]), str(args[1]), str(args[2])
self._rows[f"{run_id}\x00{namespace}"] = state_json
class FakePoolAcquire:
@@ -376,7 +381,7 @@ async def test_set_state_acquires_exactly_one_pool_connection() -> None:
assert pool.acquire_count == 1
assert pool.release_count == 1
assert "run-conn-count" in pool.rows
assert "run-conn-count\x00" in pool.rows
async def test_from_dict_empty_raises() -> None:
@@ -389,3 +394,48 @@ async def test_from_dict_no_pool_raises() -> None:
PostgresStateStore.from_dict(
{"store_type": "postgres", "run_id": "x"}, JsonSerializer()
)
# -- Per-namespace records --
@pytest.mark.docker
async def test_namespace_round_trip_and_isolation(pool: asyncpg.Pool) -> None:
"""Root and child namespaces persist as independent rows under one run."""
root: PostgresStateStore[DictState] = PostgresStateStore(
pool=pool, run_id="run-ns", schema=SCHEMA
)
child: PostgresStateStore[DictState] = PostgresStateStore(
pool=pool, run_id="run-ns", namespace=("child",), schema=SCHEMA
)
await root.set("k", "root-val")
await child.set("k", "child-val")
root2: PostgresStateStore[DictState] = PostgresStateStore(
pool=pool, run_id="run-ns", schema=SCHEMA
)
child2: PostgresStateStore[DictState] = PostgresStateStore(
pool=pool, run_id="run-ns", namespace=("child",), schema=SCHEMA
)
assert await root2.get("k") == "root-val"
assert await child2.get("k") == "child-val"
@pytest.mark.docker
async def test_pre_migration_root_row_reads_as_root(pool: asyncpg.Pool) -> None:
"""A row written without a namespace (additive default '') reads as root."""
state_json = json.dumps({"_data": {"k": JsonSerializer().serialize("legacy")}})
async with pool.acquire() as conn:
await conn.execute(
f"INSERT INTO {SCHEMA}.workflow_state "
"(run_id, state_json, state_type, state_module, created_at, updated_at) "
"VALUES ($1, $2, 'DictState', 'workflows.context.state_store', now(), now())",
"legacy-run",
state_json,
)
store: PostgresStateStore[DictState] = PostgresStateStore(
pool=pool, run_id="legacy-run", schema=SCHEMA
)
assert await store.get("k") == "legacy"
packages/llama-agents-server/tests/server/test_runtime_decorators.py
+6 -2
View File
@@ -53,7 +53,9 @@ class StubInternalAdapter(InternalRunAdapter):
async def close(self) -> None:
self.closed = True
def get_state_store(self) -> StateStore[Any] | None:
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any] | None:
return None
@@ -77,7 +79,9 @@ class StubExternalAdapter(ExternalRunAdapter):
async def get_result(self) -> StopEvent:
return StopEvent(result="done")
def get_state_store(self) -> StateStore[Any] | None:
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any] | None:
return None
packages/llama-agents-server/tests/server/test_server_runtime.py
+6 -2
View File
@@ -67,7 +67,9 @@ class StubInternalAdapter(InternalRunAdapter):
async def close(self) -> None:
self.closed = True
def get_state_store(self) -> StateStore[Any] | None:
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any] | None:
return None
@@ -91,7 +93,9 @@ class StubExternalAdapter(ExternalRunAdapter):
async def get_result(self) -> StopEvent:
return StopEvent(result="done")
def get_state_store(self) -> StateStore[Any] | None:
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any] | None:
return None
packages/llama-agents-server/tests/server/test_sqlite_state_store.py
+109
View File
@@ -5,6 +5,7 @@
from __future__ import annotations
import asyncio
import json
import sqlite3
from pathlib import Path
from typing import Any
@@ -13,6 +14,7 @@ import pytest
from llama_agents.server import HandlerQuery, SqliteWorkflowStore
from llama_agents.server._store.sqlite.migrate import run_migrations
from llama_agents.server._store.sqlite.sqlite_state_store import (
SqliteSerializedState,
SqliteStateStore,
)
from pydantic import BaseModel
@@ -491,3 +493,110 @@ async def test_sqlite_workflow_store_single_connection_opens_existing_regular_db
await state_store.set("x", 1)
assert await state_store.get("x") == 1
# -- Per-namespace records --
@pytest.mark.asyncio
async def test_namespace_round_trip_and_isolation(db_path: str) -> None:
"""Root and child namespaces persist as independent rows under one run."""
root: SqliteStateStore[DictState] = SqliteStateStore(
db_path=db_path, run_id="run-ns"
)
child: SqliteStateStore[DictState] = SqliteStateStore(
db_path=db_path, run_id="run-ns", namespace=("child",)
)
await root.set("k", "root-val")
await child.set("k", "child-val")
# Fresh facades read straight from their own row — no cross-namespace bleed.
root2: SqliteStateStore[DictState] = SqliteStateStore(
db_path=db_path, run_id="run-ns"
)
child2: SqliteStateStore[DictState] = SqliteStateStore(
db_path=db_path, run_id="run-ns", namespace=("child",)
)
assert await root2.get("k") == "root-val"
assert await child2.get("k") == "child-val"
@pytest.mark.asyncio
async def test_pre_migration_root_row_reads_as_root(tmp_path: Path) -> None:
"""A row written under the pre-namespace schema reads back as root."""
path = str(tmp_path / "legacy.db")
serializer = JsonSerializer()
legacy_state_json = json.dumps({"_data": {"k": serializer.serialize("legacy")}})
conn = sqlite3.connect(path)
try:
# Recreate the schema as it stood at migration 4 (single-column PK,
# no namespace) and mark migrations 1-4 applied.
conn.executescript(
"""
CREATE TABLE schema_migrations (
package TEXT NOT NULL,
version INTEGER NOT NULL,
applied_at TEXT NOT NULL DEFAULT (datetime('now')),
PRIMARY KEY (package, version)
);
CREATE TABLE workflow_state (
run_id TEXT PRIMARY KEY,
state_json TEXT NOT NULL DEFAULT '{}',
state_type TEXT NOT NULL DEFAULT 'DictState',
state_module TEXT NOT NULL DEFAULT 'workflows.context.state_store',
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
);
"""
)
for version in range(1, 5):
conn.execute(
"INSERT INTO schema_migrations (package, version) VALUES ('server', ?)",
(version,),
)
conn.execute(
"INSERT INTO workflow_state "
"(run_id, state_json, state_type, state_module, created_at, updated_at) "
"VALUES (?, ?, 'DictState', 'workflows.context.state_store', '', '')",
("legacy-run", legacy_state_json),
)
conn.commit()
# Apply the pending namespace migration (table rebuild).
run_migrations(conn)
conn.commit()
finally:
conn.close()
store: SqliteStateStore[DictState] = SqliteStateStore(
db_path=path, run_id="legacy-run"
)
assert await store.get("k") == "legacy"
@pytest.mark.asyncio
async def test_copy_reproduces_every_namespace(db_path: str) -> None:
"""Copying a run via its handle reproduces root and child namespaces."""
src_root: SqliteStateStore[DictState] = SqliteStateStore(
db_path=db_path, run_id="run-src"
)
src_child: SqliteStateStore[DictState] = SqliteStateStore(
db_path=db_path, run_id="run-src", namespace=("child",)
)
await src_root.set("k", "root-val")
await src_child.set("k", "child-val")
# Seed the destination root facade from the source handle; ensure_seeded
# drives copy_from_handle, which spans every namespace of the source run.
serializer = JsonSerializer()
dst_root: SqliteStateStore[DictState] = SqliteStateStore(
db_path=db_path, run_id="run-dst"
)
dst_root.add_seed(SqliteSerializedState(run_id="run-src").model_dump(), serializer)
assert await dst_root.get("k") == "root-val"
dst_child: SqliteStateStore[DictState] = SqliteStateStore(
db_path=db_path, run_id="run-dst", namespace=("child",)
)
assert await dst_child.get("k") == "child-val"
packages/llama-agents-server/tests/server/test_workflow_service.py
+2 -2
View File
@@ -329,8 +329,8 @@ async def test_context_from_handler_id_falls_back_to_legacy_state_snapshot(
started_at=datetime.now(timezone.utc),
)
)
state_stores = cast(dict[str, Any], memory_store.state_stores)
state_stores["plugin-run"] = ToDictOnlyStateStore()
state_stores = cast(dict[Any, Any], memory_store.state_stores)
state_stores[("plugin-run", ())] = ToDictOnlyStateStore()
async with server.contextmanager():
ctx = await server._service._context_from_handler_id(
packages/llama-index-workflows/src/workflows/context/external_context.py
+2 -2
View File
@@ -109,7 +109,7 @@ class ExternalContext(Generic[MODEL_T, RunResultT]):
@property
def store(self) -> StateStore[MODEL_T]:
"""Access workflow state store."""
state_store = self._external_adapter.get_state_store()
state_store = self._external_adapter.get_state_store(())
if state_store is None:
raise RuntimeError("State store not available from adapter")
return state_store # type: ignore[return-value]
@@ -162,7 +162,7 @@ class ExternalContext(Generic[MODEL_T, RunResultT]):
# Fetch state store from adapter and serialize
state_data = {}
state_store = self._external_adapter.get_state_store()
state_store = self._external_adapter.get_state_store(())
if state_store is not None:
state_data = state_store.to_dict(active_serializer)
packages/llama-index-workflows/src/workflows/context/internal_context.py
+1 -1
View File
@@ -142,7 +142,7 @@ class InternalContext(Generic[MODEL_T]):
@property
def store(self) -> StateStore[MODEL_T]:
"""Access workflow state store."""
state_store = self._internal_adapter.get_state_store()
state_store = self._internal_adapter.get_state_store(())
if state_store is None:
raise RuntimeError("State store not available from adapter")
return state_store # type: ignore[return-value]
packages/llama-index-workflows/src/workflows/plugins/basic.py
+10 -2
View File
@@ -137,7 +137,11 @@ class InternalAsyncioAdapter(InternalRunAdapter, SnapshottableAdapter):
def replay(self) -> list[WorkflowTick]:
return self._queues.ticks
def get_state_store(self) -> StateStore[Any] | None:
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any] | None:
if namespace:
return None
return self._queues.state_store
@@ -178,7 +182,11 @@ class ExternalAsyncioAdapter(
def replay(self) -> list[WorkflowTick]:
return self._queues.ticks
def get_state_store(self) -> StateStore[Any] | None:
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any] | None:
if namespace:
return None
return self._queues.state_store
async def get_result(self) -> StopEvent:
packages/llama-index-workflows/src/workflows/runtime/runtime_decorators.py
+8 -4
View File
@@ -135,8 +135,10 @@ class BaseInternalRunAdapterDecorator(InternalRunAdapter):
async def close(self) -> None:
await self._decorated.close()
def get_state_store(self) -> StateStore[Any] | None:
return self._decorated.get_state_store()
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any] | None:
return self._decorated.get_state_store(namespace)
async def finalize_step(self) -> None:
await self._decorated.finalize_step()
@@ -188,5 +190,7 @@ class BaseExternalRunAdapterDecorator(ExternalRunAdapter):
async def cancel(self) -> None:
await self._decorated.cancel()
def get_state_store(self) -> StateStore[Any] | None:
return self._decorated.get_state_store()
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any] | None:
return self._decorated.get_state_store(namespace)
packages/llama-index-workflows/src/workflows/runtime/types/plugin.py
+12 -7
View File
@@ -181,11 +181,14 @@ class InternalRunAdapter(ABC):
"""
pass
def get_state_store(self) -> StateStore[Any] | None:
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any] | None:
"""
Get the state store for this workflow run.
Get the per-namespace state store for this workflow run.
Returns the state store from the runtime, or None if not initialized.
``namespace`` selects the namespace's own record (``()`` is the root);
each namespace owns an isolated store. Returns None if not initialized.
Default implementation returns None.
"""
return None
@@ -334,12 +337,14 @@ class ExternalRunAdapter(ABC):
"""
await self.send_event(TickCancelRun())
def get_state_store(self) -> StateStore[Any] | None:
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> StateStore[Any] | None:
"""
Get the state store for this workflow run.
Get the per-namespace state store for this workflow run.
Returns the state store if this adapter owns it, or None if state
is managed externally. Default implementation returns None.
``namespace`` selects the namespace's own record (``()`` is the root).
Returns None if this adapter doesn't own state. Default returns None.
"""
return None
packages/llama-index-workflows/tests/runtime/conftest.py
+3 -1
View File
@@ -192,7 +192,9 @@ class MockRunAdapter(
def has_stream_events(self) -> bool:
return not self._event_stream.empty()
def get_state_store(self) -> "InMemoryStateStore[Any] | None":
def get_state_store(
self, namespace: tuple[str, ...] = ()
) -> "InMemoryStateStore[Any] | None":
return self._state_store
def set_state_store(self, state_store: "InMemoryStateStore[Any]") -> None: