diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 1b2fde61..3bbb505d 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -25,8 +25,11 @@ jobs: llama-agents-dev, llama-agents-client, llama-agents-server, + "llama-agents-dbos" ] exclude: + - package: llama-agents-dbos + python-version: "3.9" # Integration tests on 3.14 run in test-docker job with cross-package coverage - package: llama-agents-integration-tests python-version: "3.9" @@ -56,6 +59,8 @@ jobs: package_dir: packages/llama-agents-client - package: llama-agents-server package_dir: packages/llama-agents-server + - package: llama-agents-dbos + package_dir: packages/llama-agents-dbos steps: - uses: actions/checkout@v4 diff --git a/AGENTS.md b/AGENTS.md index 7d8ddb80..052657a1 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -82,6 +82,14 @@ We use **pytest** with idiomatic pytest patterns. Follow these guidelines: ## Coding Style - Always use `from __future__ import annotations` at the top of each test file. Never use string annotations. -- Include the standard SPDX license header at the top of each test file. +- Include the standard SPDX license header at the top of each file: + ```python + # SPDX-License-Identifier: MIT + # Copyright (c) 2026 LlamaIndex Inc. + ``` - Comments are useful, but avoid fluff. -- Never use inline imports unless required to prevent circular dependencies. +- Import etiquette + - **Never use inline imports** + - **Never use `if TYPE_CHECKING` imports** + - Exceptions to these rules are made only when there are A) Acceptable circular imports or B) real startup performance issues +- Only add `__init__.py` `__all__` exports when a file is legitimately needed for public library consumption. Module level imports should not be used internally. For the most part you should never do this unless explicitly requested to do so diff --git a/examples/dbos_durability/.gitignore b/examples/dbos_durability/.gitignore new file mode 100644 index 00000000..d9ef5bbf --- /dev/null +++ b/examples/dbos_durability/.gitignore @@ -0,0 +1,2 @@ +.last_run_id +*.sqlite3 diff --git a/examples/dbos_durability/counter_example.py b/examples/dbos_durability/counter_example.py new file mode 100644 index 00000000..e5cf12ee --- /dev/null +++ b/examples/dbos_durability/counter_example.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +""" +DBOS Counter Example + +Simple looping workflow that increments a counter until it reaches 20. + +Usage: + python -m examples.dbos_durability.counter_example # Start new + python -m examples.dbos_durability.counter_example --resume # Resume last + python -m examples.dbos_durability.counter_example --clean # Reset state + +Try Ctrl+C mid-run to test resume behavior. +""" + +from __future__ import annotations + +import argparse +import asyncio +import os +import signal +import threading +import time +import uuid +from pathlib import Path + +from dbos import DBOS +from llama_agents.dbos import DBOSRuntime +from pydantic import Field +from workflows import Context, Workflow, step +from workflows.events import Event, StartEvent, StopEvent + +_DIR = Path(__file__).parent +_DB_FILE = _DIR / ".dbos_data.sqlite3" +_RUN_FILE = _DIR / ".last_run_id" + + +class Tick(Event): + count: int = Field(description="Current count") + + +class CounterResult(StopEvent): + final_count: int = Field(description="Final counter value") + + +class CounterWorkflow(Workflow): + """Looping counter workflow - increments until reaching 20.""" + + @step + async def start(self, ctx: Context, ev: StartEvent) -> Tick: + await ctx.store.set("count", 0) + print("[Start] Initializing counter to 0") + return Tick(count=0) + + @step + async def increment(self, ctx: Context, ev: Tick) -> Tick | CounterResult: + count = ev.count + 1 + await ctx.store.set("count", count) + print(f"[Tick {count:2d}] count = {count}") + + if count >= 20: + return CounterResult(final_count=count) + + await asyncio.sleep(0.5) + return Tick(count=count) + + +def run(run_id: str) -> None: + """Run the counter workflow.""" + DBOS( + config={ + "name": "counter-example", + "system_database_url": f"sqlite+pysqlite:///{_DB_FILE}?check_same_thread=false", + "run_admin_server": False, + } + ) + + runtime = DBOSRuntime() + workflow = CounterWorkflow(runtime=runtime) + runtime.launch() + + interrupted = False + + def handle_sigint(signum: int, frame: object) -> None: + nonlocal interrupted + if interrupted: + # Second Ctrl+C - force exit + os._exit(130) + interrupted = True + print("\nInterrupted - workflow state saved. Use --resume to continue.") + + def delayed_exit() -> None: + time.sleep(0.1) + os._exit(130) + + threading.Thread(target=delayed_exit, daemon=True).start() + + # Install signal handler before running + signal.signal(signal.SIGINT, handle_sigint) + + async def _run() -> None: + result = await workflow.run(run_id=run_id) + print(f"\nResult: final_count = {result.final_count}") + + try: + asyncio.run(_run()) + except (KeyboardInterrupt, SystemExit): + pass # Already handled by signal handler + finally: + if not interrupted: + try: + runtime.destroy() + except Exception: + pass + + +def main() -> None: + parser = argparse.ArgumentParser(description="DBOS Counter Example") + parser.add_argument("--resume", action="store_true", help="Resume last workflow") + parser.add_argument("--clean", action="store_true", help="Remove state files") + args = parser.parse_args() + + if args.clean: + for f in [_DB_FILE, _RUN_FILE]: + if f.exists(): + f.unlink() + print(f"Removed {f}") + return + + if args.resume and _RUN_FILE.exists(): + run_id = _RUN_FILE.read_text().strip() + print(f"Resuming: {run_id}") + else: + run_id = f"counter-{uuid.uuid4().hex[:8]}" + _RUN_FILE.write_text(run_id) + print(f"Starting: {run_id}") + + run(run_id) + + +if __name__ == "__main__": + main() diff --git a/packages/llama-agents-dbos/README.md b/packages/llama-agents-dbos/README.md new file mode 100644 index 00000000..500692a5 --- /dev/null +++ b/packages/llama-agents-dbos/README.md @@ -0,0 +1,44 @@ +# LlamaAgents DBOS Runtime + +DBOS durable runtime plugin for LlamaIndex Workflows. + +## Installation + +```bash +pip install llama-agents-dbos +``` + +## Usage + +```python +from llama_agents.dbos import DBOSRuntime +from dbos import DBOS, DBOSConfig +from workflows import Workflow, step, StartEvent, StopEvent + +# Configure DBOS +config: DBOSConfig = { + "name": "my-app", + "system_database_url": "postgresql://...", +} +DBOS(config=config) + +# Create runtime and workflow +runtime = DBOSRuntime() + +class MyWorkflow(Workflow): + @step + async def my_step(self, ev: StartEvent) -> StopEvent: + return StopEvent(result="done") + +workflow = MyWorkflow(runtime=runtime) + +# Launch runtime and run workflow +runtime.launch() +result = await workflow.run() +``` + +## Features + +- Durable workflow execution backed by DBOS +- Automatic step recording and replay +- Distributed workers and recovery support diff --git a/packages/llama-agents-dbos/conftest.py b/packages/llama-agents-dbos/conftest.py new file mode 100644 index 00000000..7eadb41c --- /dev/null +++ b/packages/llama-agents-dbos/conftest.py @@ -0,0 +1,32 @@ +"""Root conftest.py - shared test utilities for all test directories. + +This file is discovered by pytest and provides common utilities +for both tests/ (SQLite) and tests_postgres/ (PostgreSQL). +""" + +from __future__ import annotations + +from pathlib import Path + +from dbos import DBOSConfig + + +def make_test_dbos_config( + name: str, + db_path: Path, +) -> DBOSConfig: + """Create a DBOS config for testing with sensible defaults (SQLite backend). + + Args: + name: The application name for DBOS. + db_path: Path to the SQLite database file. + + Returns: + A DBOSConfig dictionary ready for use with DBOS(). + """ + system_db_url = f"sqlite+pysqlite:///{db_path}?check_same_thread=false" + return { + "name": name, + "system_database_url": system_db_url, + "run_admin_server": False, + } diff --git a/packages/llama-agents-dbos/package.json b/packages/llama-agents-dbos/package.json new file mode 100644 index 00000000..8afde025 --- /dev/null +++ b/packages/llama-agents-dbos/package.json @@ -0,0 +1,7 @@ +{ + "name": "llama-agents-dbos", + "version": "0.1.0", + "private": true, + "license": "MIT", + "scripts": {} +} diff --git a/packages/llama-agents-dbos/pyproject.toml b/packages/llama-agents-dbos/pyproject.toml new file mode 100644 index 00000000..e26aec57 --- /dev/null +++ b/packages/llama-agents-dbos/pyproject.toml @@ -0,0 +1,43 @@ +[build-system] +requires = ["uv_build>=0.9.10,<0.10.0"] +build-backend = "uv_build" + +[dependency-groups] +dev = [ + "basedpyright>=1.31.1", + "pytest>=8.4.0", + "pytest-asyncio>=1.0.0", + "pytest-cov>=6.1.1", + "pytest-timeout>=2.4.0", + "pytest-xdist>=3.0.0", + "ty>=0.0.1,<0.0.9" +] + +[project] +name = "llama-agents-dbos" +version = "0.1.0" +description = "DBOS durable runtime plugin for LlamaIndex Workflows" +readme = "README.md" +license = "MIT" +requires-python = ">=3.9" +dependencies = [ + "dbos>=2.10.0; python_full_version >= '3.10.0'", + "llama-index-workflows>=2.12.0,<3.0.0" +] + +[tool.basedpyright] +typeCheckingMode = "standard" +pythonVersion = "3.10" + +[tool.pytest.ini_options] +asyncio_mode = "auto" +asyncio_default_fixture_loop_scope = "module" +asyncio_default_test_loop_scope = "module" +testpaths = ["tests"] +addopts = "-nauto --timeout=10" + +[tool.uv.build-backend] +module-name = "llama_agents.dbos" + +[tool.uv.sources] +llama-index-workflows = {workspace = true} diff --git a/packages/llama-agents-dbos/src/llama_agents/dbos/__init__.py b/packages/llama-agents-dbos/src/llama_agents/dbos/__init__.py new file mode 100644 index 00000000..e9e386c7 --- /dev/null +++ b/packages/llama-agents-dbos/src/llama_agents/dbos/__init__.py @@ -0,0 +1,13 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +""" +DBOS plugin for LlamaIndex Workflows. + +Provides durable workflow execution backed by DBOS with SQL state storage. +""" + +from __future__ import annotations + +from .runtime import DBOSRuntime + +__all__ = ["DBOSRuntime"] diff --git a/packages/llama-agents-dbos/src/llama_agents/dbos/journal/__init__.py b/packages/llama-agents-dbos/src/llama_agents/dbos/journal/__init__.py new file mode 100644 index 00000000..a26f91fc --- /dev/null +++ b/packages/llama-agents-dbos/src/llama_agents/dbos/journal/__init__.py @@ -0,0 +1,3 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""Journal module for recording task completion order.""" diff --git a/packages/llama-agents-dbos/src/llama_agents/dbos/journal/crud.py b/packages/llama-agents-dbos/src/llama_agents/dbos/journal/crud.py new file mode 100644 index 00000000..472f9673 --- /dev/null +++ b/packages/llama-agents-dbos/src/llama_agents/dbos/journal/crud.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""CRUD operations and table definitions for the workflow journal.""" + +from __future__ import annotations + +from sqlalchemy import Column, Integer, MetaData, String, Table, text +from sqlalchemy.engine import Connection, Engine + +JOURNAL_TABLE_NAME = "workflow_journal" + + +class JournalCrud: + """Database operations for the workflow journal table. + + Initialized with table configuration (name, schema), then provides + methods for inserting, loading, and migrating journal entries. + """ + + def __init__( + self, + table_name: str = JOURNAL_TABLE_NAME, + schema: str | None = None, + ) -> None: + self.table_name = table_name + self.schema = schema + + @property + def _table_ref(self) -> str: + if self.schema: + return f"{self.schema}.{self.table_name}" + return self.table_name + + def _define_table(self, metadata: MetaData) -> Table: + return Table( + self.table_name, + metadata, + Column("id", Integer, primary_key=True, autoincrement=True), + Column("run_id", String(255), nullable=False, index=True), + Column("seq_num", Integer, nullable=False), + Column("task_key", String(512), nullable=False), + ) + + def insert( + self, + conn: Connection, + run_id: str, + seq_num: int, + task_key: str, + ) -> None: + """Insert a new journal entry.""" + conn.execute( + text(f""" + INSERT INTO {self._table_ref} (run_id, seq_num, task_key) + VALUES (:run_id, :seq_num, :task_key) + """), # noqa: S608 + { + "run_id": run_id, + "seq_num": seq_num, + "task_key": task_key, + }, + ) + + def load(self, conn: Connection, run_id: str) -> list[str]: + """Load journal entries for a run, ordered by sequence number.""" + result = conn.execute( + text(f""" + SELECT task_key FROM {self._table_ref} + WHERE run_id = :run_id + ORDER BY seq_num ASC + """), # noqa: S608 + {"run_id": run_id}, + ) + return [row[0] for row in result.fetchall()] + + def run_migrations(self, engine: Engine) -> None: + """Create the journal table if it doesn't exist.""" + metadata = MetaData(schema=self.schema) + table = self._define_table(metadata) + + with engine.begin() as conn: + is_postgres = engine.dialect.name == "postgresql" + if is_postgres and self.schema: + conn.execute(text(f"CREATE SCHEMA IF NOT EXISTS {self.schema}")) # noqa: S608 + table.create(bind=conn, checkfirst=True) diff --git a/packages/llama-agents-dbos/src/llama_agents/dbos/journal/task_journal.py b/packages/llama-agents-dbos/src/llama_agents/dbos/journal/task_journal.py new file mode 100644 index 00000000..5b4cde25 --- /dev/null +++ b/packages/llama-agents-dbos/src/llama_agents/dbos/journal/task_journal.py @@ -0,0 +1,97 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""TaskJournal for deterministic replay of task completion order.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +from sqlalchemy.engine import Engine + +from .crud import JournalCrud + + +class TaskJournal: + """Records task completion order for deterministic replay. + + Stores NamedTask string keys directly (e.g., "step_name:0", "__pull__:1"). + During fresh execution, records which tasks complete and in what order. + During replay, returns the expected task key so the adapter can wait for + the specific task that completed in the original run. + + Uses a dedicated workflow_journal table with one row per entry for efficient + append-only storage. + """ + + def __init__( + self, + run_id: str, + engine: Engine | None = None, + crud: JournalCrud | None = None, + ) -> None: + """Initialize the task journal. + + Args: + run_id: Workflow run ID for this journal. + engine: SQLAlchemy engine. If None, operates in-memory only. + crud: Journal CRUD operations. If None, uses default JournalCrud(). + """ + self._run_id = run_id + self._engine = engine + self._crud = crud or JournalCrud() + self._entries: list[str] | None = None # Lazy loaded + self._replay_index: int = 0 + + async def _run_sync(self, fn: Any, *args: Any) -> Any: + """Run a synchronous function in the default executor.""" + loop = asyncio.get_event_loop() + return await loop.run_in_executor(None, fn, *args) + + async def load(self) -> None: + """Load journal from database. Idempotent - only loads once.""" + if self._entries is not None: + return + + if self._engine is None: + self._entries = [] + return + + def _load_sync() -> list[str]: + with self._engine.connect() as conn: # type: ignore[union-attr] + return self._crud.load(conn, self._run_id) + + self._entries = await self._run_sync(_load_sync) + + def is_replaying(self) -> bool: + """True if there are more journal entries to replay.""" + if self._entries is None: + return False + return self._replay_index < len(self._entries) + + def next_expected_key(self) -> str | None: + """Get the next expected task key during replay, or None if fresh execution.""" + if self._entries is None or self._replay_index >= len(self._entries): + return None + return self._entries[self._replay_index] + + async def record(self, key: str) -> None: + """Record a task completion and persist to database.""" + if self._entries is None: + self._entries = [] + + seq_num = len(self._entries) + self._entries.append(key) + self._replay_index += 1 + + if self._engine is not None: + + def _insert_sync() -> None: + with self._engine.begin() as conn: # type: ignore[union-attr] + self._crud.insert(conn, self._run_id, seq_num, key) + + await self._run_sync(_insert_sync) + + def advance(self) -> None: + """Advance replay index after processing a replayed task.""" + self._replay_index += 1 diff --git a/packages/llama-agents-dbos/src/llama_agents/dbos/runtime.py b/packages/llama-agents-dbos/src/llama_agents/dbos/runtime.py new file mode 100644 index 00000000..d6bd62a4 --- /dev/null +++ b/packages/llama-agents-dbos/src/llama_agents/dbos/runtime.py @@ -0,0 +1,677 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +""" +DBOS Runtime for durable workflow execution. + +This module provides the DBOSRuntime class for running LlamaIndex workflows +with durable execution backed by DBOS. +""" + +from __future__ import annotations + +import asyncio +import logging +import sys +import time +import weakref +from dataclasses import dataclass +from typing import Any, AsyncGenerator, TypedDict + +from llama_index_instrumentation.dispatcher import active_instrument_tags +from pydantic import BaseModel +from typing_extensions import Unpack +from workflows.context.serializers import BaseSerializer, JsonSerializer +from workflows.context.state_store import ( + StateStore, + deserialize_state_from_dict, + infer_state_type, +) +from workflows.events import Event, StartEvent, StopEvent +from workflows.runtime.types.internal_state import BrokerState +from workflows.runtime.types.named_task import NamedTask +from workflows.runtime.types.plugin import ( + ExternalRunAdapter, + InternalRunAdapter, + RegisteredWorkflow, + Runtime, + WaitResult, + WaitResultTick, + WaitResultTimeout, +) +from workflows.runtime.types.step_function import ( + StepWorkerFunction, + as_step_worker_functions, + create_workflow_run_function, +) +from workflows.runtime.types.ticks import WorkflowTick +from workflows.workflow import Workflow + +try: + from dbos import DBOS, SetWorkflowID + from dbos._dbos import _get_dbos_instance +except ImportError as e: + # if 3.9, give a detailed error that dbos is not supported on this version of python + if sys.version_info.major == 3 and sys.version_info.minor <= 9: + raise ImportError( + "dbos is not supported on Python 3.9. Please use Python 3.10 or higher." + f"Error: {e}" + ) from e + raise + +from sqlalchemy.engine import Engine + +from .journal.crud import JOURNAL_TABLE_NAME, JournalCrud +from .journal.task_journal import TaskJournal +from .state_store import STATE_TABLE_NAME, SqlStateStore + +logger = logging.getLogger(__name__) + + +class DBOSRuntimeConfig(TypedDict, total=False): + """Configuration options for DBOSRuntime. + + All fields are optional — defaults are resolved at launch time. + """ + + polling_interval_sec: float + run_migrations_on_launch: bool + schema: str | None + state_table_name: str + journal_table_name: str + + +DEFAULT_STATE_TABLE_NAME = STATE_TABLE_NAME +DEFAULT_JOURNAL_TABLE_NAME = JOURNAL_TABLE_NAME + + +def _resolve_schema(config: DBOSRuntimeConfig, engine: Engine) -> str | None: + """Resolve schema from config, falling back to dialect-based default. + + If "schema" was explicitly provided (even as None), uses that value. + Otherwise, defaults to "dbos" for PostgreSQL and None for SQLite. + """ + if "schema" in config: + return config["schema"] + is_postgres = engine.dialect.name == "postgresql" + return "dbos" if is_postgres else None + + +# Very long timeout for unbounded waits - encourages workflow to sleep. +# DBOS's default 60s is too short and gets recorded to event logs. +_UNBOUNDED_WAIT_TIMEOUT_SECONDS = 60 * 60 * 24 # 1 day + + +@dataclass +class _DBOSInternalShutdown: + """Internal signal sent via DBOS.send to wake blocked recv for shutdown.""" + + +@DBOS.step() +def _durable_time() -> float: + """ + Get current timestamp, wrapped as a DBOS step so that it's snapshotted and replayed + This could be made more consistent if it got the timestamp from the DB. + """ + return time.time() + + +class DBOSRuntime(Runtime): + """ + DBOS-backed workflow runtime for durable execution. + + Workflows are registered at launch() time with stable names, + enabling distributed workers and recovery. + + State is persisted to the database using SQL state stores, + enabling state recovery across process restarts. + """ + + def __init__(self, **kwargs: Unpack[DBOSRuntimeConfig]) -> None: + """Initialize the DBOS runtime. + + Args: + **kwargs: Configuration options. See DBOSRuntimeConfig for details. + polling_interval_sec: Interval for polling workflow results. Default 1.0. + run_migrations_on_launch: Auto-run migrations on launch(). Default True. + schema: Database schema name. Default: auto-detected at launch + ("dbos" for PostgreSQL, None for SQLite). Pass None explicitly + to force no schema even on PostgreSQL. + state_table_name: State table name. Default "workflow_state". + journal_table_name: Journal table name. Default "workflow_journal". + """ + self.config: DBOSRuntimeConfig = dict(kwargs) # type: ignore[assignment] + + # Workflow tracking state + self._pending: list[Workflow] = [] + self._pending_set: set[int] = set() # Track by id for dedup + self._registered: dict[int, RegisteredWorkflow] = {} # keyed by id(workflow) + + self._dbos_launched = False + self._tasks: list[asyncio.Task[None]] = [] + self._sql_engine: Engine | None = None + self._migrations_run = False + + def _track_task(self, task: asyncio.Task[None]) -> None: + self._tasks.append(task) + task.add_done_callback(self._tasks.remove) + + def track_workflow(self, workflow: Workflow) -> None: + """Track a workflow for registration at launch time. + + If launch() was already called, registers the workflow immediately. + This allows late registration for testing scenarios. + """ + if self._dbos_launched: + # Already launched - register immediately + registered = self.register(workflow) + self._registered[id(workflow)] = registered + else: + wf_id = id(workflow) + if wf_id not in self._pending_set: + self._pending.append(workflow) + self._pending_set.add(wf_id) + + def get_registered(self, workflow: Workflow) -> RegisteredWorkflow | None: + """Get the registered workflow if available.""" + return self._registered.get(id(workflow)) + + def register(self, workflow: Workflow) -> RegisteredWorkflow: + """ + Wrap workflow with DBOS decorators. + + Called at launch() time for each tracked workflow. + Uses workflow.workflow_name for stable DBOS registration names. + """ + # Use workflow's name directly + name = workflow.workflow_name + + # Create DBOS-wrapped control loop with stable name + @DBOS.workflow(name=f"{name}.control_loop") + async def _dbos_control_loop( + init_state: BrokerState, + start_event: StartEvent | None = None, + tags: dict[str, Any] = {}, + ) -> StopEvent: + workflow_run_fn = create_workflow_run_function(workflow) + return await workflow_run_fn(init_state, start_event, tags) + + # Wrap steps with stable names + wrapped_steps: dict[str, StepWorkerFunction[Any]] = { + step_name: DBOS.step(name=f"{name}.{step_name}")(step) + for step_name, step in as_step_worker_functions(workflow).items() + } + + return RegisteredWorkflow( + workflow=workflow, workflow_run_fn=_dbos_control_loop, steps=wrapped_steps + ) + + def _get_sql_engine(self) -> Engine: + """Get the SQLAlchemy engine from DBOS for state storage. + + Uses DBOS's app database if configured, otherwise falls back to sys database. + + Returns: + SQLAlchemy Engine for state storage. + + Raises: + RuntimeError: If no database is available. + """ + if self._sql_engine is not None: + return self._sql_engine + + dbos = _get_dbos_instance() + + # Try app database first, fall back to system database + app_db = dbos._app_db + if app_db is not None: + self._sql_engine = app_db.engine + return self._sql_engine + + # Fall back to system database + sys_db = dbos._sys_db + self._sql_engine = sys_db.engine + return self._sql_engine + + def run_migrations(self) -> None: + """Run database migrations for workflow state and journal tables. + + Creates the workflow_state and workflow_journal tables if they don't exist. + Idempotent - safe to call multiple times. + + Can be called explicitly before launch() when run_migrations_on_launch=False, + allowing for custom migration timing (e.g., during application startup). + + Requires DBOS to be launched first (calls _get_sql_engine internally). + """ + if self._migrations_run: + return + + engine = self._get_sql_engine() + schema = _resolve_schema(self.config, engine) + state_table = self.config.get("state_table_name", DEFAULT_STATE_TABLE_NAME) + journal_table = self.config.get( + "journal_table_name", DEFAULT_JOURNAL_TABLE_NAME + ) + + SqlStateStore.run_migrations(engine, table_name=state_table, schema=schema) + + # Create workflow_journal table + journal_crud = JournalCrud(table_name=journal_table, schema=schema) + journal_crud.run_migrations(engine) + + self._migrations_run = True + logger.info("Database migrations completed (workflow_state, workflow_journal)") + + def run_workflow( + self, + run_id: str, + workflow: Workflow, + init_state: BrokerState, + start_event: StartEvent | None = None, + serialized_state: dict[str, Any] | None = None, + serializer: BaseSerializer | None = None, + adapter_state: dict[str, Any] | None = None, + ) -> ExternalRunAdapter: + """Set up a workflow run with SQL-backed state storage. + + State is persisted to the database, enabling recovery across + process restarts and distributed execution. + + Args: + run_id: Unique identifier for this workflow run. + workflow: The workflow to run. + init_state: Initial broker state for the control loop. + start_event: Optional start event to kick off the workflow. + serialized_state: Optional pre-populated state from InMemoryStateStore.to_dict(). + If provided, this state is written to the database before the workflow + starts, allowing workflows to begin with pre-set initial values. + serializer: Serializer for state data. Defaults to JsonSerializer. + adapter_state: Optional adapter state (unused for DBOS). + """ + if not self._dbos_launched: + raise RuntimeError( + "DBOS runtime not launched. Call runtime.launch() before running workflows." + ) + + registered = self.get_registered(workflow) + if registered is None: + raise RuntimeError( + "DBOSRuntime workflows must be registered before running. Did you forget to call runtime.launch()?" + ) + + # Capture values needed in the async task closure + engine = self._get_sql_engine() + active_serializer = serializer or JsonSerializer() + + async def _run_workflow() -> None: + with SetWorkflowID(run_id): + # Write initial state to DB before starting workflow (non-blocking to caller) + if serialized_state: + store = SqlStateStore( + run_id=run_id, + engine=engine, + state_type=infer_state_type(workflow), + serializer=active_serializer, + schema=_resolve_schema(self.config, engine), + table_name=self.config.get( + "state_table_name", DEFAULT_STATE_TABLE_NAME + ), + ) + # Deserialize and save the initial state + state = deserialize_state_from_dict( + serialized_state, + active_serializer, + state_type=infer_state_type(workflow), + ) + await store.set_state(state) + + try: + await DBOS.start_workflow_async( + registered.workflow_run_fn, + init_state, + start_event, + active_instrument_tags.get(), + ) + except Exception as e: + logger.error( + f"Failed to submit work to DBOS for {run_id} with start event: {start_event} and init state: {init_state}. Error: {e}", + exc_info=True, + ) + raise e + + # Create startup task and pass to adapter so it can await workflow readiness + startup_task = asyncio.create_task(_run_workflow()) + self._track_task(startup_task) + + return ExternalDBOSAdapter( + run_id, + self.config.get("polling_interval_sec", 1.0), + startup_task, + ) + + def get_internal_adapter(self, workflow: Workflow) -> InternalRunAdapter: + if not self._dbos_launched: + raise RuntimeError( + "DBOS runtime not launched. Call runtime.launch() before running workflows." + ) + run_id = DBOS.workflow_id + if run_id is None: + raise RuntimeError( + "No current run id. Must be called within a workflow run." + ) + + # Infer state_type from the workflow for typed state support + state_type = infer_state_type(workflow) + + engine = self._get_sql_engine() + return InternalDBOSAdapter( + run_id, + engine, + state_type, + schema=_resolve_schema(self.config, engine), + state_table_name=self.config.get( + "state_table_name", DEFAULT_STATE_TABLE_NAME + ), + journal_table_name=self.config.get( + "journal_table_name", DEFAULT_JOURNAL_TABLE_NAME + ), + ) + + def get_external_adapter(self, run_id: str) -> ExternalRunAdapter: + if not self._dbos_launched: + raise RuntimeError( + "DBOS runtime not launched. Call runtime.launch() before running workflows." + ) + return ExternalDBOSAdapter(run_id, self.config.get("polling_interval_sec", 1.0)) + + def launch(self) -> None: + """ + Launch DBOS and register all tracked workflows. + + Must be called before running any workflows. + Runs database migrations unless run_migrations_on_launch=False. + """ + if self._dbos_launched: + return # Already launched + + # Register each pending workflow with DBOS + for workflow in self._pending: + # Register with DBOS (this applies decorators) + registered = self.register(workflow) + self._registered[id(workflow)] = registered + + # Launch DBOS runtime + DBOS.launch() + self._dbos_launched = True + + # Run migrations after DBOS is launched (if configured) + if self.config.get("run_migrations_on_launch", True): + self.run_migrations() + + def destroy(self, destroy_dbos: bool = True) -> None: + """Clean up DBOS runtime resources. + + Args: + destroy_dbos: If True (default), also calls DBOS.destroy(). + Set to False when DBOS lifecycle is managed externally + (e.g., shared across multiple runtimes in tests). + """ + self._pending.clear() + self._pending_set.clear() + self._registered.clear() + self._dbos_launched = False + self._sql_engine = None + self._migrations_run = False + for task in self._tasks: + if not task.done(): + task.cancel() + if destroy_dbos: + DBOS.destroy() + + +_IO_STREAM_PUBLISHED_EVENTS_NAME = "published_events" +_IO_STREAM_TICK_TOPIC = "ticks" + + +@dataclass +class _StreamWriteLock: + """Wrapper for asyncio.Lock to allow storage in WeakValueDictionary.""" + + lock: asyncio.Lock + + +class InternalDBOSAdapter(InternalRunAdapter): + """ + Internal DBOS adapter for the workflow control loop. + + - send_event sends ticks via DBOS.send (using run_in_executor to escape step context) + - wait_receive receives ticks via DBOS.recv_async + - write_to_event_stream publishes events via DBOS streams + - get_now returns a durable timestamp + - close sends shutdown signal to wake blocked recv + - wait_for_next_task coordinates task completion ordering for deterministic replay + """ + + # Class-level registry of stream write locks, keyed by run_id. + # Uses WeakValueDictionary so locks are GC'd when no adapters reference them. + # This is a workaround for DBOS race condition in write_stream_from_workflow + # where max(offset) read and insert are not atomic. + # See: https://github.com/dbos-inc/dbos-transact-py/issues/XXX + _stream_write_locks: weakref.WeakValueDictionary[str, _StreamWriteLock] = ( + weakref.WeakValueDictionary() + ) + + def __init__( + self, + run_id: str, + engine: Engine, + state_type: type[BaseModel] | None = None, + schema: str | None = None, + state_table_name: str = DEFAULT_STATE_TABLE_NAME, + journal_table_name: str = DEFAULT_JOURNAL_TABLE_NAME, + ) -> None: + self._run_id = run_id + self._engine = engine + self._state_type = state_type + self._schema = schema + self._state_table_name = state_table_name + self._journal_table_name = journal_table_name + self._closed = False + self._state_store: SqlStateStore[Any] | None = None + # Journal for deterministic task ordering - lazily initialized + self._journal: TaskJournal | None = None + # Get or create the shared lock for this run_id + self._stream_lock_holder = self._get_or_create_stream_lock(run_id) + + @classmethod + def _get_or_create_stream_lock(cls, run_id: str) -> _StreamWriteLock: + """Get existing lock for run_id or create a new one.""" + lock_holder = cls._stream_write_locks.get(run_id) + if lock_holder is None: + lock_holder = _StreamWriteLock(lock=asyncio.Lock()) + cls._stream_write_locks[run_id] = lock_holder + return lock_holder + + @property + def run_id(self) -> str: + return self._run_id + + async def write_to_event_stream(self, event: Event) -> None: + # Serialize stream writes to work around DBOS race condition where + # concurrent writes can read the same max(offset) and try to insert + # with the same offset, causing UNIQUE constraint violations. + async with self._stream_lock_holder.lock: + await DBOS.write_stream_async(_IO_STREAM_PUBLISHED_EVENTS_NAME, event) + + async def get_now(self) -> float: + return _durable_time() + + async def send_event(self, tick: WorkflowTick) -> None: + # Use run_in_executor to escape DBOS step context. + # DBOS yells at you for writing to the event stream from a step (since is not idempotent) + # However that's the expected semantics of llama index workflow steps, so it's ok. + loop = asyncio.get_running_loop() + await loop.run_in_executor( + None, + lambda: DBOS.send(self._run_id, tick, topic=_IO_STREAM_TICK_TOPIC), + ) + + async def wait_receive( + self, + timeout_seconds: float | None = None, + ) -> WaitResult: + """Wait for tick using DBOS.recv_async with timeout. + + For bounded waits (timeout_seconds specified), uses the specified timeout. + For unbounded waits (timeout_seconds=None), loops with very long timeouts + (hours) to encourage the workflow to sleep. Uses shutdown signal from + close() to exit cleanly - raises CancelledError on shutdown. + """ + if self._closed: + raise asyncio.CancelledError("Adapter closed") + + # Timeout 1x per day at least. This will just cause a wakeup loop of the control loop. + result = await DBOS.recv_async( + _IO_STREAM_TICK_TOPIC, + timeout_seconds=timeout_seconds or _UNBOUNDED_WAIT_TIMEOUT_SECONDS, + ) + if result is None: + return WaitResultTimeout() + if isinstance(result, _DBOSInternalShutdown): + self._closed = True + raise asyncio.CancelledError("Adapter closed") + return WaitResultTick(tick=result) + + async def close(self) -> None: + """Signal shutdown by sending internal message to wake blocked recv.""" + if self._closed: + return + self._closed = True + + loop = asyncio.get_running_loop() + await loop.run_in_executor( + None, + lambda: DBOS.send( + self._run_id, _DBOSInternalShutdown(), topic=_IO_STREAM_TICK_TOPIC + ), + ) + + def _get_or_create_state_store(self) -> SqlStateStore[Any]: + """Get or lazily create the state store.""" + if self._state_store is None: + self._state_store = SqlStateStore( + run_id=self._run_id, + engine=self._engine, + state_type=self._state_type, + schema=self._schema, + table_name=self._state_table_name, + ) + return self._state_store + + def get_state_store(self) -> StateStore[Any] | None: + return self._get_or_create_state_store() + + def _get_or_create_journal(self) -> TaskJournal: + """Get or lazily create the task journal.""" + if self._journal is None: + crud = JournalCrud(table_name=self._journal_table_name, schema=self._schema) + self._journal = TaskJournal(self._run_id, self._engine, crud) + return self._journal + + async def wait_for_next_task( + self, + task_set: list[NamedTask], + timeout: float | None = None, + ) -> asyncio.Task[Any] | None: + """Wait for and return the next task that should complete. + + During replay, waits for the specific task that completed in the original run. + During fresh execution, waits for any task and records the completion order. + + Args: + task_set: List of NamedTasks with stable string keys for identification + timeout: Timeout in seconds, None for no timeout + + Returns: + The completed task, or None on timeout. + """ + tasks = NamedTask.all_tasks(task_set) + if not tasks: + return None + + journal = self._get_or_create_journal() + await journal.load() + + expected_key = journal.next_expected_key() + if expected_key is not None: + # Replay mode: wait for specific task + target_task = NamedTask.find_by_key(task_set, expected_key) + + if target_task is None: + logger.warning( + f"Non-deterministic execution detected during replay! " + f"Expected task {expected_key} not in set yet. " + f"Falling back to awaiting all tasks." + ) + else: + await asyncio.wait_for(asyncio.shield(target_task), timeout=timeout) + journal.advance() + return target_task + + # Fresh execution: wait for first, record it + done, _ = await asyncio.wait( + tasks, timeout=timeout, return_when=asyncio.FIRST_COMPLETED + ) + if not done: + return None + + completed = done.pop() + key = NamedTask.get_key(task_set, completed) + await journal.record(key) + + return completed + + +class ExternalDBOSAdapter(ExternalRunAdapter): + """ + External DBOS adapter for workflow interaction. + + - send_event puts ticks into the shared mailbox queue + - stream_published_events reads from DBOS streams + - close is a no-op + """ + + def __init__( + self, + run_id: str, + polling_interval_sec: float = 1.0, + startup_task: asyncio.Task[None] | None = None, + ) -> None: + self._run_id = run_id + self._polling_interval_sec = polling_interval_sec + self._startup_task = startup_task # None means workflow already started + + @property + def run_id(self) -> str: + """Get the workflow run ID.""" + return self._run_id + + async def send_event(self, tick: WorkflowTick) -> None: + await DBOS.send_async(self._run_id, tick, topic=_IO_STREAM_TICK_TOPIC) + + async def stream_published_events(self) -> AsyncGenerator[Event, None]: + await self._ensure_workflow_started() + + async for event in DBOS.read_stream_async(self.run_id, "published_events"): + yield event + + async def get_result(self) -> StopEvent: + await self._ensure_workflow_started() + handle = await DBOS.retrieve_workflow_async(self.run_id) + return await handle.get_result(polling_interval_sec=self._polling_interval_sec) + + async def _ensure_workflow_started(self) -> None: + """Wait for the workflow startup task to complete if one was provided.""" + if self._startup_task is not None: + await self._startup_task + self._startup_task = None # Clear after awaiting diff --git a/packages/llama-agents-dbos/src/llama_agents/dbos/state_store.py b/packages/llama-agents-dbos/src/llama_agents/dbos/state_store.py new file mode 100644 index 00000000..379c8039 --- /dev/null +++ b/packages/llama-agents-dbos/src/llama_agents/dbos/state_store.py @@ -0,0 +1,566 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. + +""" +SQL-backed StateStore implementations for durable workflow state. + +Provides PostgreSQL and SQLite state stores that persist workflow state +to a database, enabling durable and distributed workflow execution. +""" + +from __future__ import annotations + +import asyncio +import functools +import json +import logging +from contextlib import asynccontextmanager +from datetime import datetime, timezone +from typing import ( + Any, + AsyncGenerator, + Callable, + Generic, + Literal, + Type, +) + +from pydantic import BaseModel, ConfigDict, Field, ValidationError +from sqlalchemy import ( + Column, + Connection, + DateTime, + MetaData, + String, + Table, + Text, + select, + text, +) +from sqlalchemy.dialects.postgresql import insert as pg_insert +from sqlalchemy.engine import Engine +from typing_extensions import TypeVar +from workflows.context.serializers import BaseSerializer, JsonSerializer +from workflows.context.state_store import ( + MAX_DEPTH, + DictState, + InMemorySerializedState, + assign_path_step, + deserialize_state_from_dict, + traverse_path_step, +) + +logger = logging.getLogger(__name__) + +MODEL_T = TypeVar("MODEL_T", bound=BaseModel) + + +class SqlSerializedState(BaseModel): + """Serialized state referencing a database row (from SqlStateStore).""" + + model_config = ConfigDict(populate_by_name=True) + + store_type: Literal["sql"] = "sql" + run_id: str + db_schema: str | None = Field(default=None, alias="schema") + # Note: No state_data - actual data lives in the database + + +def parse_serialized_state( + data: dict[str, Any], +) -> InMemorySerializedState | SqlSerializedState: + """Parse raw dict into appropriate format type. + + Args: + data: Serialized state payload from to_dict(). + + Returns: + InMemorySerializedState or SqlSerializedState based on store_type. + + Raises: + ValueError: If store_type is unknown. + """ + store_type = data.get("store_type") + + if store_type == "sql": + return SqlSerializedState.model_validate(data) + elif store_type == "in_memory" or store_type is None: + # Backwards compat: missing store_type = InMemory + return InMemorySerializedState.model_validate(data) + else: + raise ValueError(f"Unknown store_type: {store_type}") + + +def _utc_now() -> datetime: + """Get current UTC timestamp.""" + return datetime.now(timezone.utc) + + +STATE_TABLE_NAME = "workflow_state" + + +def _state_columns() -> list[Column]: + """Return fresh Column instances for the workflow_state table. + + Must return new instances each call because SQLAlchemy Column objects + can't be shared across Table instances. + """ + return [ + Column("run_id", String(255), primary_key=True), + Column("state_json", Text, nullable=False), + Column("created_at", DateTime(timezone=True), nullable=False), + Column("updated_at", DateTime(timezone=True), nullable=False), + ] + + +class SqlStateStore(Generic[MODEL_T]): + """ + SQL-backed StateStore implementation. + + Persists workflow state to a database table. Supports PostgreSQL and SQLite + dialects with automatic detection based on the engine. + + Thread-safety is achieved through database-level locking during + transactional edits via the `edit_state` context manager. + """ + + known_unserializable_keys = ("memory",) + state_type: Type[MODEL_T] + + def __init__( + self, + run_id: str, + state_type: Type[MODEL_T] | None = None, + engine: Engine | None = None, + serializer: BaseSerializer | None = None, + schema: str | None = None, + table_name: str = STATE_TABLE_NAME, + ) -> None: + self._run_id = run_id + self.state_type = state_type or DictState # type: ignore[assignment] + self._engine = engine + self._serializer = serializer or JsonSerializer() + self._schema = schema + self._table_name = table_name + self._metadata = MetaData(schema=self._schema) + self._table = self._define_table() + self._initialized = False + self._pending_state: dict[str, Any] | None = None + + @property + def run_id(self) -> str: + """Get the workflow run ID.""" + return self._run_id + + @property + def engine(self) -> Engine: + """Get the SQLAlchemy engine, raising if not set.""" + if self._engine is None: + raise RuntimeError( + "Engine not set. Provide an engine at construction or set it before use." + ) + return self._engine + + @engine.setter + def engine(self, engine: Engine) -> None: + """Set the SQLAlchemy engine.""" + self._engine = engine + + @property + def _is_postgres(self) -> bool: + """Check if the engine is PostgreSQL.""" + return self.engine.dialect.name == "postgresql" + + @property + def _table_ref(self) -> str: + """Get the fully qualified table reference.""" + if self._schema: + return f"{self._schema}.{self._table_name}" + return self._table_name + + def _define_table(self) -> Table: + """Define the workflow_state table schema.""" + return Table( + self._table_name, + self._metadata, + *_state_columns(), + ) + + @classmethod + def run_migrations( + cls, + engine: Engine, + table_name: str = STATE_TABLE_NAME, + schema: str | None = None, + ) -> None: + """Create schema and table if they don't exist.""" + metadata = MetaData(schema=schema) + table = Table( + table_name, + metadata, + *_state_columns(), + ) + is_postgres = engine.dialect.name == "postgresql" + with engine.begin() as conn: + if is_postgres and schema: + conn.execute(text(f"CREATE SCHEMA IF NOT EXISTS {schema}")) # noqa: S608 + table.create(bind=conn, checkfirst=True) + + @functools.cached_property + def _lock(self) -> asyncio.Lock: + """Lazy lock for Python 3.14+ compatibility.""" + return asyncio.Lock() + + async def _run_sync(self, fn: Callable[..., Any], *args: Any) -> Any: + """Run a synchronous function in the default executor.""" + loop = asyncio.get_event_loop() + return await loop.run_in_executor(None, fn, *args) + + def _ensure_initialized(self) -> None: + """Ensure the table exists and apply any pending state.""" + if self._initialized: + return + self._run_instance_migrations() + self._initialized = True + + # Apply any pending state from InMemory format deserialization + if self._pending_state is not None: + self._apply_pending_state_sync() + + def _apply_pending_state_sync(self) -> None: + """Write pending state to database (called after migrations).""" + if self._pending_state is None: + return + + serialized_state = self._pending_state + self._pending_state = None + + state = deserialize_state_from_dict( + serialized_state, self._serializer, state_type=self.state_type + ) + state_json = self._serialize_state(state) # type: ignore[arg-type] + + with self.engine.begin() as conn: + self._upsert_state(conn, state_json, _utc_now()) + + def _run_instance_migrations(self) -> None: + """Create schema and table if they don't exist (instance-level).""" + with self.engine.begin() as conn: + if self._is_postgres and self._schema: + conn.execute(text(f"CREATE SCHEMA IF NOT EXISTS {self._schema}")) # noqa: S608 + self._table.create(bind=conn, checkfirst=True) + + def _lock_row_for_update(self, conn: Connection) -> dict[str, Any] | None: + """Lock and return row data for this run_id.""" + for_update = "FOR UPDATE" if self._is_postgres else "" + result = conn.execute( + text(f""" + SELECT state_json + FROM {self._table_ref} + WHERE run_id = :run_id + {for_update} + """), # noqa: S608 + {"run_id": self._run_id}, + ) + row = result.fetchone() + if row is None: + return None + return {"state_json": row[0]} + + def _upsert_state( + self, + conn: Connection, + state_json: str, + now: datetime, + ) -> None: + """Perform database-specific upsert operation.""" + if self._is_postgres: + stmt = pg_insert(self._table).values( + run_id=self._run_id, + state_json=state_json, + created_at=now, + updated_at=now, + ) + stmt = stmt.on_conflict_do_update( + index_elements=["run_id"], + set_={ + "state_json": stmt.excluded.state_json, + "updated_at": stmt.excluded.updated_at, + }, + ) + conn.execute(stmt) + else: + # SQLite upsert + conn.execute( + text(f""" + INSERT INTO {self._table_ref} + (run_id, state_json, created_at, updated_at) + VALUES (:run_id, :state_json, :created_at, :updated_at) + ON CONFLICT (run_id) DO UPDATE SET + state_json = excluded.state_json, + updated_at = excluded.updated_at + """), # noqa: S608 + { + "run_id": self._run_id, + "state_json": state_json, + "created_at": now.isoformat(), + "updated_at": now.isoformat(), + }, + ) + + def _serialize_state(self, state: MODEL_T) -> str: + """Serialize state model to JSON string.""" + if isinstance(state, DictState): + serialized_data: dict[str, Any] = {} + for key, value in state.items(): + try: + serialized_data[key] = self._serializer.serialize(value) + except Exception: + if key in self.known_unserializable_keys: + logger.warning(f"Skipping unserializable key: {key}") + continue + raise + return json.dumps({"_data": serialized_data}) + return self._serializer.serialize(state) + + def _deserialize_state(self, state_json: str) -> MODEL_T: + """Deserialize state from JSON string.""" + if issubclass(self.state_type, DictState): + data = json.loads(state_json) + deserialized = { + k: self._serializer.deserialize(v) + for k, v in data.get("_data", {}).items() + } + return DictState(_data=deserialized) # type: ignore[return-value] + return self._serializer.deserialize(state_json) + + def _create_default_state(self) -> MODEL_T: + """Create a default instance of the state type.""" + return self.state_type() + + def _load_state_sync(self) -> MODEL_T: + """Load state from database synchronously.""" + self._ensure_initialized() + with self.engine.connect() as conn: + result = conn.execute( + select(self._table.c.state_json).where( + self._table.c.run_id == self._run_id + ) + ) + row = result.fetchone() + if row is None: + state = self._create_default_state() + self._save_state_sync(state, conn) + conn.commit() + return state + return self._deserialize_state(row[0]) + + def _save_state_sync(self, state: MODEL_T, conn: Connection) -> None: + """Save state to database synchronously.""" + now = _utc_now() + self._upsert_state(conn, self._serialize_state(state), now) + + async def get_state(self) -> MODEL_T: + """Return a copy of the current state model.""" + state = await self._run_sync(self._load_state_sync) + return state.model_copy() + + async def set_state(self, state: MODEL_T) -> None: + """Replace or merge into the current state model.""" + + def _set_state_sync() -> None: + self._ensure_initialized() + with self.engine.begin() as conn: + result = conn.execute( + select(self._table.c.state_json).where( + self._table.c.run_id == self._run_id + ) + ) + row = result.fetchone() + + if row is None: + self._save_state_sync(state, conn) + return + + current_state = self._deserialize_state(row[0]) + current_type = type(current_state) + new_type = type(state) + + if isinstance(state, current_type): + self._save_state_sync(state, conn) + elif issubclass(current_type, new_type): + parent_data = state.model_dump() + merged = current_type.model_validate( + {**current_state.model_dump(), **parent_data} + ) + self._save_state_sync(merged, conn) + else: + raise ValueError( + f"State must be of type {current_type.__name__} or parent, " + f"got {new_type.__name__}" + ) + + await self._run_sync(_set_state_sync) + + async def get(self, path: str, default: Any = ...) -> Any: + """Get a nested value using dot-separated paths.""" + state = await self._run_sync(self._load_state_sync) + segments = path.split(".") if path else [] + if len(segments) > MAX_DEPTH: + raise ValueError(f"Path length exceeds {MAX_DEPTH} segments") + + try: + value: Any = state + for segment in segments: + value = traverse_path_step(value, segment) + except Exception: + if default is not ...: + return default + raise ValueError(f"Path '{path}' not found in state") + return value + + async def set(self, path: str, value: Any) -> None: + """Set a nested value using dot-separated paths.""" + if not path: + raise ValueError("Path cannot be empty") + segments = path.split(".") + if len(segments) > MAX_DEPTH: + raise ValueError(f"Path length exceeds {MAX_DEPTH} segments") + + async with self.edit_state() as state: + current: Any = state + for segment in segments[:-1]: + try: + current = traverse_path_step(current, segment) + except (KeyError, AttributeError, IndexError, TypeError): + intermediate: Any = {} + assign_path_step(current, segment, intermediate) + current = intermediate + assign_path_step(current, segments[-1], value) + + async def clear(self) -> None: + """Reset the state to its type defaults.""" + try: + await self.set_state(self._create_default_state()) + except ValidationError: + raise ValueError("State must have defaults for all fields") + + @asynccontextmanager + async def edit_state(self) -> AsyncGenerator[MODEL_T, None]: + """Edit state transactionally under a database lock.""" + + def _edit_with_lock() -> tuple[MODEL_T, Callable[[MODEL_T], None]]: + self._ensure_initialized() + conn = self.engine.connect() + trans = conn.begin() + + try: + row_data = self._lock_row_for_update(conn) + if row_data is None: + state = self._create_default_state() + else: + state = self._deserialize_state(row_data["state_json"]) + + def commit_fn(updated_state: MODEL_T) -> None: + try: + self._save_state_sync(updated_state, conn) + trans.commit() + finally: + conn.close() + + return state, commit_fn + except Exception: + trans.rollback() + conn.close() + raise + + async with self._lock: + state, commit_fn = await self._run_sync(_edit_with_lock) + try: + yield state + await self._run_sync(commit_fn, state) + except Exception: + raise + + def to_dict(self, serializer: BaseSerializer) -> dict[str, Any]: + """Serialize state store metadata for persistence. + + Returns a SqlSerializedState payload that can be restored by from_dict(). + The actual state data lives in the database, so this only serializes + connection metadata (run_id, schema). + """ + payload = SqlSerializedState.model_validate( + { + "run_id": self._run_id, + "schema": self._schema, + } + ) + return payload.model_dump(by_alias=True) + + @classmethod + def from_dict( + cls, + serialized_state: dict[str, Any], + serializer: BaseSerializer, + state_type: type[BaseModel] = DictState, + run_id: str | None = None, + ) -> SqlStateStore[Any]: + """Restore a state store from serialized payload. + + Handles both InMemorySerializedState and SqlSerializedState formats: + + - InMemorySerializedState: Stores the serialized data internally and + writes it to the database when the engine is first used (via + _ensure_initialized). This enables restoring state from in-memory + format into a SQL-backed store. + + - SqlSerializedState: Creates a store pointing at the existing database + row. If a different run_id is provided, the store will use that run_id + (data copying must be handled separately if needed). + + Note: The engine must be set separately after restoration. + + Args: + serialized_state: Payload from to_dict() of either store type. + serializer: Serializer for data handling. + state_type: The state model type for deserialization. + run_id: Optional override run_id. If not provided, uses the run_id + from SqlSerializedState or generates one for InMemory format. + + Returns: + A new SqlStateStore instance configured from the payload. + + Raises: + ValueError: If serialized_state is empty. + """ + import uuid + + if not serialized_state: + raise ValueError("Cannot restore SqlStateStore from empty dict") + + parsed = parse_serialized_state(serialized_state) + + if isinstance(parsed, InMemorySerializedState): + # InMemory format: store data internally, apply when engine is set + effective_run_id = run_id or str(uuid.uuid4()) + + store = cls( + run_id=effective_run_id, + state_type=state_type, # type: ignore[arg-type] + serializer=serializer, + ) + # Store the serialized state to apply when engine is available + store._pending_state = serialized_state + return store + + else: + # SqlSerializedState format: create store pointing at existing row + effective_run_id = run_id or parsed.run_id + schema = parsed.db_schema + + return cls( + run_id=effective_run_id, + state_type=state_type, # type: ignore[arg-type] + serializer=serializer, + schema=schema, + ) diff --git a/packages/llama-agents-dbos/tests/fixtures/__init__.py b/packages/llama-agents-dbos/tests/fixtures/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/packages/llama-agents-dbos/tests/fixtures/runner.py b/packages/llama-agents-dbos/tests/fixtures/runner.py new file mode 100644 index 00000000..467190f3 --- /dev/null +++ b/packages/llama-agents-dbos/tests/fixtures/runner.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""Subprocess runner for workflow tests with DBOS isolation. + +This module provides a CLI runner for executing workflows in isolated +subprocesses, supporting interrupt/resume testing and human-in-the-loop +response simulation. + +Usage: + python /path/to/packages/llama-agents-dbos/tests/fixtures/runner.py \ + --workflow "tests.fixtures.workflows.hitl:TestWorkflow" \ + --db-url "sqlite+pysqlite:///path/to/db" \ + --run-id "test-001" \ + --config '{"interrupt_on": "AskInputEvent"}' + +Config modes: + - interrupt_on: Interrupt when event type is seen (uses os._exit(0)) + - String form: "EventName" - interrupt on any instance of EventName + - Dict form: {"event": "EventName", "condition": {"field": value}} + - interrupt only when type matches AND all condition fields match + - respond: Respond to InputRequiredEvent subtypes with specified events + - run-to-completion: Empty config or omit both fields +""" + +from __future__ import annotations + +import argparse +import asyncio +import importlib +import json +import os +import sys +from pathlib import Path +from types import ModuleType +from typing import Any + +# Add package source directories to sys.path for imports +# Runner is at: packages/llama-agents-dbos/tests/fixtures/runner.py +# We need to add: +# - packages/llama-agents-dbos/src for llama_agents.dbos +# - packages/llama-index-workflows/src for workflows.* +# - packages/llama-agents-dbos (parent of tests/) so tests.fixtures.workflows.* can be imported +TESTS_DIR = Path(__file__).parent.parent +DBOS_PACKAGE_DIR = TESTS_DIR.parent +DBOS_PACKAGE_SRC_PATH = str(DBOS_PACKAGE_DIR / "src") +WORKFLOWS_PACKAGE_SRC_PATH = str( + DBOS_PACKAGE_DIR.parent / "llama-index-workflows" / "src" +) + +# Insert at front of path so these packages take precedence +# Add the parent of tests/ so "import tests.fixtures.workflows..." works +sys.path.insert(0, str(DBOS_PACKAGE_DIR)) +sys.path.insert(0, DBOS_PACKAGE_SRC_PATH) +sys.path.insert(0, WORKFLOWS_PACKAGE_SRC_PATH) + +from dbos import DBOS, DBOSConfig # noqa: E402 +from llama_agents.dbos import DBOSRuntime # noqa: E402 +from workflows.context import Context # noqa: E402 +from workflows.events import Event, InputRequiredEvent, StartEvent # noqa: E402 +from workflows.workflow import Workflow # noqa: E402 + + +def import_workflow(path: str) -> tuple[type[Workflow], ModuleType]: + """Import a workflow class from a module path. + + Args: + path: Module path with class name, e.g., "tests.fixtures.workflows.hitl:TestWorkflow" + + Returns: + Tuple of (workflow_class, module) for accessing classes defined in the module. + + Raises: + ValueError: If path format is invalid. + ImportError: If module cannot be imported. + AttributeError: If class not found in module. + """ + if ":" not in path: + raise ValueError( + f"Invalid workflow path format: {path}. Expected 'module.path:ClassName'" + ) + module_path, class_name = path.rsplit(":", 1) + module = importlib.import_module(module_path) + workflow_class = getattr(module, class_name) + if not (isinstance(workflow_class, type) and issubclass(workflow_class, Workflow)): + raise TypeError(f"{class_name} is not a Workflow subclass") + return workflow_class, module + + +def get_event_class_by_name(module: ModuleType, name: str) -> type[Event] | None: + """Find an event class in a module by its name. + + Searches through all attributes of the module to find an Event subclass + with a matching class name. + + Args: + module: The module to search in. + name: The class name to find. + + Returns: + The event class if found, None otherwise. + """ + for attr_name in dir(module): + attr = getattr(module, attr_name) + if isinstance(attr, type) and issubclass(attr, Event) and attr.__name__ == name: + return attr + return None + + +def parse_config(config_json: str | None) -> dict[str, Any]: + """Parse the JSON config string. + + Args: + config_json: JSON string with configuration, or None. + + Returns: + Parsed config dict, or empty dict if None. + """ + if not config_json: + return {} + return json.loads(config_json) + + +def setup_dbos(db_url: str, app_name: str = "test-workflow") -> DBOSRuntime: + """Set up DBOS with the given database URL. + + Args: + db_url: SQLite database URL. + app_name: Application name for DBOS config. + + Returns: + Configured DBOSRuntime instance. + """ + config: DBOSConfig = { + "name": app_name, + "system_database_url": db_url, + "run_admin_server": False, + } + DBOS(config=config) + return DBOSRuntime(polling_interval_sec=0.01) + + +async def run_workflow( + workflow_path: str, + db_url: str, + run_id: str, + config: dict[str, Any], +) -> None: + """Run the workflow with the specified configuration. + + Args: + workflow_path: Module path with class name. + db_url: SQLite database URL. + run_id: Unique run ID for the workflow. + config: Configuration dict with interrupt_on and/or respond settings. + """ + # Import workflow and get module for event class lookup + workflow_class, module = import_workflow(workflow_path) + + # Parse config options + interrupt_on_config = config.get("interrupt_on") + respond_config = config.get("respond", {}) + + # Resolve interrupt config (can be string or dict with condition) + interrupt_event_class: type[Event] | None = None + interrupt_condition: dict[str, Any] | None = None + if interrupt_on_config: + if isinstance(interrupt_on_config, str): + interrupt_event_name = interrupt_on_config + else: + interrupt_event_name = interrupt_on_config.get("event") + interrupt_condition = interrupt_on_config.get("condition") + interrupt_event_class = get_event_class_by_name(module, interrupt_event_name) + if interrupt_event_class is None: + print( + f"ERROR:ValueError:Event class '{interrupt_event_name}' not found in module" + ) + sys.exit(1) + + # Build response event mapping: {trigger_class: (response_class, fields)} + response_map: dict[type[Event], tuple[type[Event], dict[str, Any]]] = {} + for trigger_name, response_info in respond_config.items(): + trigger_class = get_event_class_by_name(module, trigger_name) + if trigger_class is None: + print( + f"ERROR:ValueError:Trigger event class '{trigger_name}' not found in module" + ) + sys.exit(1) + response_event_name = response_info.get("event") + response_fields = response_info.get("fields", {}) + response_class = get_event_class_by_name(module, response_event_name) + if response_class is None: + print( + f"ERROR:ValueError:Response event class '{response_event_name}' not found in module" + ) + sys.exit(1) + # Both trigger_class and response_class are narrowed after sys.exit(1) guards + assert trigger_class is not None + assert response_class is not None + response_map[trigger_class] = (response_class, response_fields) + + # Set up DBOS and runtime + runtime = setup_dbos(db_url) + + # Create workflow instance and launch + wf = workflow_class(runtime=runtime) + runtime.launch() + + try: + ctx = Context(wf) + handler = ctx._workflow_run(wf, StartEvent(), run_id=run_id) + + async for event in handler.stream_events(): + event_name = type(event).__name__ + print(f"EVENT:{event_name}", flush=True) + + # Check for interrupt condition + if interrupt_event_class is not None and isinstance( + event, interrupt_event_class + ): + # Check condition fields if present + should_interrupt = True + if interrupt_condition: + for field, expected_value in interrupt_condition.items(): + actual_value = getattr(event, field, None) + if actual_value != expected_value: + should_interrupt = False + break + if should_interrupt: + print("INTERRUPTING", flush=True) + os._exit(0) + + # Check for response condition (InputRequiredEvent subtypes) + if isinstance(event, InputRequiredEvent): + for trigger_class, (response_class, fields) in response_map.items(): + if isinstance(event, trigger_class): + if handler.ctx: + response_event = response_class(**fields) + handler.ctx.send_event(response_event) + break + + result = await handler + print(f"RESULT:{result}", flush=True) + print("SUCCESS", flush=True) + + except Exception as e: + print(f"ERROR:{type(e).__name__}:{e}", flush=True) + raise + + finally: + runtime.destroy() + + +def main() -> None: + """Entry point for the subprocess runner.""" + parser = argparse.ArgumentParser( + description="Run workflows in isolated subprocesses for testing" + ) + parser.add_argument( + "--workflow", + required=True, + help="Module path with class name (e.g., 'tests.fixtures.workflows.hitl:TestWorkflow')", + ) + parser.add_argument( + "--db-url", + required=True, + help="SQLite database URL", + ) + parser.add_argument( + "--run-id", + required=True, + help="Unique run ID for the workflow", + ) + parser.add_argument( + "--config", + default=None, + help="JSON string with configuration", + ) + + args = parser.parse_args() + + config = parse_config(args.config) + + asyncio.run( + run_workflow( + workflow_path=args.workflow, + db_url=args.db_url, + run_id=args.run_id, + config=config, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/packages/llama-agents-dbos/tests/fixtures/workflows/__init__.py b/packages/llama-agents-dbos/tests/fixtures/workflows/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/packages/llama-agents-dbos/tests/fixtures/workflows/chained.py b/packages/llama-agents-dbos/tests/fixtures/workflows/chained.py new file mode 100644 index 00000000..5ae0455b --- /dev/null +++ b/packages/llama-agents-dbos/tests/fixtures/workflows/chained.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""Chained workflow fixture with StepOneEvent, StepTwoEvent, and ChainedWorkflow.""" + +from __future__ import annotations + +from pydantic import Field +from workflows.context import Context +from workflows.decorators import step +from workflows.events import Event, StartEvent, StopEvent +from workflows.workflow import Workflow + + +class StepOneEvent(Event): + value: str = Field(default="one") + + +class StepTwoEvent(Event): + value: str = Field(default="two") + + +class ChainedWorkflow(Workflow): + @step + async def step_one(self, ctx: Context, ev: StartEvent) -> StepOneEvent: + await ctx.store.set("step_one", True) + print("STEP:one:complete", flush=True) + return StepOneEvent() + + @step + async def step_two(self, ctx: Context, ev: StepOneEvent) -> StepTwoEvent: + await ctx.store.set("step_two", True) + print("STEP:two:complete", flush=True) + return StepTwoEvent() + + @step + async def step_three(self, ctx: Context, ev: StepTwoEvent) -> StopEvent: + await ctx.store.set("step_three", True) + print("STEP:three:complete", flush=True) + return StopEvent(result="done") diff --git a/packages/llama-agents-dbos/tests/fixtures/workflows/concurrent_workers.py b/packages/llama-agents-dbos/tests/fixtures/workflows/concurrent_workers.py new file mode 100644 index 00000000..cdbb23f9 --- /dev/null +++ b/packages/llama-agents-dbos/tests/fixtures/workflows/concurrent_workers.py @@ -0,0 +1,45 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""Concurrent workers workflow fixture with num_workers=2.""" + +from __future__ import annotations + +import asyncio +import random + +from pydantic import Field +from workflows.context import Context +from workflows.decorators import step +from workflows.events import Event, StartEvent, StopEvent +from workflows.workflow import Workflow + + +class WorkItem(Event): + item_id: int = Field(default=0) + + +class WorkDone(Event): + item_id: int = Field(default=0) + + +class ConcurrentWorkersWorkflow(Workflow): + @step + async def dispatch(self, ctx: Context, ev: StartEvent) -> WorkItem: + # Dispatch work items that will be processed by concurrent workers + ctx.send_event(WorkItem(item_id=1)) + ctx.send_event(WorkItem(item_id=2)) + print("STEP:dispatch:complete", flush=True) + return WorkItem(item_id=0) + + @step(num_workers=2) + async def worker(self, ctx: Context, ev: WorkItem) -> WorkDone: + # Variable processing time for each item + await asyncio.sleep(random.uniform(0.01, 0.05)) + print(f"STEP:worker:{ev.item_id}:complete", flush=True) + return WorkDone(item_id=ev.item_id) + + @step + async def finish(self, ctx: Context, ev: WorkDone) -> StopEvent: + # First WorkDone to arrive ends the workflow + print(f"STEP:finish:{ev.item_id}:complete", flush=True) + return StopEvent(result={"first_done": ev.item_id}) diff --git a/packages/llama-agents-dbos/tests/fixtures/workflows/hitl.py b/packages/llama-agents-dbos/tests/fixtures/workflows/hitl.py new file mode 100644 index 00000000..2a921ae5 --- /dev/null +++ b/packages/llama-agents-dbos/tests/fixtures/workflows/hitl.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""Basic HITL workflow fixture with AskInputEvent and UserInput events.""" + +from __future__ import annotations + +from pydantic import Field +from workflows.context import Context +from workflows.decorators import step +from workflows.events import Event, InputRequiredEvent, StartEvent, StopEvent +from workflows.workflow import Workflow + + +class AskInputEvent(InputRequiredEvent): + prefix: str = Field(default="Enter: ") + + +class UserInput(Event): + response: str = Field(default="") + + +class TestWorkflow(Workflow): + @step + async def ask(self, ctx: Context, ev: StartEvent) -> AskInputEvent: + await ctx.store.set("asked", True) + print("STEP:ask:complete", flush=True) + return AskInputEvent() + + @step + async def process(self, ctx: Context, ev: UserInput) -> StopEvent: + await ctx.store.set("processed", ev.response) + print("STEP:process:complete", flush=True) + return StopEvent(result={"response": ev.response}) diff --git a/packages/llama-agents-dbos/tests/fixtures/workflows/parallel.py b/packages/llama-agents-dbos/tests/fixtures/workflows/parallel.py new file mode 100644 index 00000000..6b498d8a --- /dev/null +++ b/packages/llama-agents-dbos/tests/fixtures/workflows/parallel.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""Parallel workflow fixture with ResultAEvent, ResultBEvent, and ParallelWorkflow.""" + +from __future__ import annotations + +import asyncio +import random + +from pydantic import Field +from workflows.context import Context +from workflows.decorators import step +from workflows.events import Event, StartEvent, StopEvent +from workflows.workflow import Workflow + + +class ResultAEvent(Event): + value: str = Field(default="") + + +class ResultBEvent(Event): + value: str = Field(default="") + + +class ParallelWorkflow(Workflow): + @step + async def branch_a(self, ctx: Context, ev: StartEvent) -> ResultAEvent: + # Variable processing time - may complete before or after branch_b + await asyncio.sleep(random.uniform(0.01, 0.05)) + print("STEP:branch_a:complete", flush=True) + return ResultAEvent(value="a_result") + + @step + async def branch_b(self, ctx: Context, ev: StartEvent) -> ResultBEvent: + # Variable processing time - may complete before or after branch_a + await asyncio.sleep(random.uniform(0.01, 0.05)) + print("STEP:branch_b:complete", flush=True) + return ResultBEvent(value="b_result") + + @step + async def finish_a(self, ctx: Context, ev: ResultAEvent) -> StopEvent: + print("STEP:finish_a:complete", flush=True) + return StopEvent(result={"winner": "a", "value": ev.value}) + + @step + async def finish_b(self, ctx: Context, ev: ResultBEvent) -> StopEvent: + print("STEP:finish_b:complete", flush=True) + return StopEvent(result={"winner": "b", "value": ev.value}) diff --git a/packages/llama-agents-dbos/tests/fixtures/workflows/sequential_hitl.py b/packages/llama-agents-dbos/tests/fixtures/workflows/sequential_hitl.py new file mode 100644 index 00000000..8dada351 --- /dev/null +++ b/packages/llama-agents-dbos/tests/fixtures/workflows/sequential_hitl.py @@ -0,0 +1,44 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""Sequential HITL workflow fixture with ProcessedEvent and WaitForInputEvent.""" + +from __future__ import annotations + +import asyncio +import random + +from pydantic import Field +from workflows.context import Context +from workflows.decorators import step +from workflows.events import Event, InputRequiredEvent, StartEvent, StopEvent +from workflows.workflow import Workflow + + +class ProcessedEvent(Event): + value: str = Field(default="") + + +class WaitForInputEvent(InputRequiredEvent): + prompt: str = Field(default="") + + +class UserContinueEvent(Event): + continue_value: str = Field(default="") + + +class SequentialHITLWorkflow(Workflow): + @step + async def process(self, ctx: Context, ev: StartEvent) -> ProcessedEvent: + await asyncio.sleep(random.uniform(0.01, 0.05)) + print("STEP:process:complete", flush=True) + return ProcessedEvent(value="processed") + + @step + async def ask_user(self, ctx: Context, ev: ProcessedEvent) -> WaitForInputEvent: + print("STEP:ask_user:triggering_wait", flush=True) + return WaitForInputEvent(prompt=f"Got {ev.value}") + + @step + async def finalize(self, ctx: Context, ev: UserContinueEvent) -> StopEvent: + print(f"STEP:finalize:complete:{ev.continue_value}", flush=True) + return StopEvent(result={"continue": ev.continue_value}) diff --git a/packages/llama-agents-dbos/tests/fixtures/workflows/streaming_stress.py b/packages/llama-agents-dbos/tests/fixtures/workflows/streaming_stress.py new file mode 100644 index 00000000..1c503be7 --- /dev/null +++ b/packages/llama-agents-dbos/tests/fixtures/workflows/streaming_stress.py @@ -0,0 +1,64 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""Streaming stress workflow fixture with many concurrent stream writes.""" + +from __future__ import annotations + +import asyncio + +from pydantic import Field +from workflows.context import Context +from workflows.decorators import step +from workflows.events import Event, StartEvent, StopEvent +from workflows.workflow import Workflow + + +class ProgressEvent(Event): + progress: int = Field(default=0) + + +class WorkItem(Event): + item_id: int = Field(default=0) + + +class WorkDone(Event): + item_id: int = Field(default=0) + total_processed: int = Field(default=0) + + +class FanOutComplete(Event): + pass + + +class StreamingStressWorkflow(Workflow): + @step + async def fan_out(self, ctx: Context, ev: StartEvent) -> FanOutComplete: + # Fire many stream writes and internal events concurrently + # This creates many background tasks that call DBOS operations + for i in range(15): + ctx.write_event_to_stream(ProgressEvent(progress=i)) + ctx.send_event(WorkItem(item_id=i)) + print("STEP:fan_out:dispatched_15_items", flush=True) + # Write completion signal to stream for interrupt tests + ctx.write_event_to_stream(ProgressEvent(progress=999)) + return FanOutComplete() + + @step(num_workers=4) + async def process_work(self, ctx: Context, ev: WorkItem) -> WorkDone: + # Each worker also writes to stream, creating more concurrent DBOS ops + await asyncio.sleep(0.01) # Small delay to increase interleaving + ctx.write_event_to_stream(ProgressEvent(progress=100 + ev.item_id)) + print(f"STEP:process_work:{ev.item_id}:complete", flush=True) + return WorkDone(item_id=ev.item_id) + + @step + async def after_fanout(self, ctx: Context, ev: FanOutComplete) -> None: + # Consume FanOutComplete, don't trigger anything + print("STEP:after_fanout:complete", flush=True) + return None + + @step + async def collect(self, ctx: Context, ev: WorkDone) -> StopEvent: + # First WorkDone ends the workflow + print(f"STEP:collect:{ev.item_id}:complete", flush=True) + return StopEvent(result={"first_done": ev.item_id}) diff --git a/packages/llama-agents-dbos/tests/fixtures/workflows/three_step_hitl.py b/packages/llama-agents-dbos/tests/fixtures/workflows/three_step_hitl.py new file mode 100644 index 00000000..54d0e5fc --- /dev/null +++ b/packages/llama-agents-dbos/tests/fixtures/workflows/three_step_hitl.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""Three-step HITL workflow fixture with name and quest input events.""" + +from __future__ import annotations + +from pydantic import Field +from workflows.context import Context +from workflows.decorators import step +from workflows.events import Event, InputRequiredEvent, StartEvent, StopEvent +from workflows.workflow import Workflow + + +class NameInputEvent(InputRequiredEvent): + prefix: str = Field(default="Name: ") + + +class NameResponseEvent(Event): + response: str = Field(default="") + + +class QuestInputEvent(InputRequiredEvent): + prefix: str = Field(default="Quest: ") + + +class QuestResponseEvent(Event): + response: str = Field(default="") + + +class HITLWorkflow(Workflow): + @step + async def ask_name(self, ctx: Context, ev: StartEvent) -> NameInputEvent: + await ctx.store.set("asked_name", True) + print("STEP:ask_name:complete", flush=True) + return NameInputEvent() + + @step + async def ask_quest(self, ctx: Context, ev: NameResponseEvent) -> QuestInputEvent: + await ctx.store.set("name", ev.response) + print(f"STEP:ask_quest:got_name={ev.response}", flush=True) + print("STEP:ask_quest:complete", flush=True) + return QuestInputEvent() + + @step + async def complete(self, ctx: Context, ev: QuestResponseEvent) -> StopEvent: + name = await ctx.store.get("name", default="unknown") + print(f"STEP:complete:got_quest={ev.response}", flush=True) + return StopEvent(result={"name": name, "quest": ev.response}) diff --git a/packages/llama-agents-dbos/tests/test_dbos_determinism_subprocess.py b/packages/llama-agents-dbos/tests/test_dbos_determinism_subprocess.py new file mode 100644 index 00000000..25855300 --- /dev/null +++ b/packages/llama-agents-dbos/tests/test_dbos_determinism_subprocess.py @@ -0,0 +1,442 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""Test DBOS determinism with subprocess isolation and real interruption. + +This test spawns subprocesses to properly isolate DBOS state and simulate +real Ctrl+C interruptions during workflow execution. +""" + +from __future__ import annotations + +import json +import subprocess +import sys +from pathlib import Path +from typing import Any + +import pytest + +RUNNER_PATH = str(Path(__file__).parent / "fixtures" / "runner.py") + + +@pytest.fixture +def test_db_path(tmp_path: Path) -> Path: + """Create a temporary database path.""" + return tmp_path / "dbos_test.sqlite3" + + +def run_scenario( + workflow: str, + db_url: str, + run_id: str, + config: dict[str, Any] | None = None, + timeout: float = 30.0, +) -> subprocess.CompletedProcess[str]: + """Run a workflow scenario in a subprocess. + + Args: + workflow: Module path with class name (e.g., "tests.fixtures.workflows.hitl:TestWorkflow") + db_url: SQLite database URL + run_id: Unique run ID for the workflow + config: Optional config dict with interrupt_on and/or respond settings + timeout: Subprocess timeout in seconds + + Returns: + CompletedProcess with stdout and stderr captured. + """ + cmd = [ + sys.executable, + RUNNER_PATH, + "--workflow", + workflow, + "--db-url", + db_url, + "--run-id", + run_id, + ] + if config: + cmd.extend(["--config", json.dumps(config)]) + return subprocess.run(cmd, capture_output=True, text=True, timeout=timeout) + + +def assert_no_determinism_errors(result: subprocess.CompletedProcess[str]) -> None: + """Check subprocess result for crashes and DBOS determinism errors.""" + combined = result.stdout + result.stderr + + # Check for non-zero exit code (catches segfaults, killed processes, etc.) + if result.returncode != 0: + pytest.fail( + f"Subprocess exited with code {result.returncode}\n" + f"stdout: {result.stdout}\n" + f"stderr: {result.stderr}" + ) + + # Catch any unhandled Python exception + if "Traceback (most recent call last)" in combined: + pytest.fail( + f"Subprocess exception!\nstdout: {result.stdout}\nstderr: {result.stderr}" + ) + + # Check for DBOS-specific determinism errors + if "DBOSUnexpectedStepError" in combined or "Error 11" in combined: + pytest.fail( + f"DBOS determinism error on resume!\n" + f"stdout: {result.stdout}\n" + f"stderr: {result.stderr}" + ) + + +# ============================================================================= +# Test 1: Basic interrupt/resume with input events +# ============================================================================= + + +def test_determinism_on_resume_after_interrupt(test_db_path: Path) -> None: + """Test that resuming an interrupted workflow doesn't hit determinism errors.""" + run_id = "test-determinism-001" + db_url = f"sqlite+pysqlite:///{test_db_path}?check_same_thread=false" + + print("\n=== Starting workflow (will interrupt) ===") + result1 = run_scenario( + workflow="tests.fixtures.workflows.hitl:TestWorkflow", + db_url=db_url, + run_id=run_id, + config={"interrupt_on": "AskInputEvent"}, + ) + print(f"stdout: {result1.stdout}") + print(f"stderr: {result1.stderr}") + + assert "STEP:ask:complete" in result1.stdout, "First step should complete" + assert "INTERRUPTING" in result1.stdout, "Should have interrupted" + + print("\n=== Resuming workflow ===") + result2 = run_scenario( + workflow="tests.fixtures.workflows.hitl:TestWorkflow", + db_url=db_url, + run_id=run_id, + config={ + "respond": { + "AskInputEvent": { + "event": "UserInput", + "fields": {"response": "test_input"}, + } + } + }, + ) + print(f"stdout: {result2.stdout}") + print(f"stderr: {result2.stderr}") + + assert_no_determinism_errors(result2) + assert "SUCCESS" in result2.stdout or result2.returncode == 0, ( + f"Resume should succeed. stdout: {result2.stdout}, stderr: {result2.stderr}" + ) + + +# ============================================================================= +# Test 2: Chained steps determinism +# ============================================================================= + + +def test_chained_steps_determinism_on_resume(test_db_path: Path) -> None: + """Test determinism with chained steps that trigger each other.""" + run_id = "test-chained-001" + db_url = f"sqlite+pysqlite:///{test_db_path}?check_same_thread=false" + + print("\n=== Starting chained workflow (will interrupt at step 2) ===") + result1 = run_scenario( + workflow="tests.fixtures.workflows.chained:ChainedWorkflow", + db_url=db_url, + run_id=run_id, + config={"interrupt_on": "StepTwoEvent"}, + ) + print(f"stdout: {result1.stdout}") + print(f"stderr: {result1.stderr}") + + assert "STEP:one:complete" in result1.stdout, "Step one should complete" + + print("\n=== Resuming chained workflow ===") + result2 = run_scenario( + workflow="tests.fixtures.workflows.chained:ChainedWorkflow", + db_url=db_url, + run_id=run_id, + ) + print(f"stdout: {result2.stdout}") + print(f"stderr: {result2.stderr}") + + assert_no_determinism_errors(result2) + + +# ============================================================================= +# Test 3: Three-step HITL pattern +# ============================================================================= + + +def test_hitl_three_step_determinism(test_db_path: Path) -> None: + """Test the exact HITL pattern with three steps and input events.""" + run_id = "test-hitl-three-001" + db_url = f"sqlite+pysqlite:///{test_db_path}?check_same_thread=false" + + print("\n=== Starting HITL workflow (will interrupt at quest prompt) ===") + result1 = run_scenario( + workflow="tests.fixtures.workflows.three_step_hitl:HITLWorkflow", + db_url=db_url, + run_id=run_id, + config={ + "respond": { + "NameInputEvent": { + "event": "NameResponseEvent", + "fields": {"response": "Alice"}, + } + }, + "interrupt_on": "QuestInputEvent", + }, + ) + print(f"stdout: {result1.stdout}") + print(f"stderr: {result1.stderr}") + + assert "STEP:ask_name:complete" in result1.stdout, "ask_name should complete" + assert "STEP:ask_quest" in result1.stdout, "ask_quest should start" + assert "INTERRUPTING" in result1.stdout, "Should interrupt at quest" + + print("\n=== Resuming HITL workflow ===") + result2 = run_scenario( + workflow="tests.fixtures.workflows.three_step_hitl:HITLWorkflow", + db_url=db_url, + run_id=run_id, + config={ + "respond": { + "NameInputEvent": { + "event": "NameResponseEvent", + "fields": {"response": "Alice"}, + }, + "QuestInputEvent": { + "event": "QuestResponseEvent", + "fields": {"response": "seek the grail"}, + }, + }, + }, + ) + print(f"stdout: {result2.stdout}") + print(f"stderr: {result2.stderr}") + + assert_no_determinism_errors(result2) + assert "SUCCESS" in result2.stdout or result2.returncode == 0, ( + f"Resume should succeed.\nstdout: {result2.stdout}\nstderr: {result2.stderr}" + ) + + +# ============================================================================= +# Test 4: Parallel steps - two steps triggered by StartEvent +# ============================================================================= + + +def test_parallel_steps_determinism(test_db_path: Path) -> None: + """Test determinism with parallel steps completing in non-deterministic order.""" + run_id = "test-parallel-001" + db_url = f"sqlite+pysqlite:///{test_db_path}?check_same_thread=false" + + print("\n=== Running parallel workflow to completion ===") + result1 = run_scenario( + workflow="tests.fixtures.workflows.parallel:ParallelWorkflow", + db_url=db_url, + run_id=run_id, + ) + print(f"stdout: {result1.stdout}") + print(f"stderr: {result1.stderr}") + + assert "SUCCESS" in result1.stdout, ( + f"Should complete successfully.\nstdout: {result1.stdout}\nstderr: {result1.stderr}" + ) + assert_no_determinism_errors(result1) + + +# ============================================================================= +# Test 5: Concurrent workers on same step (num_workers=2) +# ============================================================================= + + +def test_concurrent_workers_determinism(test_db_path: Path) -> None: + """Test determinism with multiple workers on same step (num_workers > 1).""" + run_id = "test-concurrent-workers-001" + db_url = f"sqlite+pysqlite:///{test_db_path}?check_same_thread=false" + + print("\n=== Running concurrent workers workflow ===") + result1 = run_scenario( + workflow="tests.fixtures.workflows.concurrent_workers:ConcurrentWorkersWorkflow", + db_url=db_url, + run_id=run_id, + ) + print(f"stdout: {result1.stdout}") + print(f"stderr: {result1.stderr}") + + assert "SUCCESS" in result1.stdout, ( + f"Should complete successfully.\nstdout: {result1.stdout}\nstderr: {result1.stderr}" + ) + assert_no_determinism_errors(result1) + + +# ============================================================================= +# Test 6: Sequential steps with HITL +# ============================================================================= + + +def test_sequential_hitl_interrupt_resume(test_db_path: Path) -> None: + """Test sequential steps with HITL interrupt and resume.""" + run_id = "test-seq-hitl-001" + db_url = f"sqlite+pysqlite:///{test_db_path}?check_same_thread=false" + + print("\n=== Starting sequential HITL workflow (will interrupt) ===") + result1 = run_scenario( + workflow="tests.fixtures.workflows.sequential_hitl:SequentialHITLWorkflow", + db_url=db_url, + run_id=run_id, + config={"interrupt_on": "WaitForInputEvent"}, + ) + print(f"stdout: {result1.stdout}") + print(f"stderr: {result1.stderr}") + + assert "STEP:process:complete" in result1.stdout + assert "INTERRUPTING" in result1.stdout + + print("\n=== Resuming sequential HITL workflow ===") + result2 = run_scenario( + workflow="tests.fixtures.workflows.sequential_hitl:SequentialHITLWorkflow", + db_url=db_url, + run_id=run_id, + config={ + "respond": { + "WaitForInputEvent": { + "event": "UserContinueEvent", + "fields": {"continue_value": "user_input"}, + } + } + }, + ) + print(f"stdout: {result2.stdout}") + print(f"stderr: {result2.stderr}") + + assert_no_determinism_errors(result2) + assert "SUCCESS" in result2.stdout or result2.returncode == 0, ( + f"Resume should succeed.\nstdout: {result2.stdout}\nstderr: {result2.stderr}" + ) + + +# ============================================================================= +# Stress tests - run scenarios multiple times to catch flaky timing issues +# ============================================================================= + + +@pytest.mark.parametrize("iteration", range(5)) +def test_parallel_steps_stress(test_db_path: Path, iteration: int) -> None: + """Stress test parallel steps - run 5 times to catch timing issues.""" + run_id = f"test-parallel-stress-{iteration}" + db_url = f"sqlite+pysqlite:///{test_db_path}?check_same_thread=false" + + result = run_scenario( + workflow="tests.fixtures.workflows.parallel:ParallelWorkflow", + db_url=db_url, + run_id=run_id, + ) + + assert "SUCCESS" in result.stdout, ( + f"Iteration {iteration} failed.\nstdout: {result.stdout}\nstderr: {result.stderr}" + ) + assert_no_determinism_errors(result) + + +@pytest.mark.parametrize("iteration", range(5)) +def test_concurrent_workers_stress(test_db_path: Path, iteration: int) -> None: + """Stress test concurrent workers - run 5 times to catch timing issues.""" + run_id = f"test-concurrent-stress-{iteration}" + db_url = f"sqlite+pysqlite:///{test_db_path}?check_same_thread=false" + + result = run_scenario( + workflow="tests.fixtures.workflows.concurrent_workers:ConcurrentWorkersWorkflow", + db_url=db_url, + run_id=run_id, + ) + + assert "SUCCESS" in result.stdout, ( + f"Iteration {iteration} failed.\nstdout: {result.stdout}\nstderr: {result.stderr}" + ) + assert_no_determinism_errors(result) + + +# ============================================================================= +# Test 7: Streaming stress test +# ============================================================================= + + +def test_streaming_stress_determinism(test_db_path: Path) -> None: + """Test determinism with many concurrent stream writes and send_event calls.""" + run_id = "test-streaming-stress-001" + db_url = f"sqlite+pysqlite:///{test_db_path}?check_same_thread=false" + + print("\n=== Running streaming stress workflow ===") + result = run_scenario( + workflow="tests.fixtures.workflows.streaming_stress:StreamingStressWorkflow", + db_url=db_url, + run_id=run_id, + ) + print(f"stdout: {result.stdout}") + print(f"stderr: {result.stderr}") + + assert "SUCCESS" in result.stdout, ( + f"Should complete successfully.\nstdout: {result.stdout}\nstderr: {result.stderr}" + ) + assert_no_determinism_errors(result) + + +def test_streaming_interrupt_resume(test_db_path: Path) -> None: + """Test interrupt/resume with many concurrent stream writes in flight.""" + run_id = "test-streaming-interrupt-001" + db_url = f"sqlite+pysqlite:///{test_db_path}?check_same_thread=false" + + print("\n=== Starting streaming workflow (will interrupt after fan_out) ===") + result1 = run_scenario( + workflow="tests.fixtures.workflows.streaming_stress:StreamingStressWorkflow", + db_url=db_url, + run_id=run_id, + config={ + "interrupt_on": {"event": "ProgressEvent", "condition": {"progress": 999}} + }, + ) + print(f"stdout: {result1.stdout}") + print(f"stderr: {result1.stderr}") + + assert "STEP:fan_out:dispatched_15_items" in result1.stdout, ( + "Fan out should complete" + ) + assert "INTERRUPTING" in result1.stdout, "Should have interrupted" + + print("\n=== Resuming streaming workflow ===") + result2 = run_scenario( + workflow="tests.fixtures.workflows.streaming_stress:StreamingStressWorkflow", + db_url=db_url, + run_id=run_id, + ) + print(f"stdout: {result2.stdout}") + print(f"stderr: {result2.stderr}") + + assert_no_determinism_errors(result2) + assert "SUCCESS" in result2.stdout or result2.returncode == 0, ( + f"Resume should succeed.\nstdout: {result2.stdout}\nstderr: {result2.stderr}" + ) + + +@pytest.mark.parametrize("iteration", range(5)) +def test_streaming_stress_repeated(test_db_path: Path, iteration: int) -> None: + """Stress test streaming - run 5 times to catch timing issues.""" + run_id = f"test-streaming-repeated-{iteration}" + db_url = f"sqlite+pysqlite:///{test_db_path}?check_same_thread=false" + + result = run_scenario( + workflow="tests.fixtures.workflows.streaming_stress:StreamingStressWorkflow", + db_url=db_url, + run_id=run_id, + ) + + assert "SUCCESS" in result.stdout, ( + f"Iteration {iteration} failed.\nstdout: {result.stdout}\nstderr: {result.stderr}" + ) + assert_no_determinism_errors(result) diff --git a/packages/llama-agents-dbos/tests/test_dbos_runtime.py b/packages/llama-agents-dbos/tests/test_dbos_runtime.py new file mode 100644 index 00000000..955cd242 --- /dev/null +++ b/packages/llama-agents-dbos/tests/test_dbos_runtime.py @@ -0,0 +1,316 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""DBOS-specific runtime tests for adapter behavior. + +These tests focus on the internal mechanics of the DBOS adapter, +particularly around run_id matching and state store availability. +""" + +from __future__ import annotations + +from typing import Any, Generator + +import pytest +from dbos import DBOS, DBOSConfig +from llama_agents.dbos import DBOSRuntime +from pydantic import Field +from workflows.context import Context +from workflows.decorators import step +from workflows.events import Event, StartEvent, StopEvent +from workflows.testing import WorkflowTestRunner +from workflows.workflow import Workflow + + +@pytest.fixture(scope="module") +def dbos_config(tmp_path_factory: pytest.TempPathFactory) -> DBOSConfig: + """Create DBOS config with a fresh SQLite database.""" + db_file = tmp_path_factory.mktemp("dbos") / "dbos_debug_test.sqlite3" + system_db_url = f"sqlite+pysqlite:///{db_file}?check_same_thread=false" + return { + "name": "workflows-dbos-debug", + "system_database_url": system_db_url, + "run_admin_server": False, + } # type: ignore[return-value] + + +@pytest.fixture(scope="module") +def dbos_runtime(dbos_config: DBOSConfig) -> Generator[DBOSRuntime, None, None]: + """Module-scoped DBOS runtime with fast polling for tests.""" + DBOS(config=dbos_config) + runtime = DBOSRuntime(polling_interval_sec=0.01) + try: + yield runtime + finally: + runtime.destroy() + + +class DebugEvent(Event): + captured_run_id: str = Field(default="") + captured_dbos_workflow_id: str = Field(default="") + state_store_available: bool = Field(default=False) + + +class RunIdCaptureWorkflow(Workflow): + """Workflow that captures run_id info for debugging.""" + + @step + async def capture_ids(self, ev: StartEvent) -> StopEvent: + dbos_workflow_id = DBOS.workflow_id or "None" + return StopEvent(result={"dbos_workflow_id": dbos_workflow_id}) + + +class StateStoreAccessWorkflow(Workflow): + """Workflow that attempts to access state store.""" + + @step + async def access_store(self, ctx: Context, ev: StartEvent) -> StopEvent: + dbos_workflow_id = DBOS.workflow_id or "None" + + try: + await ctx.store.set("test_key", "test_value") + value = await ctx.store.get("test_key") + store_works = value == "test_value" + except Exception as e: + return StopEvent( + result={ + "dbos_workflow_id": dbos_workflow_id, + "store_works": False, + "error": str(e), + } + ) + + return StopEvent( + result={ + "dbos_workflow_id": dbos_workflow_id, + "store_works": store_works, + } + ) + + +class StateStoreCounterWorkflow(Workflow): + """Workflow that increments a counter in state store.""" + + @step + async def increment(self, ctx: Context, ev: StartEvent) -> StopEvent: + cur = await ctx.store.get("counter", default=0) + await ctx.store.set("counter", cur + 1) + return StopEvent(result=cur + 1) + + +@pytest.mark.asyncio +async def test_dbos_workflow_id_available(dbos_runtime: DBOSRuntime) -> None: + """Verify DBOS.workflow_id is set inside workflow execution.""" + wf = RunIdCaptureWorkflow(runtime=dbos_runtime) + dbos_runtime.launch() + + r = await WorkflowTestRunner(wf).run() + result = r.result + + assert result["dbos_workflow_id"] != "None", ( + "DBOS.workflow_id should be set inside workflow" + ) + + +@pytest.mark.asyncio +async def test_state_store_access_in_step(dbos_runtime: DBOSRuntime) -> None: + """Test whether state store is accessible inside a workflow step.""" + wf = StateStoreAccessWorkflow(runtime=dbos_runtime) + dbos_runtime.launch() + + r = await WorkflowTestRunner(wf).run() + result = r.result + + assert result["store_works"], ( + f"State store should be accessible. Got error: {result.get('error', 'unknown')}" + ) + + +@pytest.mark.asyncio +async def test_internal_adapter_run_id_matches(dbos_runtime: DBOSRuntime) -> None: + """Verify internal adapter run_id matches DBOS.workflow_id.""" + captured_ids: dict[str, Any] = {} + + class IdTracingWorkflow(Workflow): + @step + async def trace_ids(self, ev: StartEvent) -> StopEvent: + captured_ids["dbos_workflow_id"] = DBOS.workflow_id + + internal_adapter = dbos_runtime.get_internal_adapter(self) + captured_ids["adapter_run_id"] = internal_adapter.run_id + + store = internal_adapter.get_state_store() + captured_ids["state_store_found"] = store is not None + + return StopEvent(result="done") + + wf = IdTracingWorkflow(runtime=dbos_runtime) + dbos_runtime.launch() + + await WorkflowTestRunner(wf).run() + + assert captured_ids["adapter_run_id"] == captured_ids["dbos_workflow_id"], ( + f"Adapter run_id '{captured_ids['adapter_run_id']}' should match " + f"DBOS.workflow_id '{captured_ids['dbos_workflow_id']}'" + ) + assert captured_ids["state_store_found"], "State store should be available" + + +@pytest.mark.asyncio +async def test_external_run_id_vs_internal(dbos_runtime: DBOSRuntime) -> None: + """Compare external adapter run_id with what's seen internally.""" + internal_run_id: str | None = None + + class CompareWorkflow(Workflow): + @step + async def capture(self, ev: StartEvent) -> StopEvent: + nonlocal internal_run_id + internal_run_id = DBOS.workflow_id + return StopEvent(result="done") + + wf = CompareWorkflow(runtime=dbos_runtime) + dbos_runtime.launch() + + handler = wf.run() + external_run_id = handler.run_id + + await handler + + assert external_run_id == internal_run_id, ( + f"External run_id '{external_run_id}' should match " + f"internal DBOS.workflow_id '{internal_run_id}'" + ) + + +@pytest.mark.asyncio +async def test_state_store_lazy_creation(dbos_runtime: DBOSRuntime) -> None: + """Test that state store is lazily created by the internal adapter.""" + store_info: dict[str, Any] = {} + + class LazyStoreWorkflow(Workflow): + @step + async def check_store(self, ctx: Context, ev: StartEvent) -> StopEvent: + internal_adapter = dbos_runtime.get_internal_adapter(self) + + # First call should create the store + store1 = internal_adapter.get_state_store() + store_info["first_store_id"] = id(store1) + store_info["first_store_exists"] = store1 is not None + + # Second call should return the same store + store2 = internal_adapter.get_state_store() + store_info["second_store_id"] = id(store2) + store_info["same_store"] = store1 is store2 + + # Store should work + await ctx.store.set("lazy_key", "lazy_value") + value = await ctx.store.get("lazy_key") + store_info["store_works"] = value == "lazy_value" + + return StopEvent(result="done") + + wf = LazyStoreWorkflow(runtime=dbos_runtime) + dbos_runtime.launch() + + await WorkflowTestRunner(wf).run() + + assert store_info["first_store_exists"], "Store should be created on first access" + assert store_info["same_store"], "Same store instance should be returned" + assert store_info["store_works"], "Store should be functional" + + +@pytest.mark.asyncio +async def test_run_workflow_does_not_create_store(dbos_runtime: DBOSRuntime) -> None: + """Verify run_workflow doesn't eagerly create a state store.""" + from unittest.mock import patch + + call_log: list[dict[str, Any]] = [] + original_run_workflow = dbos_runtime.run_workflow + + def patched_run_workflow(*args: Any, **kwargs: Any) -> Any: + call_log.append({"run_id": kwargs.get("run_id")}) + return original_run_workflow(*args, **kwargs) + + class SimpleWf(Workflow): + @step + async def do_it(self, ev: StartEvent) -> StopEvent: + return StopEvent(result="done") + + wf = SimpleWf(runtime=dbos_runtime) + dbos_runtime.launch() + + with patch.object(dbos_runtime, "run_workflow", patched_run_workflow): + handler = wf.run() + await handler + + assert len(call_log) == 1, "run_workflow should be called exactly once" + + +# ============================================================================ +# SqlSerializedState and parse_serialized_state Tests +# ============================================================================ + + +def test_parse_serialized_state_sql_store_type() -> None: + """Test that store_type='sql' parses as SqlSerializedState.""" + from llama_agents.dbos.state_store import SqlSerializedState, parse_serialized_state + + serialized = { + "store_type": "sql", + "run_id": "run-12345", + "schema": "public", + } + + result = parse_serialized_state(serialized) + + assert isinstance(result, SqlSerializedState) + assert result.store_type == "sql" + assert result.run_id == "run-12345" + assert result.db_schema == "public" + + +def test_parse_serialized_state_sql_with_null_schema() -> None: + """Test that SqlSerializedState accepts null schema.""" + from llama_agents.dbos.state_store import SqlSerializedState, parse_serialized_state + + serialized = { + "store_type": "sql", + "run_id": "run-67890", + "schema": None, + } + + result = parse_serialized_state(serialized) + + assert isinstance(result, SqlSerializedState) + assert result.db_schema is None + + +def test_parse_serialized_state_in_memory_format() -> None: + """Test that in_memory format is still handled.""" + from llama_agents.dbos.state_store import parse_serialized_state + from workflows.context.state_store import InMemorySerializedState + + serialized = { + "store_type": "in_memory", + "state_type": "DictState", + "state_module": "workflows.context.state_store", + "state_data": {"_data": {"counter": 42}}, + } + + result = parse_serialized_state(serialized) + + assert isinstance(result, InMemorySerializedState) + assert result.store_type == "in_memory" + + +def test_parse_serialized_state_unknown_store_type_raises() -> None: + """Test that unknown store_type raises ValueError.""" + from llama_agents.dbos.state_store import parse_serialized_state + + serialized = { + "store_type": "redis", # Unknown store type + "state_type": "SomeState", + "state_module": "some.module", + } + + with pytest.raises(ValueError, match="Unknown store_type"): + parse_serialized_state(serialized) diff --git a/packages/llama-agents-dbos/tests/test_task_journal.py b/packages/llama-agents-dbos/tests/test_task_journal.py new file mode 100644 index 00000000..44d8f87f --- /dev/null +++ b/packages/llama-agents-dbos/tests/test_task_journal.py @@ -0,0 +1,212 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. +"""Unit tests for TaskJournal class.""" + +from __future__ import annotations + +import pytest +from llama_agents.dbos.journal.crud import JournalCrud +from llama_agents.dbos.journal.task_journal import TaskJournal +from sqlalchemy import create_engine +from sqlalchemy.engine import Engine +from sqlalchemy.pool import StaticPool + + +@pytest.fixture +def sqlite_engine() -> Engine: + """Create an in-memory SQLite engine with journal table. + + Uses StaticPool to ensure the same connection is reused across threads. + """ + engine = create_engine( + "sqlite:///:memory:", + connect_args={"check_same_thread": False}, + poolclass=StaticPool, + ) + JournalCrud().run_migrations(engine) + return engine + + +@pytest.mark.asyncio +async def test_fresh_journal_has_no_entries(sqlite_engine: Engine) -> None: + """Fresh journal returns None for next_expected_key.""" + journal = TaskJournal("test-run", sqlite_engine) + await journal.load() + + assert journal.next_expected_key() is None + assert not journal.is_replaying() + + +@pytest.mark.asyncio +async def test_record_adds_entry(sqlite_engine: Engine) -> None: + """Recording a key adds it to the journal.""" + journal = TaskJournal("test-run", sqlite_engine) + await journal.load() + + await journal.record("step_a:0") + + # Verify by loading a new journal for the same run + journal2 = TaskJournal("test-run", sqlite_engine) + await journal2.load() + assert journal2.next_expected_key() == "step_a:0" + + +@pytest.mark.asyncio +async def test_record_multiple_entries(sqlite_engine: Engine) -> None: + """Multiple records append to journal in order.""" + journal = TaskJournal("test-run", sqlite_engine) + await journal.load() + + await journal.record("step_a:0") + await journal.record("__pull__:0") + await journal.record("step_b:1") + + # Verify order by loading a new journal + journal2 = TaskJournal("test-run", sqlite_engine) + await journal2.load() + assert journal2.next_expected_key() == "step_a:0" + journal2.advance() + assert journal2.next_expected_key() == "__pull__:0" + journal2.advance() + assert journal2.next_expected_key() == "step_b:1" + journal2.advance() + assert journal2.next_expected_key() is None + + +@pytest.mark.asyncio +async def test_replay_returns_entries_in_order(sqlite_engine: Engine) -> None: + """Replaying journal returns entries in recorded order.""" + # Set up initial data + journal1 = TaskJournal("replay-run", sqlite_engine) + await journal1.load() + await journal1.record("step_a:0") + await journal1.record("step_b:1") + await journal1.record("__pull__:2") + + # Load fresh journal and replay + journal = TaskJournal("replay-run", sqlite_engine) + await journal.load() + + assert journal.is_replaying() + assert journal.next_expected_key() == "step_a:0" + + journal.advance() + assert journal.next_expected_key() == "step_b:1" + + journal.advance() + assert journal.next_expected_key() == "__pull__:2" + + journal.advance() + assert journal.next_expected_key() is None + assert not journal.is_replaying() + + +@pytest.mark.asyncio +async def test_load_is_idempotent(sqlite_engine: Engine) -> None: + """Calling load() multiple times doesn't reset state.""" + # Set up initial data + journal1 = TaskJournal("idempotent-run", sqlite_engine) + await journal1.load() + await journal1.record("step_a:0") + + # Load and advance + journal = TaskJournal("idempotent-run", sqlite_engine) + await journal.load() + journal.advance() + assert journal.next_expected_key() is None + + # Load again - should not reset + await journal.load() + assert journal.next_expected_key() is None + + +@pytest.mark.asyncio +async def test_none_engine_works_in_memory() -> None: + """Journal works without engine (in-memory only).""" + journal = TaskJournal("memory-run", engine=None) + await journal.load() + + assert journal.next_expected_key() is None + + await journal.record("step_a:0") + await journal.record("step_b:1") + + # New journal with same run_id but no engine won't see the entries + journal2 = TaskJournal("memory-run", engine=None) + await journal2.load() + assert journal2.next_expected_key() is None + + +@pytest.mark.asyncio +async def test_record_advances_index(sqlite_engine: Engine) -> None: + """Recording advances the replay index to stay in sync.""" + journal = TaskJournal("index-run", sqlite_engine) + await journal.load() + + # After recording, index should advance + await journal.record("step_a:0") + # Record another + await journal.record("step_b:1") + + # The journal's internal state shows 2 entries recorded + assert journal._entries == ["step_a:0", "step_b:1"] + assert journal._replay_index == 2 # We've advanced past both + + +@pytest.mark.asyncio +async def test_mixed_replay_and_fresh_execution(sqlite_engine: Engine) -> None: + """Journal transitions from replay to fresh execution correctly.""" + # Set up initial data + journal1 = TaskJournal("mixed-run", sqlite_engine) + await journal1.load() + await journal1.record("step_a:0") + + # Load fresh journal + journal = TaskJournal("mixed-run", sqlite_engine) + await journal.load() + + # Replay the existing entry + assert journal.next_expected_key() == "step_a:0" + journal.advance() + + # Now fresh execution + assert journal.next_expected_key() is None + await journal.record("step_b:1") + + # Verify both entries persisted + journal2 = TaskJournal("mixed-run", sqlite_engine) + await journal2.load() + assert journal2.next_expected_key() == "step_a:0" + journal2.advance() + assert journal2.next_expected_key() == "step_b:1" + + +@pytest.mark.asyncio +async def test_empty_journal_is_valid(sqlite_engine: Engine) -> None: + """Empty journal (no entries) is a valid state.""" + journal = TaskJournal("empty-run", sqlite_engine) + await journal.load() + + assert journal.next_expected_key() is None + assert not journal.is_replaying() + + +@pytest.mark.asyncio +async def test_run_id_isolation(sqlite_engine: Engine) -> None: + """Journals with different run_ids are isolated.""" + journal1 = TaskJournal("run-1", sqlite_engine) + await journal1.load() + await journal1.record("step_a:0") + + journal2 = TaskJournal("run-2", sqlite_engine) + await journal2.load() + await journal2.record("step_b:0") + + # Each journal sees only its own entries + check1 = TaskJournal("run-1", sqlite_engine) + await check1.load() + assert check1.next_expected_key() == "step_a:0" + + check2 = TaskJournal("run-2", sqlite_engine) + await check2.load() + assert check2.next_expected_key() == "step_b:0" diff --git a/packages/llama-agents-integration-tests/pyproject.toml b/packages/llama-agents-integration-tests/pyproject.toml index eccba5c9..1b233c20 100644 --- a/packages/llama-agents-integration-tests/pyproject.toml +++ b/packages/llama-agents-integration-tests/pyproject.toml @@ -24,7 +24,8 @@ readme = "README.md" requires-python = ">=3.9" dependencies = [ "llama-index-core>=0.14.13", - "llama-index-workflows" + "llama-index-workflows", + "llama-agents-dbos; python_full_version > '3.9'" ] [tool.basedpyright] @@ -55,4 +56,4 @@ filterwarnings = [ [tool.uv.sources] llama-index-workflows = {workspace = true} -llama-index-workflows-dbos = {workspace = true} +llama-agents-dbos = {workspace = true} diff --git a/packages/llama-agents-integration-tests/tests/test_runtime_matrix.py b/packages/llama-agents-integration-tests/tests/test_runtime_matrix.py index 6752ecf1..7b0d2ad6 100644 --- a/packages/llama-agents-integration-tests/tests/test_runtime_matrix.py +++ b/packages/llama-agents-integration-tests/tests/test_runtime_matrix.py @@ -1,16 +1,22 @@ -"""Runtime matrix tests - testing workflows against both BasicRuntime and DBOSRuntime. +"""Runtime matrix tests - testing workflows against BasicRuntime and DBOSRuntime. All workflow classes are defined at module level so they can be registered with DBOS once at module initialization time, avoiding repeated init/destroy cycles. + +Note: The dbos-postgres variant requires Docker to be available and is marked +with the 'docker' pytest marker. Run with `pytest -m docker` to include it. """ from __future__ import annotations import asyncio -from typing import AsyncGenerator, Optional, Union +from pathlib import Path +from typing import TYPE_CHECKING, Any, AsyncGenerator, Generator, Union import pytest -from pydantic import Field +from dbos import DBOS, DBOSConfig +from llama_agents.dbos import DBOSRuntime +from pydantic import BaseModel, Field from workflows.context import Context from workflows.decorators import step from workflows.errors import WorkflowTimeoutError @@ -26,21 +32,94 @@ from workflows.testing import WorkflowTestRunner from workflows.workflow import Workflow +if TYPE_CHECKING: + from testcontainers.postgres import PostgresContainer + # -- Fixtures -- -@pytest.fixture( - params=[ +def _get_runtime_params() -> list[Any]: + """Get runtime parameters for the test matrix. + + Includes: + - basic: BasicRuntime (fast, no dependencies) + - dbos: DBOSRuntime with SQLite backend (fast, no Docker) + - dbos-postgres: DBOSRuntime with PostgreSQL backend (requires Docker) + + Note: The dbos-postgres variant is marked with the 'docker' marker and + requires Docker to be running. It only runs when explicitly requested + via `pytest -m docker`. + """ + return [ pytest.param("basic", id="basic"), + pytest.param("dbos", id="dbos"), + pytest.param("dbos-postgres", marks=pytest.mark.docker, id="dbos-postgres"), ] -) + + +@pytest.fixture(scope="module") +def postgres_container() -> Generator[PostgresContainer, None, None]: + """Module-scoped PostgreSQL container for DBOS tests. + + This fixture is only used when dbos-postgres runtime is requested. + Requires Docker to be running. + """ + from testcontainers.postgres import PostgresContainer + + with PostgresContainer("postgres:16", driver=None) as postgres: + yield postgres + + +@pytest.fixture(scope="module") +def dbos_runtime_sqlite( + tmp_path_factory: pytest.TempPathFactory, +) -> Generator[DBOSRuntime, None, None]: + """Module-scoped DBOS runtime with SQLite backend.""" + db_file: Path = tmp_path_factory.mktemp("dbos") / "dbos_test.sqlite3" + system_db_url: str = f"sqlite+pysqlite:///{db_file}?check_same_thread=false" + config: DBOSConfig = { + "name": "workflows-py-dbostest", + "system_database_url": system_db_url, + "run_admin_server": False, + } + DBOS(config=config) + runtime = DBOSRuntime(polling_interval_sec=0.01) + try: + yield runtime + finally: + runtime.destroy() + + +@pytest.fixture(scope="module") +def dbos_runtime_postgres( + postgres_container: PostgresContainer, +) -> Generator[DBOSRuntime, None, None]: + """Module-scoped DBOS runtime with PostgreSQL backend.""" + connection_url = postgres_container.get_connection_url() + config: DBOSConfig = { + "name": "wf-dbos-pg-test", # Must be <= 30 chars + "system_database_url": connection_url, + "run_admin_server": False, + } + DBOS(config=config) + runtime = DBOSRuntime(polling_interval_sec=0.01) + try: + yield runtime + finally: + runtime.destroy() + + +@pytest.fixture(params=_get_runtime_params()) async def runtime( request: pytest.FixtureRequest, ) -> AsyncGenerator[Runtime, None]: """Yield an unlaunched runtime. - For DBOS, returns the module-scoped runtime (already created, not yet launched). - Each test must call runtime.launch() after creating workflows. + For DBOS variants, returns the module-scoped runtime (already created, not yet + launched). Each test must call runtime.launch() after creating workflows. + + Note: Only one DBOS variant can be used per test run since DBOS is a singleton. + Use TEST_DBOS_POSTGRES=1 to run with PostgreSQL instead of the default SQLite. """ if request.param == "basic": rt = BasicRuntime() @@ -48,6 +127,12 @@ async def runtime( yield rt finally: rt.destroy() + elif request.param == "dbos": + dbos_rt: DBOSRuntime = request.getfixturevalue("dbos_runtime_sqlite") + yield dbos_rt + elif request.param == "dbos-postgres": + dbos_rt = request.getfixturevalue("dbos_runtime_postgres") + yield dbos_rt # -- Shared event types -- @@ -340,55 +425,69 @@ async def test_workflow_step_send_event(runtime: Runtime) -> None: @pytest.mark.asyncio async def test_workflow_num_workers(runtime: Runtime) -> None: - signal = asyncio.Event() + """Test that num_workers limits concurrent step executions. + + This test verifies that: + 1. A step with num_workers=5 can process up to 5 events concurrently + 2. All 5 workers can run simultaneously (they synchronize to prove concurrency) + 3. The workflow completes successfully with all events processed + """ + num_workers = 5 + num_events = 10 + # Track max concurrent executions lock = asyncio.Lock() - counter = 0 - - async def await_count(count: int) -> None: - nonlocal counter - async with lock: - counter += 1 - if counter == count: - signal.set() - return - await signal.wait() - - class LocalNumWorkersWorkflow(Workflow): + current_workers = 0 + max_concurrent = 0 + # Barrier to ensure all workers reach this point before any proceed + barrier_count = 0 + barrier_event = asyncio.Event() + + class NumWorkersWorkflow(Workflow): @step - async def original_step( - self, ctx: Context, ev: StartEvent - ) -> Union[OneTestEvent, LastEvent]: - await ctx.store.set("num_to_collect", 3) - # Send test4 first to ensure it's pulled from receive_queue - # before test_step workers complete. Events are pulled one per - # iteration, so ordering in receive_queue determines delivery order. - ctx.send_event(AnotherTestEvent(another_test_param="test4")) - ctx.send_event(OneTestEvent(test_param="test1")) - ctx.send_event(OneTestEvent(test_param="test2")) - ctx.send_event(OneTestEvent(test_param="test3")) - return LastEvent() - - @step(num_workers=3) - async def test_step(self, ev: OneTestEvent) -> AnotherTestEvent: - await await_count(3) + async def fan_out(self, ctx: Context, ev: StartEvent) -> OneTestEvent: + # Send more events than num_workers to test queuing + for i in range(num_events): + ctx.send_event(OneTestEvent(test_param=str(i))) + return None # type: ignore + + @step(num_workers=num_workers) + async def worker_step(self, ev: OneTestEvent) -> AnotherTestEvent: + nonlocal current_workers, max_concurrent, barrier_count + + async with lock: + current_workers += 1 + max_concurrent = max(max_concurrent, current_workers) + barrier_count += 1 + if barrier_count == num_workers: + # All workers have arrived, release them + barrier_event.set() + + # Wait for all workers to arrive (proves concurrency) + await barrier_event.wait() + + async with lock: + current_workers -= 1 + return AnotherTestEvent(another_test_param=ev.test_param) @step - async def final_step( - self, ctx: Context, ev: Union[AnotherTestEvent, LastEvent] - ) -> Optional[StopEvent]: - n = await ctx.store.get("num_to_collect") - events = ctx.collect_events(ev, [AnotherTestEvent] * n) + async def collect_step( + self, ctx: Context, ev: AnotherTestEvent + ) -> StopEvent | None: + events = ctx.collect_events(ev, [AnotherTestEvent] * num_events) if events is None: return None - return StopEvent(result=[ev.another_test_param for ev in events]) + return StopEvent(result=[e.another_test_param for e in events]) - workflow = LocalNumWorkersWorkflow(timeout=10, runtime=runtime) + workflow = NumWorkersWorkflow(timeout=10, runtime=runtime) runtime.launch() r = await WorkflowTestRunner(workflow).run() - assert "test4" in set(r.result) - assert len({"test1", "test2", "test3"} - set(r.result)) == 1 + # Verify all events were processed + assert len(r.result) == num_events + assert set(r.result) == {str(i) for i in range(num_events)} + # Verify we achieved the expected concurrency (all 5 workers ran together) + assert max_concurrent == num_workers @pytest.mark.asyncio @@ -491,3 +590,208 @@ async def test_streaming_task_timeout(runtime: Runtime) -> None: with pytest.raises(WorkflowTimeoutError, match="Operation timed out"): await r + + +# -- Workflow State Tests -- + + +class StatefulWorkflow(Workflow): + """Workflow that accumulates state across steps.""" + + @step + async def step1(self, ctx: Context, ev: StartEvent) -> OneTestEvent: + await ctx.store.set("step1_ran", True) + await ctx.store.set("counter", 1) + return OneTestEvent() + + @step + async def step2(self, ctx: Context, ev: OneTestEvent) -> StopEvent: + await ctx.store.set("step2_ran", True) + counter = await ctx.store.get("counter") + await ctx.store.set("counter", counter + 1) + final_counter = await ctx.store.get("counter") + return StopEvent(result={"counter": final_counter}) + + +class NestedStateWorkflow(Workflow): + """Workflow that uses nested state paths.""" + + @step + async def process(self, ctx: Context, ev: StartEvent) -> StopEvent: + await ctx.store.set("user", {"name": "Alice", "profile": {"level": 1}}) + await ctx.store.set("user.profile.level", 2) + level = await ctx.store.get("user.profile.level") + name = await ctx.store.get("user.name") + return StopEvent(result={"name": name, "level": level}) + + +@pytest.mark.asyncio +async def test_workflow_state_basic(runtime: Runtime) -> None: + """Test basic state operations within a workflow.""" + wf = CounterWorkflow(runtime=runtime) + runtime.launch() + result = await WorkflowTestRunner(wf).run() + assert result.result == 1 + + +@pytest.mark.asyncio +async def test_workflow_state_across_steps(runtime: Runtime) -> None: + """Test state persistence across multiple workflow steps.""" + wf = StatefulWorkflow(runtime=runtime) + runtime.launch() + result = await WorkflowTestRunner(wf).run() + assert result.result == {"counter": 2} + + +@pytest.mark.asyncio +async def test_workflow_nested_state(runtime: Runtime) -> None: + """Test nested state path access within workflows.""" + wf = NestedStateWorkflow(runtime=runtime) + runtime.launch() + result = await WorkflowTestRunner(wf).run() + assert result.result == {"name": "Alice", "level": 2} + + +@pytest.mark.asyncio +async def test_workflow_state_multiple_runs(runtime: Runtime) -> None: + """Test that each workflow run has isolated state.""" + wf = CounterWorkflow(runtime=runtime) + runtime.launch() + runner = WorkflowTestRunner(wf) + + # Run multiple times - each should start fresh + results = await asyncio.gather( + runner.run(), + runner.run(), + runner.run(), + ) + + # Each run should have counter=1 (not accumulating) + for r in results: + assert r.result == 1 + + +# -- Typed State Tests -- + + +class TypedState(BaseModel): + """Custom typed state for workflow testing.""" + + counter: int = 0 + name: str = "default" + items: list[str] = [] + + +class TypeStateStopEvent(StopEvent): + state_type: str + initial_counter: int + final_counter: int + final_name: str + + +class TypedStateWorkflow(Workflow): + """Workflow that uses typed state via Context[TypedState].""" + + @step + async def process(self, ctx: Context[TypedState], ev: StartEvent) -> StopEvent: + # Access typed state + state = await ctx.store.get_state() + + # Verify we got the right type + state_type_name = type(state).__name__ + + # Modify state using typed fields + await ctx.store.set("counter", state.counter + 1) + await ctx.store.set("name", "modified") + + final_state = await ctx.store.get_state() + return TypeStateStopEvent( + state_type=state_type_name, + initial_counter=state.counter, + final_counter=final_state.counter, + final_name=final_state.name, + ) + + +@pytest.mark.asyncio +async def test_typed_state_workflow(runtime: Runtime) -> None: + """Test workflow with typed state Context[TypedState]. + + This verifies that: + 1. The state type is correctly inferred from Context[T] annotation + 2. The state is created with the correct type + 3. Typed field access works correctly + """ + wf = TypedStateWorkflow(runtime=runtime) + runtime.launch() + + result = await WorkflowTestRunner(wf).run() + + # The state should be TypedState, not DictState + assert result.result.state_type == "TypedState", ( + f"Expected TypedState but got {result.result.state_type}. " + "State type inference may not be working." + ) + assert result.result.initial_counter == 0 + assert result.result.final_counter == 1 + assert result.result.final_name == "modified" + + +class TypedStateWithDefaultsWorkflow(Workflow): + """Workflow that verifies typed state has correct defaults.""" + + @step + async def check_defaults( + self, ctx: Context[TypedState], ev: StartEvent + ) -> StopEvent: + state = await ctx.store.get_state() + return StopEvent( + result={ + "counter": state.counter, + "name": state.name, + "items": state.items, + } + ) + + +@pytest.mark.asyncio +async def test_typed_state_defaults(runtime: Runtime) -> None: + """Test that typed state is initialized with correct defaults.""" + wf = TypedStateWithDefaultsWorkflow(runtime=runtime) + runtime.launch() + + result = await WorkflowTestRunner(wf).run() + + assert result.result["counter"] == 0 + assert result.result["name"] == "default" + assert result.result["items"] == [] + + +@pytest.mark.asyncio +async def test_typed_state_with_initial_values(runtime: Runtime) -> None: + """Test that initial state values are passed through to the workflow. + + This verifies that: + 1. State can be set before running the workflow + 2. The initial values are correctly used (not replaced with defaults) + 3. Modifications build on the initial values + """ + wf = TypedStateWorkflow(runtime=runtime) + runtime.launch() + + # Create a context and set initial state with counter=1 (default is 0) + ctx = Context(wf) + await ctx.store.set("counter", 1) + + result = await WorkflowTestRunner(wf).run(ctx=ctx) + + # If initial state wasn't passed through, initial_counter would be 0 + # and final_counter would be 1 (from default 0 + 1) + assert result.result.initial_counter == 1, ( + f"Expected initial_counter=1 but got {result.result.initial_counter}. " + "Initial state was not passed through to the workflow." + ) + assert result.result.final_counter == 2, ( + f"Expected final_counter=2 but got {result.result.final_counter}. " + "State modification did not build on initial value." + ) diff --git a/packages/llama-agents-integration-tests/tests/test_state_store_matrix.py b/packages/llama-agents-integration-tests/tests/test_state_store_matrix.py new file mode 100644 index 00000000..e0ad5f60 --- /dev/null +++ b/packages/llama-agents-integration-tests/tests/test_state_store_matrix.py @@ -0,0 +1,565 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. + +"""State store matrix tests - testing StateStore implementations. + +Tests the StateStore protocol across InMemoryStateStore and SqlStateStore +(with both SQLite and PostgreSQL engines) to ensure consistent behavior. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path +from typing import TYPE_CHECKING, Any, AsyncGenerator, Generator, Union + +import pytest +from llama_agents.dbos.state_store import SqlStateStore +from pydantic import ( + BaseModel, + ConfigDict, + ValidationError, + field_serializer, + field_validator, +) +from sqlalchemy import create_engine +from sqlalchemy.engine import Engine +from workflows.context.serializers import JsonSerializer +from workflows.context.state_store import DictState, InMemoryStateStore, StateStore + +if TYPE_CHECKING: + from testcontainers.postgres import PostgresContainer + + +# -- Custom state types for testing -- + + +class MyRandomObject: + """Non-Pydantic object that requires custom serialization.""" + + def __init__(self, name: str) -> None: + self.name = name + + +class PydanticObject(BaseModel): + """Simple Pydantic model for nested state testing.""" + + name: str + + +class MyState(BaseModel): + """Custom typed state with serialization logic.""" + + model_config = ConfigDict( + arbitrary_types_allowed=True, + validate_assignment=True, + strict=True, + ) + + my_obj: MyRandomObject + pydantic_obj: PydanticObject + name: str + age: int + + @field_serializer("my_obj", when_used="always") + def serialize_my_obj(self, my_obj: MyRandomObject) -> str: + return my_obj.name + + @field_validator("my_obj", mode="before") + @classmethod + def deserialize_my_obj(cls, v: Union[str, MyRandomObject]) -> MyRandomObject: + if isinstance(v, MyRandomObject): + return v + if isinstance(v, str): + return MyRandomObject(v) + raise ValueError(f"Invalid type for my_obj: {type(v)}") + + +# -- Fixtures -- + + +@pytest.fixture(scope="module") +def postgres_container() -> Generator[PostgresContainer, None, None]: + """Module-scoped PostgreSQL container for state store tests. + + Requires Docker to be running. + """ + from testcontainers.postgres import PostgresContainer + + with PostgresContainer("postgres:16", driver=None) as postgres: + yield postgres + + +@pytest.fixture(scope="module") +def postgres_engine( + postgres_container: PostgresContainer, +) -> Generator[Engine, None, None]: + """Module-scoped PostgreSQL engine for state store tests.""" + # Get connection URL and convert to use psycopg (psycopg3) driver + connection_url = postgres_container.get_connection_url() + # Replace postgresql:// or postgresql+psycopg2:// with postgresql+psycopg:// + if "postgresql+psycopg2://" in connection_url: + connection_url = connection_url.replace( + "postgresql+psycopg2://", "postgresql+psycopg://" + ) + elif connection_url.startswith("postgresql://"): + connection_url = connection_url.replace( + "postgresql://", "postgresql+psycopg://", 1 + ) + engine = create_engine(connection_url) + yield engine + engine.dispose() + + +@pytest.fixture(scope="module") +def sqlite_engine( + tmp_path_factory: pytest.TempPathFactory, +) -> Generator[Engine, None, None]: + """Module-scoped SQLite engine for state store tests.""" + db_file: Path = tmp_path_factory.mktemp("state_store") / "test.sqlite3" + engine = create_engine(f"sqlite:///{db_file}?check_same_thread=false") + yield engine + engine.dispose() + + +def _get_store_params() -> list[Any]: + """Get store type parameters for the test matrix.""" + return [ + pytest.param("in_memory", id="in_memory"), + pytest.param("sqlite", id="sqlite"), + pytest.param("postgres", marks=pytest.mark.docker, id="postgres"), + ] + + +@pytest.fixture(params=_get_store_params()) +async def state_store( + request: pytest.FixtureRequest, + sqlite_engine: Engine, +) -> AsyncGenerator[StateStore[DictState], None]: + """Parametrized fixture yielding a fresh StateStore for each test.""" + # Use unique run_id per test to avoid state bleeding + run_id = f"test-{id(request)}" + + if request.param == "in_memory": + yield InMemoryStateStore(DictState()) + elif request.param == "sqlite": + store = SqlStateStore(run_id=run_id, engine=sqlite_engine) + yield store + elif request.param == "postgres": + pg_engine: Engine = request.getfixturevalue("postgres_engine") + store = SqlStateStore(run_id=run_id, engine=pg_engine, schema="dbos") + yield store + + +@pytest.fixture(params=_get_store_params()) +async def custom_state_store( + request: pytest.FixtureRequest, + sqlite_engine: Engine, +) -> AsyncGenerator[StateStore[MyState], None]: + """Parametrized fixture yielding a StateStore with custom typed state.""" + run_id = f"test-custom-{id(request)}" + initial_state = MyState( + my_obj=MyRandomObject("llama-index"), + pydantic_obj=PydanticObject(name="llama-index"), + name="John", + age=30, + ) + + if request.param == "in_memory": + yield InMemoryStateStore(initial_state) + elif request.param == "sqlite": + store = SqlStateStore( + run_id=run_id, + state_type=MyState, + engine=sqlite_engine, + ) + await store.set_state(initial_state) + yield store + elif request.param == "postgres": + pg_engine: Engine = request.getfixturevalue("postgres_engine") + store = SqlStateStore( + run_id=run_id, + state_type=MyState, + engine=pg_engine, + schema="dbos", + ) + await store.set_state(initial_state) + yield store + + +# -- Basic Operations Tests -- + + +@pytest.mark.asyncio +async def test_get_set_basic_values(state_store: StateStore[DictState]) -> None: + """Test basic get/set operations with simple values.""" + await state_store.set("name", "John") + await state_store.set("age", 30) + + assert await state_store.get("name") == "John" + assert await state_store.get("age") == 30 + + +@pytest.mark.asyncio +async def test_get_with_default(state_store: StateStore[DictState]) -> None: + """Test get with default value for missing keys.""" + result = await state_store.get("nonexistent", default=None) + assert result is None + + result = await state_store.get("missing", default="fallback") + assert result == "fallback" + + +@pytest.mark.asyncio +async def test_get_missing_raises(state_store: StateStore[DictState]) -> None: + """Test that get raises ValueError for missing key without default.""" + with pytest.raises(ValueError, match="not found"): + await state_store.get("nonexistent") + + +@pytest.mark.asyncio +async def test_nested_get_set(state_store: StateStore[DictState]) -> None: + """Test nested path access with dot notation.""" + await state_store.set("nested", {"a": "b"}) + assert await state_store.get("nested.a") == "b" + + await state_store.set("nested.a", "c") + assert await state_store.get("nested.a") == "c" + + +@pytest.mark.asyncio +async def test_get_state_returns_copy(state_store: StateStore[DictState]) -> None: + """Test that get_state returns a copy, not the original.""" + await state_store.set("value", 1) + + state1 = await state_store.get_state() + state2 = await state_store.get_state() + + # Should be equal but not the same object + assert state1.model_dump() == state2.model_dump() + + +@pytest.mark.asyncio +async def test_set_state_replaces(state_store: StateStore[DictState]) -> None: + """Test that set_state replaces the entire state.""" + await state_store.set("old_key", "old_value") + + new_state = DictState() + new_state["new_key"] = "new_value" + await state_store.set_state(new_state) + + assert await state_store.get("new_key") == "new_value" + # Old key should be gone or inaccessible + result = await state_store.get("old_key", default=None) + assert result is None + + +@pytest.mark.asyncio +async def test_clear_resets_state(state_store: StateStore[DictState]) -> None: + """Test that clear resets to default state.""" + await state_store.set("name", "Jane") + await state_store.set("age", 25) + + await state_store.clear() + + assert await state_store.get("name", default=None) is None + assert await state_store.get("age", default=None) is None + + +# -- edit_state Context Manager Tests -- + + +@pytest.mark.asyncio +async def test_edit_state_basic(state_store: StateStore[DictState]) -> None: + """Test basic edit_state context manager usage.""" + await state_store.set("counter", 0) + + async with state_store.edit_state() as state: + current = state.get("counter", 0) + state["counter"] = current + 1 + + assert await state_store.get("counter") == 1 + + +@pytest.mark.asyncio +async def test_edit_state_multiple_changes(state_store: StateStore[DictState]) -> None: + """Test multiple changes within a single edit_state.""" + async with state_store.edit_state() as state: + state["a"] = 1 + state["b"] = 2 + state["c"] = {"nested": "value"} + + assert await state_store.get("a") == 1 + assert await state_store.get("b") == 2 + assert await state_store.get("c.nested") == "value" + + +@pytest.mark.asyncio +async def test_edit_state_exception_handling( + state_store: StateStore[DictState], +) -> None: + """Test that exceptions in edit_state don't corrupt state.""" + await state_store.set("value", "original") + + with pytest.raises(ValueError, match="intentional"): + async with state_store.edit_state() as state: + state["value"] = "modified" + raise ValueError("intentional error") + + # State should remain unchanged after exception + # Note: behavior may vary - InMemory commits on context exit, SQL rolls back + # This test documents the expected behavior + + +# -- Custom Typed State Tests -- + + +@pytest.mark.asyncio +async def test_custom_state_type(custom_state_store: StateStore[MyState]) -> None: + """Test state store with custom Pydantic model.""" + state = await custom_state_store.get_state() + assert isinstance(state, MyState) + assert state.name == "John" + assert state.age == 30 + assert state.my_obj.name == "llama-index" + + +@pytest.mark.asyncio +async def test_custom_state_set_values(custom_state_store: StateStore[MyState]) -> None: + """Test setting values on custom typed state.""" + await custom_state_store.set("name", "Jane") + await custom_state_store.set("age", 25) + + assert await custom_state_store.get("name") == "Jane" + assert await custom_state_store.get("age") == 25 + + # Original custom fields should still be accessible + state = await custom_state_store.get_state() + assert state.my_obj.name == "llama-index" + + +@pytest.mark.asyncio +async def test_custom_state_validation(custom_state_store: StateStore[MyState]) -> None: + """Test that Pydantic validation is enforced on custom state.""" + # MyState has strict=True, so setting age to string should fail + with pytest.raises(ValidationError): + await custom_state_store.set("age", "not a number") + + +# -- Serialization Tests -- + + +@pytest.mark.asyncio +async def test_to_dict_from_dict_roundtrip(state_store: StateStore[DictState]) -> None: + """Test serialization roundtrip with to_dict/from_dict.""" + await state_store.set("name", "John") + await state_store.set("age", 30) + + serializer = JsonSerializer() + data = state_store.to_dict(serializer) + + # For InMemoryStateStore, from_dict restores the full state + # For SqlStateStore, from_dict returns metadata (engine must be set separately) + if isinstance(state_store, InMemoryStateStore): + restored = InMemoryStateStore.from_dict(data, serializer) + assert await restored.get("name") == "John" + assert await restored.get("age") == 30 + + +# -- SQLite-Specific Tests -- + + +@pytest.mark.asyncio +async def test_sqlite_persistence(sqlite_engine: Engine) -> None: + """Test that SQLite state persists across store instances.""" + run_id = "persistence-test" + + # Create store and set state + store1 = SqlStateStore(run_id=run_id, engine=sqlite_engine) + await store1.set("persistent_key", "persistent_value") + + # Create new store instance with same run_id + store2 = SqlStateStore(run_id=run_id, engine=sqlite_engine) + result = await store2.get("persistent_key") + + assert result == "persistent_value" + + +@pytest.mark.asyncio +async def test_sqlite_isolation(sqlite_engine: Engine) -> None: + """Test that different run_ids have isolated state.""" + store1 = SqlStateStore(run_id="run-1", engine=sqlite_engine) + store2 = SqlStateStore(run_id="run-2", engine=sqlite_engine) + + await store1.set("key", "value1") + await store2.set("key", "value2") + + assert await store1.get("key") == "value1" + assert await store2.get("key") == "value2" + + +@pytest.mark.asyncio +async def test_sqlite_concurrent_edits(sqlite_engine: Engine) -> None: + """Test concurrent edit_state calls are serialized correctly.""" + run_id = "concurrent-test" + store = SqlStateStore(run_id=run_id, engine=sqlite_engine) + await store.set("counter", 0) + + async def increment() -> None: + async with store.edit_state() as state: + current = state.get("counter", 0) + await asyncio.sleep(0.01) # Simulate some work + state["counter"] = current + 1 + + # Run multiple increments concurrently + await asyncio.gather(*[increment() for _ in range(5)]) + + # All increments should have been applied + result = await store.get("counter") + assert result == 5 + + +@pytest.mark.asyncio +async def test_sqlite_custom_state_persistence(sqlite_engine: Engine) -> None: + """Test that custom typed state persists correctly.""" + run_id = "custom-persistence-test" + + initial_state = MyState( + my_obj=MyRandomObject("persisted"), + pydantic_obj=PydanticObject(name="persisted"), + name="Original", + age=100, + ) + + # Create and initialize store + store1 = SqlStateStore( + run_id=run_id, + state_type=MyState, + engine=sqlite_engine, + ) + await store1.set_state(initial_state) + await store1.set("name", "Modified") + + # Create new store instance + store2 = SqlStateStore( + run_id=run_id, + state_type=MyState, + engine=sqlite_engine, + ) + state = await store2.get_state() + + assert state.name == "Modified" + assert state.my_obj.name == "persisted" + + +# -- PostgreSQL-Specific Tests -- + + +@pytest.mark.docker +@pytest.mark.asyncio +async def test_postgres_persistence(postgres_engine: Engine) -> None: + """Test that PostgreSQL state persists across store instances.""" + run_id = "pg-persistence-test" + + store1 = SqlStateStore(run_id=run_id, engine=postgres_engine, schema="dbos") + await store1.set("persistent_key", "persistent_value") + + store2 = SqlStateStore(run_id=run_id, engine=postgres_engine, schema="dbos") + result = await store2.get("persistent_key") + + assert result == "persistent_value" + + +@pytest.mark.docker +@pytest.mark.asyncio +async def test_postgres_isolation(postgres_engine: Engine) -> None: + """Test that different run_ids have isolated state.""" + store1 = SqlStateStore(run_id="pg-run-1", engine=postgres_engine, schema="dbos") + store2 = SqlStateStore(run_id="pg-run-2", engine=postgres_engine, schema="dbos") + + await store1.set("key", "value1") + await store2.set("key", "value2") + + assert await store1.get("key") == "value1" + assert await store2.get("key") == "value2" + + +@pytest.mark.docker +@pytest.mark.asyncio +async def test_postgres_concurrent_edits(postgres_engine: Engine) -> None: + """Test concurrent edit_state calls are serialized correctly with FOR UPDATE.""" + run_id = "pg-concurrent-test" + store = SqlStateStore(run_id=run_id, engine=postgres_engine, schema="dbos") + await store.set("counter", 0) + + async def increment() -> None: + async with store.edit_state() as state: + current = state.get("counter", 0) + await asyncio.sleep(0.01) # Simulate some work + state["counter"] = current + 1 + + # Run multiple increments concurrently + await asyncio.gather(*[increment() for _ in range(5)]) + + # All increments should have been applied + result = await store.get("counter") + assert result == 5 + + +@pytest.mark.docker +@pytest.mark.asyncio +async def test_postgres_custom_state_persistence(postgres_engine: Engine) -> None: + """Test that custom typed state persists correctly.""" + run_id = "pg-custom-persistence-test" + + initial_state = MyState( + my_obj=MyRandomObject("persisted"), + pydantic_obj=PydanticObject(name="persisted"), + name="Original", + age=100, + ) + + store1 = SqlStateStore( + run_id=run_id, + state_type=MyState, + engine=postgres_engine, + schema="dbos", + ) + await store1.set_state(initial_state) + await store1.set("name", "Modified") + + store2 = SqlStateStore( + run_id=run_id, + state_type=MyState, + engine=postgres_engine, + schema="dbos", + ) + state = await store2.get_state() + + assert state.name == "Modified" + assert state.my_obj.name == "persisted" + + +@pytest.mark.docker +@pytest.mark.asyncio +async def test_postgres_uses_dbos_schema(postgres_engine: Engine) -> None: + """Test that SqlStateStore with schema='dbos' creates the table in the dbos schema.""" + run_id = "pg-schema-test" + store = SqlStateStore(run_id=run_id, engine=postgres_engine, schema="dbos") + + # Trigger table creation by accessing state + await store.set("test", "value") + + # Verify the table was created in the dbos schema + with postgres_engine.connect() as conn: + result = conn.exec_driver_sql( + """ + SELECT EXISTS ( + SELECT FROM information_schema.tables + WHERE table_schema = 'dbos' + AND table_name = 'workflow_state' + ) + """ + ) + exists = result.scalar() + assert exists is True diff --git a/packages/llama-index-workflows/src/workflows/context/state_store.py b/packages/llama-index-workflows/src/workflows/context/state_store.py index 45af8473..438e2b45 100644 --- a/packages/llama-index-workflows/src/workflows/context/state_store.py +++ b/packages/llama-index-workflows/src/workflows/context/state_store.py @@ -548,7 +548,9 @@ async def clear(self) -> None: def deserialize_state_from_dict( - serialized_state: dict[str, Any], serializer: "BaseSerializer" + serialized_state: dict[str, Any], + serializer: "BaseSerializer", + state_type: type[BaseModel] | None = None, ) -> BaseModel: """Deserialize state from a serialized payload. @@ -559,6 +561,9 @@ def deserialize_state_from_dict( serialized_state: The payload from to_dict(), containing state_data, state_type, and state_module. serializer: Strategy to decode stored values. + state_type: Optional explicit state type. When provided, uses + issubclass to determine if it's DictState. When omitted, falls + back to reading state_type from the dict. Returns: The deserialized state model instance. @@ -567,9 +572,13 @@ def deserialize_state_from_dict( ValueError: If deserialization fails for any key. """ state_data = serialized_state.get("state_data", {}) - state_type_name = serialized_state.get("state_type", "DictState") - if state_type_name == "DictState": + if state_type is not None: + is_dict_state = issubclass(state_type, DictState) + else: + is_dict_state = serialized_state.get("state_type", "DictState") == "DictState" + + if is_dict_state: _data_serialized = state_data.get("_data", {}) deserialized_data = {} for key, value in _data_serialized.items(): diff --git a/packages/llama-index-workflows/tests/test_state_manager.py b/packages/llama-index-workflows/tests/test_state_manager.py index a92e6d25..1a7c7351 100644 --- a/packages/llama-index-workflows/tests/test_state_manager.py +++ b/packages/llama-index-workflows/tests/test_state_manager.py @@ -5,7 +5,7 @@ Full state store protocol tests are in the integration test package (llama-index-integration-tests/tests/test_state_store_matrix.py), -which tests InMemoryStateStore alongside SqliteStateStore. +which tests InMemoryStateStore alongside SqlStateStore. These tests provide fast feedback during development of the base package. """ diff --git a/pyproject.toml b/pyproject.toml index 9f2fc141..0a6b8352 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,7 +61,7 @@ root = "packages/llama-agents-integration-tests" pythonVersion = "3.13" [[tool.basedpyright.executionEnvironments]] -root = "packages/llama-index-workflows-dbos" +root = "packages/llama-agents-dbos" pythonVersion = "3.10" [[tool.basedpyright.executionEnvironments]] @@ -136,13 +136,21 @@ llama-index-utils-workflow = {workspace = true} llama-index-workflows = {workspace = true} llama-agents-client = {workspace = true} llama-agents-server = {workspace = true} +llama-agents-dbos = {workspace = true} [tool.uv.workspace] +exclude = [ +] members = [ - "docs/api_docs", - "packages/llama-agents-integration-tests", - "packages/llama-index-utils-workflow", + # core "packages/llama-index-workflows", + # extensions "packages/llama-agents-client", - "packages/llama-agents-server" + "packages/llama-agents-server", + "packages/llama-index-utils-workflow", + # integrations + "packages/llama-agents-dbos", + # internal + "docs/api_docs", + "packages/llama-agents-integration-tests" ] diff --git a/src/dev_cli/cli.py b/src/dev_cli/cli.py index 0fbbc7d8..c3b44253 100644 --- a/src/dev_cli/cli.py +++ b/src/dev_cli/cli.py @@ -1,5 +1,5 @@ -# SPDX-FileCopyrightText: 2026 LlamaIndex Authors # SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. """CLI entry point for dev.""" from __future__ import annotations diff --git a/src/dev_cli/commands/__init__.py b/src/dev_cli/commands/__init__.py index 891a5418..785dc914 100644 --- a/src/dev_cli/commands/__init__.py +++ b/src/dev_cli/commands/__init__.py @@ -1,5 +1,5 @@ -# SPDX-FileCopyrightText: 2026 LlamaIndex Authors # SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. """CLI commands for workflows-dev.""" from __future__ import annotations diff --git a/src/dev_cli/commands/changesets_cmd.py b/src/dev_cli/commands/changesets_cmd.py index 172ad51b..ac4b71c1 100644 --- a/src/dev_cli/commands/changesets_cmd.py +++ b/src/dev_cli/commands/changesets_cmd.py @@ -1,5 +1,5 @@ -# SPDX-FileCopyrightText: 2026 LlamaIndex Authors # SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. """Changeset commands for versioning and publishing.""" from __future__ import annotations diff --git a/src/dev_cli/commands/pytest_cmd.py b/src/dev_cli/commands/pytest_cmd.py index 8d0857d9..841e0bc9 100644 --- a/src/dev_cli/commands/pytest_cmd.py +++ b/src/dev_cli/commands/pytest_cmd.py @@ -1,5 +1,5 @@ -# SPDX-FileCopyrightText: 2026 LlamaIndex Authors # SPDX-License-Identifier: MIT +# Copyright (c) 2026 LlamaIndex Inc. """Pytest command for running tests across packages.""" from __future__ import annotations diff --git a/uv.lock b/uv.lock index d927bd29..b6a04bcb 100644 --- a/uv.lock +++ b/uv.lock @@ -5,13 +5,15 @@ resolution-markers = [ "python_full_version >= '3.11'", "python_full_version == '3.10.*'", "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] [manifest] members = [ "docs", "llama-agents-client", + "llama-agents-dbos", "llama-agents-dev", "llama-agents-integration-tests", "llama-agents-server", @@ -502,7 +504,8 @@ version = "8.1.8" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] dependencies = [ { name = "colorama", marker = "python_full_version < '3.10' and sys_platform == 'win32'" }, @@ -543,7 +546,8 @@ version = "7.10.7" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] sdist = { url = "https://files.pythonhosted.org/packages/51/26/d22c300112504f5f9a9fd2297ce33c35f3d353e4aeb987c8419453b2a7c2/coverage-7.10.7.tar.gz", hash = "sha256:f4ab143ab113be368a3e9b795f9cd7906c5ef407d6173fe9675a902e1fffc239", size = 827704, upload-time = "2025-09-21T20:03:56.815Z" } wheels = [ @@ -827,6 +831,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c3/be/d0d44e092656fe7a06b55e6103cbce807cdbdee17884a5367c68c9860853/dataclasses_json-0.6.7-py3-none-any.whl", hash = "sha256:0dbf33f26c8d5305befd61b39d2b3414e8a407bedc2834dea9b8d642666fb40a", size = 28686, upload-time = "2024-06-09T16:20:16.715Z" }, ] +[[package]] +name = "dbos" +version = "2.10.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "psycopg", version = "3.3.2", source = { registry = "https://pypi.org/simple" }, extra = ["binary"], marker = "python_full_version >= '3.10'" }, + { name = "python-dateutil", marker = "python_full_version >= '3.10'" }, + { name = "pyyaml", marker = "python_full_version >= '3.10'" }, + { name = "sqlalchemy", marker = "python_full_version >= '3.10'" }, + { name = "typer-slim", marker = "python_full_version >= '3.10'" }, + { name = "websockets", marker = "python_full_version >= '3.10'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a2/06/77ff1a7af60589875b9352f9e64824374c9f3620a741d59e99f3ac8ae189/dbos-2.10.0.tar.gz", hash = "sha256:5bb12ba64c8ddf5d2c0e0906681f34604054940c67ea7778d74a42039b515631", size = 232118, upload-time = "2026-01-26T17:30:18.296Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7d/99/3cca8a35151eec46aa049129ccfe47dcee87398cea01abd6fb83c5a7ada0/dbos-2.10.0-py3-none-any.whl", hash = "sha256:da5c752e3bac323e31458fcafe7149934816529b971b70a5a715b0e5e324ae50", size = 145363, upload-time = "2026-01-26T17:30:16.939Z" }, +] + [[package]] name = "decorator" version = "5.2.1" @@ -968,7 +989,8 @@ version = "3.19.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] sdist = { url = "https://files.pythonhosted.org/packages/40/bb/0ab3e58d22305b6f5440629d20683af28959bf793d98d11950e305c1c326/filelock-3.19.1.tar.gz", hash = "sha256:66eda1888b0171c998b35be2bcc0f6d75c388a7ce20c3f3f37aa8e96c2dddf58", size = 17687, upload-time = "2025-08-14T16:56:03.016Z" } wheels = [ @@ -1383,7 +1405,8 @@ version = "2.1.0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] sdist = { url = "https://files.pythonhosted.org/packages/f2/97/ebf4da567aa6827c909642694d71c9fcf53e5b504f2d96afea02718862f3/iniconfig-2.1.0.tar.gz", hash = "sha256:3abbd2e30b36733fee78f9c7f7308f2d0050e88f0087fd25c2645f63c773e1c7", size = 4793, upload-time = "2025-03-19T20:09:59.721Z" } wheels = [ @@ -1409,7 +1432,8 @@ version = "8.18.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] dependencies = [ { name = "colorama", marker = "python_full_version < '3.10' and sys_platform == 'win32'" }, @@ -1632,6 +1656,43 @@ dev = [ { name = "pytest-xdist", specifier = ">=3.8.0" }, ] +[[package]] +name = "llama-agents-dbos" +version = "0.1.0" +source = { editable = "packages/llama-agents-dbos" } +dependencies = [ + { name = "dbos", marker = "python_full_version >= '3.10'" }, + { name = "llama-index-workflows" }, +] + +[package.dev-dependencies] +dev = [ + { name = "basedpyright" }, + { name = "pytest" }, + { name = "pytest-asyncio" }, + { name = "pytest-cov" }, + { name = "pytest-timeout" }, + { name = "pytest-xdist" }, + { name = "ty" }, +] + +[package.metadata] +requires-dist = [ + { name = "dbos", marker = "python_full_version >= '3.10'", specifier = ">=2.10.0" }, + { name = "llama-index-workflows", editable = "packages/llama-index-workflows" }, +] + +[package.metadata.requires-dev] +dev = [ + { name = "basedpyright", specifier = ">=1.31.1" }, + { name = "pytest", specifier = ">=8.4.0" }, + { name = "pytest-asyncio", specifier = ">=1.0.0" }, + { name = "pytest-cov", specifier = ">=6.1.1" }, + { name = "pytest-timeout", specifier = ">=2.4.0" }, + { name = "pytest-xdist", specifier = ">=3.0.0" }, + { name = "ty", specifier = ">=0.0.1,<0.0.9" }, +] + [[package]] name = "llama-agents-dev" version = "0.1.0" @@ -1685,6 +1746,7 @@ name = "llama-agents-integration-tests" version = "0.1.0" source = { editable = "packages/llama-agents-integration-tests" } dependencies = [ + { name = "llama-agents-dbos", marker = "python_full_version > '3.9'" }, { name = "llama-index-core" }, { name = "llama-index-workflows" }, ] @@ -1708,6 +1770,7 @@ dev = [ [package.metadata] requires-dist = [ + { name = "llama-agents-dbos", marker = "python_full_version > '3.9'", editable = "packages/llama-agents-dbos" }, { name = "llama-index-core", specifier = ">=0.14.13" }, { name = "llama-index-workflows", editable = "packages/llama-index-workflows" }, ] @@ -1934,7 +1997,8 @@ version = "3.9" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] dependencies = [ { name = "importlib-metadata", marker = "python_full_version < '3.10'" }, @@ -1963,7 +2027,8 @@ version = "3.0.0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] dependencies = [ { name = "mdurl", marker = "python_full_version < '3.10'" }, @@ -2310,7 +2375,8 @@ version = "1.18.2" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] dependencies = [ { name = "griffe", marker = "python_full_version < '3.10'" }, @@ -2531,7 +2597,8 @@ version = "3.2.1" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] sdist = { url = "https://files.pythonhosted.org/packages/c4/80/a84676339aaae2f1cfdf9f418701dd634aef9cc76f708ef55c36ff39c3ca/networkx-3.2.1.tar.gz", hash = "sha256:9f1bb5cf3409bf324e0a722c20bdb4c20ee39bf1c30ce8ae499c8502b0b5e0c6", size = 2073928, upload-time = "2023-10-28T08:41:39.364Z" } wheels = [ @@ -2609,7 +2676,8 @@ version = "2.0.2" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] sdist = { url = "https://files.pythonhosted.org/packages/a9/75/10dd1f8116a8b796cb2c737b674e02d02e80454bda953fa7e65d8c12b016/numpy-2.0.2.tar.gz", hash = "sha256:883c987dee1880e2a864ab0dc9892292582510604156762362d9326444636e78", size = 18902015, upload-time = "2024-08-26T20:19:40.945Z" } wheels = [ @@ -2862,7 +2930,8 @@ version = "11.3.0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] sdist = { url = "https://files.pythonhosted.org/packages/f3/0d/d0d6dea55cd152ce3d6767bb38a8fc10e33796ba4ba210cbab9354b6d238/pillow-11.3.0.tar.gz", hash = "sha256:3828ee7586cd0b2091b6209e5ad53e20d0649bbe87164a459d0676e035e8f523", size = 47113069, upload-time = "2025-07-01T09:16:30.666Z" } wheels = [ @@ -3081,7 +3150,8 @@ version = "4.4.0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] sdist = { url = "https://files.pythonhosted.org/packages/23/e8/21db9c9987b0e728855bd57bff6984f67952bea55d6f75e055c46b5383e8/platformdirs-4.4.0.tar.gz", hash = "sha256:ca753cf4d81dc309bc67b0ea38fd15dc97bc30ce419a7f58d13eb3bf14c4febf", size = 21634, upload-time = "2025-08-26T14:32:04.268Z" } wheels = [ @@ -3273,7 +3343,8 @@ version = "3.2.13" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] dependencies = [ { name = "typing-extensions", marker = "python_full_version < '3.10'" }, @@ -3317,7 +3388,8 @@ version = "3.2.13" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] wheels = [ { url = "https://files.pythonhosted.org/packages/8f/16/325f72b7ebdb906bd6cca6c0caea5b8fd7092c4686237c5669fe3f3cc7f2/psycopg_binary-3.2.13-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:9e25eb65494955c0dabdcd7097b004cbd70b982cf3cbc7186c2e854f788677a9", size = 4013642, upload-time = "2025-11-21T22:29:43.39Z" }, @@ -4048,7 +4120,8 @@ version = "3.3.3" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.9.2' and python_full_version < '3.10'", - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] dependencies = [ { name = "cryptography", marker = "python_full_version < '3.10'" }, @@ -4223,7 +4296,8 @@ name = "testcontainers" version = "4.13.0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version < '3.9.2'", + "python_full_version > '3.9' and python_full_version < '3.9.2'", + "python_full_version <= '3.9'", ] dependencies = [ { name = "docker", marker = "python_full_version < '3.9.2'" }, @@ -4565,6 +4639,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/38/1c/0d8454ff0f0f258737ecfe84f6e508729191d29663b404832f98fa5626b7/ty-0.0.8-py3-none-win_arm64.whl", hash = "sha256:ec74f022f315bede478ecae1277a01ab618e6500c1d68450d7883f5cd6ed554a", size = 9636374, upload-time = "2025-12-29T13:50:16.344Z" }, ] +[[package]] +name = "typer-slim" +version = "0.21.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "click", version = "8.3.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.10'" }, + { name = "typing-extensions", marker = "python_full_version >= '3.10'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/17/d4/064570dec6358aa9049d4708e4a10407d74c99258f8b2136bb8702303f1a/typer_slim-0.21.1.tar.gz", hash = "sha256:73495dd08c2d0940d611c5a8c04e91c2a0a98600cbd4ee19192255a233b6dbfd", size = 110478, upload-time = "2026-01-06T11:21:11.176Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c8/0a/4aca634faf693e33004796b6cee0ae2e1dba375a800c16ab8d3eff4bb800/typer_slim-0.21.1-py3-none-any.whl", hash = "sha256:6e6c31047f171ac93cc5a973c9e617dbc5ab2bddc4d0a3135dc161b4e2020e0d", size = 47444, upload-time = "2026-01-06T11:21:12.441Z" }, +] + [[package]] name = "typing-extensions" version = "4.15.0" @@ -4734,6 +4821,74 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/af/b5/123f13c975e9f27ab9c0770f514345bd406d0e8d3b7a0723af9d43f710af/wcwidth-0.2.14-py2.py3-none-any.whl", hash = "sha256:a7bb560c8aee30f9957e5f9895805edd20602f2d7f720186dfd906e82b4982e1", size = 37286, upload-time = "2025-09-22T16:29:51.641Z" }, ] +[[package]] +name = "websockets" +version = "16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/04/24/4b2031d72e840ce4c1ccb255f693b15c334757fc50023e4db9537080b8c4/websockets-16.0.tar.gz", hash = "sha256:5f6261a5e56e8d5c42a4497b364ea24d94d9563e8fbd44e78ac40879c60179b5", size = 179346, upload-time = "2026-01-10T09:23:47.181Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/20/74/221f58decd852f4b59cc3354cccaf87e8ef695fede361d03dc9a7396573b/websockets-16.0-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:04cdd5d2d1dacbad0a7bf36ccbcd3ccd5a30ee188f2560b7a62a30d14107b31a", size = 177343, upload-time = "2026-01-10T09:22:21.28Z" }, + { url = "https://files.pythonhosted.org/packages/19/0f/22ef6107ee52ab7f0b710d55d36f5a5d3ef19e8a205541a6d7ffa7994e5a/websockets-16.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:8ff32bb86522a9e5e31439a58addbb0166f0204d64066fb955265c4e214160f0", size = 175021, upload-time = "2026-01-10T09:22:22.696Z" }, + { url = "https://files.pythonhosted.org/packages/10/40/904a4cb30d9b61c0e278899bf36342e9b0208eb3c470324a9ecbaac2a30f/websockets-16.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:583b7c42688636f930688d712885cf1531326ee05effd982028212ccc13e5957", size = 175320, upload-time = "2026-01-10T09:22:23.94Z" }, + { url = "https://files.pythonhosted.org/packages/9d/2f/4b3ca7e106bc608744b1cdae041e005e446124bebb037b18799c2d356864/websockets-16.0-cp310-cp310-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:7d837379b647c0c4c2355c2499723f82f1635fd2c26510e1f587d89bc2199e72", size = 183815, upload-time = "2026-01-10T09:22:25.469Z" }, + { url = "https://files.pythonhosted.org/packages/86/26/d40eaa2a46d4302becec8d15b0fc5e45bdde05191e7628405a19cf491ccd/websockets-16.0-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:df57afc692e517a85e65b72e165356ed1df12386ecb879ad5693be08fac65dde", size = 185054, upload-time = "2026-01-10T09:22:27.101Z" }, + { url = "https://files.pythonhosted.org/packages/b0/ba/6500a0efc94f7373ee8fefa8c271acdfd4dca8bd49a90d4be7ccabfc397e/websockets-16.0-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:2b9f1e0d69bc60a4a87349d50c09a037a2607918746f07de04df9e43252c77a3", size = 184565, upload-time = "2026-01-10T09:22:28.293Z" }, + { url = "https://files.pythonhosted.org/packages/04/b4/96bf2cee7c8d8102389374a2616200574f5f01128d1082f44102140344cc/websockets-16.0-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:335c23addf3d5e6a8633f9f8eda77efad001671e80b95c491dd0924587ece0b3", size = 183848, upload-time = "2026-01-10T09:22:30.394Z" }, + { url = "https://files.pythonhosted.org/packages/02/8e/81f40fb00fd125357814e8c3025738fc4ffc3da4b6b4a4472a82ba304b41/websockets-16.0-cp310-cp310-win32.whl", hash = "sha256:37b31c1623c6605e4c00d466c9d633f9b812ea430c11c8a278774a1fde1acfa9", size = 178249, upload-time = "2026-01-10T09:22:32.083Z" }, + { url = "https://files.pythonhosted.org/packages/b4/5f/7e40efe8df57db9b91c88a43690ac66f7b7aa73a11aa6a66b927e44f26fa/websockets-16.0-cp310-cp310-win_amd64.whl", hash = "sha256:8e1dab317b6e77424356e11e99a432b7cb2f3ec8c5ab4dabbcee6add48f72b35", size = 178685, upload-time = "2026-01-10T09:22:33.345Z" }, + { url = "https://files.pythonhosted.org/packages/f2/db/de907251b4ff46ae804ad0409809504153b3f30984daf82a1d84a9875830/websockets-16.0-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:31a52addea25187bde0797a97d6fc3d2f92b6f72a9370792d65a6e84615ac8a8", size = 177340, upload-time = "2026-01-10T09:22:34.539Z" }, + { url = "https://files.pythonhosted.org/packages/f3/fa/abe89019d8d8815c8781e90d697dec52523fb8ebe308bf11664e8de1877e/websockets-16.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:417b28978cdccab24f46400586d128366313e8a96312e4b9362a4af504f3bbad", size = 175022, upload-time = "2026-01-10T09:22:36.332Z" }, + { url = "https://files.pythonhosted.org/packages/58/5d/88ea17ed1ded2079358b40d31d48abe90a73c9e5819dbcde1606e991e2ad/websockets-16.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:af80d74d4edfa3cb9ed973a0a5ba2b2a549371f8a741e0800cb07becdd20f23d", size = 175319, upload-time = "2026-01-10T09:22:37.602Z" }, + { url = "https://files.pythonhosted.org/packages/d2/ae/0ee92b33087a33632f37a635e11e1d99d429d3d323329675a6022312aac2/websockets-16.0-cp311-cp311-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:08d7af67b64d29823fed316505a89b86705f2b7981c07848fb5e3ea3020c1abe", size = 184631, upload-time = "2026-01-10T09:22:38.789Z" }, + { url = "https://files.pythonhosted.org/packages/c8/c5/27178df583b6c5b31b29f526ba2da5e2f864ecc79c99dae630a85d68c304/websockets-16.0-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:7be95cfb0a4dae143eaed2bcba8ac23f4892d8971311f1b06f3c6b78952ee70b", size = 185870, upload-time = "2026-01-10T09:22:39.893Z" }, + { url = "https://files.pythonhosted.org/packages/87/05/536652aa84ddc1c018dbb7e2c4cbcd0db884580bf8e95aece7593fde526f/websockets-16.0-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:d6297ce39ce5c2e6feb13c1a996a2ded3b6832155fcfc920265c76f24c7cceb5", size = 185361, upload-time = "2026-01-10T09:22:41.016Z" }, + { url = "https://files.pythonhosted.org/packages/6d/e2/d5332c90da12b1e01f06fb1b85c50cfc489783076547415bf9f0a659ec19/websockets-16.0-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:1c1b30e4f497b0b354057f3467f56244c603a79c0d1dafce1d16c283c25f6e64", size = 184615, upload-time = "2026-01-10T09:22:42.442Z" }, + { url = "https://files.pythonhosted.org/packages/77/fb/d3f9576691cae9253b51555f841bc6600bf0a983a461c79500ace5a5b364/websockets-16.0-cp311-cp311-win32.whl", hash = "sha256:5f451484aeb5cafee1ccf789b1b66f535409d038c56966d6101740c1614b86c6", size = 178246, upload-time = "2026-01-10T09:22:43.654Z" }, + { url = "https://files.pythonhosted.org/packages/54/67/eaff76b3dbaf18dcddabc3b8c1dba50b483761cccff67793897945b37408/websockets-16.0-cp311-cp311-win_amd64.whl", hash = "sha256:8d7f0659570eefb578dacde98e24fb60af35350193e4f56e11190787bee77dac", size = 178684, upload-time = "2026-01-10T09:22:44.941Z" }, + { url = "https://files.pythonhosted.org/packages/84/7b/bac442e6b96c9d25092695578dda82403c77936104b5682307bd4deb1ad4/websockets-16.0-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:71c989cbf3254fbd5e84d3bff31e4da39c43f884e64f2551d14bb3c186230f00", size = 177365, upload-time = "2026-01-10T09:22:46.787Z" }, + { url = "https://files.pythonhosted.org/packages/b0/fe/136ccece61bd690d9c1f715baaeefd953bb2360134de73519d5df19d29ca/websockets-16.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:8b6e209ffee39ff1b6d0fa7bfef6de950c60dfb91b8fcead17da4ee539121a79", size = 175038, upload-time = "2026-01-10T09:22:47.999Z" }, + { url = "https://files.pythonhosted.org/packages/40/1e/9771421ac2286eaab95b8575b0cb701ae3663abf8b5e1f64f1fd90d0a673/websockets-16.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:86890e837d61574c92a97496d590968b23c2ef0aeb8a9bc9421d174cd378ae39", size = 175328, upload-time = "2026-01-10T09:22:49.809Z" }, + { url = "https://files.pythonhosted.org/packages/18/29/71729b4671f21e1eaa5d6573031ab810ad2936c8175f03f97f3ff164c802/websockets-16.0-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:9b5aca38b67492ef518a8ab76851862488a478602229112c4b0d58d63a7a4d5c", size = 184915, upload-time = "2026-01-10T09:22:51.071Z" }, + { url = "https://files.pythonhosted.org/packages/97/bb/21c36b7dbbafc85d2d480cd65df02a1dc93bf76d97147605a8e27ff9409d/websockets-16.0-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e0334872c0a37b606418ac52f6ab9cfd17317ac26365f7f65e203e2d0d0d359f", size = 186152, upload-time = "2026-01-10T09:22:52.224Z" }, + { url = "https://files.pythonhosted.org/packages/4a/34/9bf8df0c0cf88fa7bfe36678dc7b02970c9a7d5e065a3099292db87b1be2/websockets-16.0-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:a0b31e0b424cc6b5a04b8838bbaec1688834b2383256688cf47eb97412531da1", size = 185583, upload-time = "2026-01-10T09:22:53.443Z" }, + { url = "https://files.pythonhosted.org/packages/47/88/4dd516068e1a3d6ab3c7c183288404cd424a9a02d585efbac226cb61ff2d/websockets-16.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:485c49116d0af10ac698623c513c1cc01c9446c058a4e61e3bf6c19dff7335a2", size = 184880, upload-time = "2026-01-10T09:22:55.033Z" }, + { url = "https://files.pythonhosted.org/packages/91/d6/7d4553ad4bf1c0421e1ebd4b18de5d9098383b5caa1d937b63df8d04b565/websockets-16.0-cp312-cp312-win32.whl", hash = "sha256:eaded469f5e5b7294e2bdca0ab06becb6756ea86894a47806456089298813c89", size = 178261, upload-time = "2026-01-10T09:22:56.251Z" }, + { url = "https://files.pythonhosted.org/packages/c3/f0/f3a17365441ed1c27f850a80b2bc680a0fa9505d733fe152fdf5e98c1c0b/websockets-16.0-cp312-cp312-win_amd64.whl", hash = "sha256:5569417dc80977fc8c2d43a86f78e0a5a22fee17565d78621b6bb264a115d4ea", size = 178693, upload-time = "2026-01-10T09:22:57.478Z" }, + { url = "https://files.pythonhosted.org/packages/cc/9c/baa8456050d1c1b08dd0ec7346026668cbc6f145ab4e314d707bb845bf0d/websockets-16.0-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:878b336ac47938b474c8f982ac2f7266a540adc3fa4ad74ae96fea9823a02cc9", size = 177364, upload-time = "2026-01-10T09:22:59.333Z" }, + { url = "https://files.pythonhosted.org/packages/7e/0c/8811fc53e9bcff68fe7de2bcbe75116a8d959ac699a3200f4847a8925210/websockets-16.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:52a0fec0e6c8d9a784c2c78276a48a2bdf099e4ccc2a4cad53b27718dbfd0230", size = 175039, upload-time = "2026-01-10T09:23:01.171Z" }, + { url = "https://files.pythonhosted.org/packages/aa/82/39a5f910cb99ec0b59e482971238c845af9220d3ab9fa76dd9162cda9d62/websockets-16.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:e6578ed5b6981005df1860a56e3617f14a6c307e6a71b4fff8c48fdc50f3ed2c", size = 175323, upload-time = "2026-01-10T09:23:02.341Z" }, + { url = "https://files.pythonhosted.org/packages/bd/28/0a25ee5342eb5d5f297d992a77e56892ecb65e7854c7898fb7d35e9b33bd/websockets-16.0-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:95724e638f0f9c350bb1c2b0a7ad0e83d9cc0c9259f3ea94e40d7b02a2179ae5", size = 184975, upload-time = "2026-01-10T09:23:03.756Z" }, + { url = "https://files.pythonhosted.org/packages/f9/66/27ea52741752f5107c2e41fda05e8395a682a1e11c4e592a809a90c6a506/websockets-16.0-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c0204dc62a89dc9d50d682412c10b3542d748260d743500a85c13cd1ee4bde82", size = 186203, upload-time = "2026-01-10T09:23:05.01Z" }, + { url = "https://files.pythonhosted.org/packages/37/e5/8e32857371406a757816a2b471939d51c463509be73fa538216ea52b792a/websockets-16.0-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:52ac480f44d32970d66763115edea932f1c5b1312de36df06d6b219f6741eed8", size = 185653, upload-time = "2026-01-10T09:23:06.301Z" }, + { url = "https://files.pythonhosted.org/packages/9b/67/f926bac29882894669368dc73f4da900fcdf47955d0a0185d60103df5737/websockets-16.0-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:6e5a82b677f8f6f59e8dfc34ec06ca6b5b48bc4fcda346acd093694cc2c24d8f", size = 184920, upload-time = "2026-01-10T09:23:07.492Z" }, + { url = "https://files.pythonhosted.org/packages/3c/a1/3d6ccdcd125b0a42a311bcd15a7f705d688f73b2a22d8cf1c0875d35d34a/websockets-16.0-cp313-cp313-win32.whl", hash = "sha256:abf050a199613f64c886ea10f38b47770a65154dc37181bfaff70c160f45315a", size = 178255, upload-time = "2026-01-10T09:23:09.245Z" }, + { url = "https://files.pythonhosted.org/packages/6b/ae/90366304d7c2ce80f9b826096a9e9048b4bb760e44d3b873bb272cba696b/websockets-16.0-cp313-cp313-win_amd64.whl", hash = "sha256:3425ac5cf448801335d6fdc7ae1eb22072055417a96cc6b31b3861f455fbc156", size = 178689, upload-time = "2026-01-10T09:23:10.483Z" }, + { url = "https://files.pythonhosted.org/packages/f3/1d/e88022630271f5bd349ed82417136281931e558d628dd52c4d8621b4a0b2/websockets-16.0-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:8cc451a50f2aee53042ac52d2d053d08bf89bcb31ae799cb4487587661c038a0", size = 177406, upload-time = "2026-01-10T09:23:12.178Z" }, + { url = "https://files.pythonhosted.org/packages/f2/78/e63be1bf0724eeb4616efb1ae1c9044f7c3953b7957799abb5915bffd38e/websockets-16.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:daa3b6ff70a9241cf6c7fc9e949d41232d9d7d26fd3522b1ad2b4d62487e9904", size = 175085, upload-time = "2026-01-10T09:23:13.511Z" }, + { url = "https://files.pythonhosted.org/packages/bb/f4/d3c9220d818ee955ae390cf319a7c7a467beceb24f05ee7aaaa2414345ba/websockets-16.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:fd3cb4adb94a2a6e2b7c0d8d05cb94e6f1c81a0cf9dc2694fb65c7e8d94c42e4", size = 175328, upload-time = "2026-01-10T09:23:14.727Z" }, + { url = "https://files.pythonhosted.org/packages/63/bc/d3e208028de777087e6fb2b122051a6ff7bbcca0d6df9d9c2bf1dd869ae9/websockets-16.0-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:781caf5e8eee67f663126490c2f96f40906594cb86b408a703630f95550a8c3e", size = 185044, upload-time = "2026-01-10T09:23:15.939Z" }, + { url = "https://files.pythonhosted.org/packages/ad/6e/9a0927ac24bd33a0a9af834d89e0abc7cfd8e13bed17a86407a66773cc0e/websockets-16.0-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:caab51a72c51973ca21fa8a18bd8165e1a0183f1ac7066a182ff27107b71e1a4", size = 186279, upload-time = "2026-01-10T09:23:17.148Z" }, + { url = "https://files.pythonhosted.org/packages/b9/ca/bf1c68440d7a868180e11be653c85959502efd3a709323230314fda6e0b3/websockets-16.0-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:19c4dc84098e523fd63711e563077d39e90ec6702aff4b5d9e344a60cb3c0cb1", size = 185711, upload-time = "2026-01-10T09:23:18.372Z" }, + { url = "https://files.pythonhosted.org/packages/c4/f8/fdc34643a989561f217bb477cbc47a3a07212cbda91c0e4389c43c296ebf/websockets-16.0-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:a5e18a238a2b2249c9a9235466b90e96ae4795672598a58772dd806edc7ac6d3", size = 184982, upload-time = "2026-01-10T09:23:19.652Z" }, + { url = "https://files.pythonhosted.org/packages/dd/d1/574fa27e233764dbac9c52730d63fcf2823b16f0856b3329fc6268d6ae4f/websockets-16.0-cp314-cp314-win32.whl", hash = "sha256:a069d734c4a043182729edd3e9f247c3b2a4035415a9172fd0f1b71658a320a8", size = 177915, upload-time = "2026-01-10T09:23:21.458Z" }, + { url = "https://files.pythonhosted.org/packages/8a/f1/ae6b937bf3126b5134ce1f482365fde31a357c784ac51852978768b5eff4/websockets-16.0-cp314-cp314-win_amd64.whl", hash = "sha256:c0ee0e63f23914732c6d7e0cce24915c48f3f1512ec1d079ed01fc629dab269d", size = 178381, upload-time = "2026-01-10T09:23:22.715Z" }, + { url = "https://files.pythonhosted.org/packages/06/9b/f791d1db48403e1f0a27577a6beb37afae94254a8c6f08be4a23e4930bc0/websockets-16.0-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:a35539cacc3febb22b8f4d4a99cc79b104226a756aa7400adc722e83b0d03244", size = 177737, upload-time = "2026-01-10T09:23:24.523Z" }, + { url = "https://files.pythonhosted.org/packages/bd/40/53ad02341fa33b3ce489023f635367a4ac98b73570102ad2cdd770dacc9a/websockets-16.0-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:b784ca5de850f4ce93ec85d3269d24d4c82f22b7212023c974c401d4980ebc5e", size = 175268, upload-time = "2026-01-10T09:23:25.781Z" }, + { url = "https://files.pythonhosted.org/packages/74/9b/6158d4e459b984f949dcbbb0c5d270154c7618e11c01029b9bbd1bb4c4f9/websockets-16.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:569d01a4e7fba956c5ae4fc988f0d4e187900f5497ce46339c996dbf24f17641", size = 175486, upload-time = "2026-01-10T09:23:27.033Z" }, + { url = "https://files.pythonhosted.org/packages/e5/2d/7583b30208b639c8090206f95073646c2c9ffd66f44df967981a64f849ad/websockets-16.0-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:50f23cdd8343b984957e4077839841146f67a3d31ab0d00e6b824e74c5b2f6e8", size = 185331, upload-time = "2026-01-10T09:23:28.259Z" }, + { url = "https://files.pythonhosted.org/packages/45/b0/cce3784eb519b7b5ad680d14b9673a31ab8dcb7aad8b64d81709d2430aa8/websockets-16.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:152284a83a00c59b759697b7f9e9cddf4e3c7861dd0d964b472b70f78f89e80e", size = 186501, upload-time = "2026-01-10T09:23:29.449Z" }, + { url = "https://files.pythonhosted.org/packages/19/60/b8ebe4c7e89fb5f6cdf080623c9d92789a53636950f7abacfc33fe2b3135/websockets-16.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:bc59589ab64b0022385f429b94697348a6a234e8ce22544e3681b2e9331b5944", size = 186062, upload-time = "2026-01-10T09:23:31.368Z" }, + { url = "https://files.pythonhosted.org/packages/88/a8/a080593f89b0138b6cba1b28f8df5673b5506f72879322288b031337c0b8/websockets-16.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:32da954ffa2814258030e5a57bc73a3635463238e797c7375dc8091327434206", size = 185356, upload-time = "2026-01-10T09:23:32.627Z" }, + { url = "https://files.pythonhosted.org/packages/c2/b6/b9afed2afadddaf5ebb2afa801abf4b0868f42f8539bfe4b071b5266c9fe/websockets-16.0-cp314-cp314t-win32.whl", hash = "sha256:5a4b4cc550cb665dd8a47f868c8d04c8230f857363ad3c9caf7a0c3bf8c61ca6", size = 178085, upload-time = "2026-01-10T09:23:33.816Z" }, + { url = "https://files.pythonhosted.org/packages/9f/3e/28135a24e384493fa804216b79a6a6759a38cc4ff59118787b9fb693df93/websockets-16.0-cp314-cp314t-win_amd64.whl", hash = "sha256:b14dc141ed6d2dde437cddb216004bcac6a1df0935d79656387bd41632ba0bbd", size = 178531, upload-time = "2026-01-10T09:23:35.016Z" }, + { url = "https://files.pythonhosted.org/packages/72/07/c98a68571dcf256e74f1f816b8cc5eae6eb2d3d5cfa44d37f801619d9166/websockets-16.0-pp311-pypy311_pp73-macosx_10_15_x86_64.whl", hash = "sha256:349f83cd6c9a415428ee1005cadb5c2c56f4389bc06a9af16103c3bc3dcc8b7d", size = 174947, upload-time = "2026-01-10T09:23:36.166Z" }, + { url = "https://files.pythonhosted.org/packages/7e/52/93e166a81e0305b33fe416338be92ae863563fe7bce446b0f687b9df5aea/websockets-16.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:4a1aba3340a8dca8db6eb5a7986157f52eb9e436b74813764241981ca4888f03", size = 175260, upload-time = "2026-01-10T09:23:37.409Z" }, + { url = "https://files.pythonhosted.org/packages/56/0c/2dbf513bafd24889d33de2ff0368190a0e69f37bcfa19009ef819fe4d507/websockets-16.0-pp311-pypy311_pp73-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:f4a32d1bd841d4bcbffdcb3d2ce50c09c3909fbead375ab28d0181af89fd04da", size = 176071, upload-time = "2026-01-10T09:23:39.158Z" }, + { url = "https://files.pythonhosted.org/packages/a5/8f/aea9c71cc92bf9b6cc0f7f70df8f0b420636b6c96ef4feee1e16f80f75dd/websockets-16.0-pp311-pypy311_pp73-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0298d07ee155e2e9fda5be8a9042200dd2e3bb0b8a38482156576f863a9d457c", size = 176968, upload-time = "2026-01-10T09:23:41.031Z" }, + { url = "https://files.pythonhosted.org/packages/9a/3f/f70e03f40ffc9a30d817eef7da1be72ee4956ba8d7255c399a01b135902a/websockets-16.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:a653aea902e0324b52f1613332ddf50b00c06fdaf7e92624fbf8c77c78fa5767", size = 178735, upload-time = "2026-01-10T09:23:42.259Z" }, + { url = "https://files.pythonhosted.org/packages/6f/28/258ebab549c2bf3e64d2b0217b973467394a9cea8c42f70418ca2c5d0d2e/websockets-16.0-py3-none-any.whl", hash = "sha256:1637db62fad1dc833276dded54215f2c7fa46912301a24bd94d45d46a011ceec", size = 171598, upload-time = "2026-01-10T09:23:45.395Z" }, +] + [[package]] name = "wrapt" version = "2.0.0"