workflows-py/examples/durable_workflows.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7bHdUvTD1MGa"
   },
   "source": [
    "Workflows are ephemeral by default, meaning that once the `run()` method returns its result, the workflow state is lost. A subsequent call to `run()` on the same workflow instance will start from a fresh state.\n",
    "\n",
    "If the use case requires to persist the workflow state  across multiple runs and possibly different processes, there are a few strategies that can be used to make workflows more durable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "N0PWAEUE1AN2"
   },
   "outputs": [],
   "source": [
    "!pip install llama-index-workflows"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "10SuevfQ1Suz"
   },
   "source": [
    "## Storing data in the workflow instance\n",
    "\n",
    "Workflows are regular Python classes, and data can be stored in class or instance variables, so that subsequent `run()` invocations can access it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "vT5JIA5V1Ugx",
    "outputId": "37ade20e-534a-4049-b67b-dcab70ba98d4"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The step ran 1 times\n",
      "The step ran 2 times\n",
      "The step ran 3 times\n"
     ]
    }
   ],
   "source": [
    "from workflows import Workflow, step\n",
    "from workflows.events import StartEvent, StopEvent\n",
    "\n",
    "\n",
    "class MyWorkflow(Workflow):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        self.counter = 0\n",
    "        super().__init__(*args, **kwargs)\n",
    "\n",
    "    @step\n",
    "    def count(self, ev: StartEvent) -> StopEvent:\n",
    "        self.counter += 1\n",
    "        return StopEvent(result=f\"The step ran {self.counter} times\")\n",
    "\n",
    "\n",
    "w = MyWorkflow()\n",
    "for _ in range(3):\n",
    "    print(await w.run())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "HhIoDaeT22n5"
   },
   "source": [
    "## Storing data in the context object\n",
    "\n",
    "Each workflow comes with a special object responsible for its runtime operations called `Context`. The context instance is available to any step of a workflow and comes with a `store` property that can be used to store and load state data. Using the state store has two major advantages compared to class and instance variables:\n",
    "\n",
    "- It’s async safe and supports concurrent access\n",
    "- It can be serialized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "w5x_Z4-_5vGd",
    "outputId": "c75afe24-3aee-44e7-fe8e-743941055402"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The step ran 1 times\n",
      "The step ran 2 times\n"
     ]
    }
   ],
   "source": [
    "from workflows import Context, Workflow, step\n",
    "from workflows.events import StartEvent, StopEvent\n",
    "\n",
    "\n",
    "class MyWorkflow(Workflow):\n",
    "    @step\n",
    "    async def count(self, ctx: Context, ev: StartEvent) -> StopEvent:\n",
    "        async with ctx.store.edit_state() as state:\n",
    "            counter = state.get(\"counter\", 1)\n",
    "            retval = StopEvent(result=f\"The step ran {counter} times\")\n",
    "            state[\"counter\"] = counter + 1\n",
    "        return retval\n",
    "\n",
    "\n",
    "w = MyWorkflow()\n",
    "handler = w.run()\n",
    "print(await handler)\n",
    "\n",
    "w = MyWorkflow()\n",
    "handler = w.run(ctx=handler.ctx)\n",
    "print(await handler)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BDB7BAHB7iY_"
   },
   "source": [
    "## Using external resources to checkpoint execution\n",
    "\n",
    "To avoid any overhead, workflows don’t take snapshots of the current state automatically, so they can’t survive a fatal error on their own. However, any step can rely on some external database like Redis and snapshot the current context on sensitive parts of the code.\n",
    "\n",
    "For example, given a long running workflow processing hundreds of documents, we could save the id of the last document successfully processed in the state store:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "42MML8XG7kLT",
    "outputId": "773ad015-a40d-4d61-dc31-191dfa675526"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The step ran 1 times\n",
      "The step ran 2 times\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import sqlite3\n",
    "from typing import Annotated\n",
    "\n",
    "from workflows import Context, Workflow, step\n",
    "from workflows.context import JsonSerializer\n",
    "from workflows.events import StartEvent, StopEvent\n",
    "from workflows.resource import Resource\n",
    "\n",
    "\n",
    "def get_db() -> sqlite3.Connection:\n",
    "    return sqlite3.connect(\"mydb.db\")\n",
    "\n",
    "\n",
    "class MyWorkflow(Workflow):\n",
    "    @step\n",
    "    async def count(\n",
    "        self,\n",
    "        ctx: Context,\n",
    "        ev: StartEvent,\n",
    "        db: Annotated[sqlite3.Connection, Resource(get_db)],\n",
    "    ) -> StopEvent:\n",
    "        async with ctx.store.edit_state() as state:\n",
    "            counter = state.get(\"counter\", 1)\n",
    "            retval = StopEvent(result=f\"The step ran {counter} times\")\n",
    "            state[\"counter\"] = counter + 1\n",
    "\n",
    "        cursor = db.cursor()\n",
    "        ctx_dict = ctx.to_dict(serializer=JsonSerializer())\n",
    "        cursor.execute(\n",
    "            \"INSERT OR REPLACE INTO state VALUES (?, ?)\",\n",
    "            (\"last_ctx\", json.dumps(ctx_dict)),\n",
    "        )\n",
    "        db.commit()\n",
    "\n",
    "        return retval\n",
    "\n",
    "\n",
    "# Create a simple key-value table\n",
    "db = get_db()\n",
    "db.cursor().execute(\n",
    "    \"CREATE TABLE IF NOT EXISTS state (key TEXT PRIMARY KEY, value TEXT)\"\n",
    ")\n",
    "db.commit()\n",
    "\n",
    "\n",
    "w = MyWorkflow()\n",
    "print(await w.run())\n",
    "\n",
    "# State is stored in a DB now, we could restart the process here...\n",
    "\n",
    "w = MyWorkflow()\n",
    "cursor = db.cursor()\n",
    "cursor.execute(\"SELECT value FROM state WHERE key=?\", (\"last_ctx\",))\n",
    "ctx_json = cursor.fetchone()[0]\n",
    "restored_ctx = Context.from_dict(w, json.loads(ctx_json), serializer=JsonSerializer())\n",
    "print(await w.run(ctx=restored_ctx))"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "name": "python3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}