Extending workflows (#15573)

docs/docs/understanding/workflows/concurrent_execution.md

@@ -122,4 +122,4 @@ The visualization of this workflow is quite pleasing:
 
 ![A concurrent workflow](./different_events.png)
 
-Finally, let's take a look at [observability and debugging](observability.md) in workflows.
+Now let's look at how [we can extend workflows](extending.md).

docs/docs/understanding/workflows/nested.md

@@ -0,0 +1,61 @@
+# Nested workflows
+
+Another way to extend workflows is to nest additional workflows. It's possible to create explicit slots in existing flows where you can supply an entire additional workflow. For example, let's say we had a query that used an LLM to reflect on the quality of that query. The author might expect that you would want to modify the reflection step, and leave a slot for you to do that.
+
+Here's our base workflow:
+
+```python
+from llama_index.core.workflow import (
+    StartEvent,
+    StopEvent,
+    Workflow,
+    step,
+    Event,
+    Context,
+)
+from llama_index.utils.workflow import draw_all_possible_flows
+
+
+class Step2Event(Event):
+    query: str
+
+
+class MainWorkflow(Workflow):
+    @step
+    async def start(
+        self, ctx: Context, ev: StartEvent, reflection_workflow: Workflow
+    ) -> Step2Event:
+        print("Need to run reflection")
+        res = await reflection_workflow.run(query=ev.query)
+
+        return Step2Event(query=res)
+
+    @step
+    async def step_two(self, ctx: Context, ev: Step2Event) -> StopEvent:
+        print("Query is ", ev.query)
+        # do something with the query here
+        return StopEvent(result=ev.query)
+```
+
+This workflow by itself will not run; it needs a valid workflow for the reflection step. Let's create one:
+
+```python
+class ReflectionFlow(Workflow):
+    @step
+    async def sub_start(self, ctx: Context, ev: StartEvent) -> StopEvent:
+        print("Doing custom reflection")
+        return StopEvent(result="Improved query")
+```
+
+Now we can run the main workflow by supplying this custom reflection nested flow using the `add_workflows` method, to which we pass an instance of the `ReflectionFlow` class:
+
+```python
+w = MainWorkflow(timeout=10, verbose=False)
+w.add_workflows(reflection_workflow=ReflectionFlow())
+result = await w.run(query="Initial query")
+print(result)
+```
+
+Note that because the nested flow is a totally different workflow rather than a step, `draw_all_possible_flows` will only draw the flow of `MainWorkflow`.
+
+Finally, let's take a look at [observability and debugging](observability.md) in workflows.

docs/docs/understanding/workflows/state.md

@@ -2,9 +2,9 @@
 
 In our examples so far, we have passed data from step to step using properties of custom events. This is a powerful way to pass data around, but it has limitations. For example, if you want to pass data between steps that are not directly connected, you need to pass the data through all the steps in between. This can make your code harder to read and maintain.
 
-To avoid this pitfall, we have a `Context` object available to every step in the workflow. To use it, you must explicitly tell the step decorator to pass the context to the step. Here's how you do that.
+To avoid this pitfall, we have a `Context` object available to every step in the workflow. To use it, declare an argument of type `Context` to your step. Here's how you do that.
 
-We need one new import, the `Context`:
+We need one new import, the `Context` type:
 
 ```python
 from llama_index.core.workflow import (

docs/docs/understanding/workflows/subclass.md

@@ -0,0 +1,99 @@
+# Subclassing workflows
+
+Another great feature of workflows is their extensibility. You can take workflows written by others or built-ins from LlamaIndex and extend them to customize them to your needs. We'll look at two ways to do that.
+
+The first is subclassing: workflows are just regular Python classes, which means you can subclass them to add new functionality. For example, let's say you have an agentic workflow that does some processing and then sends an email. You can subclass the workflow to add an extra step to send a text message as well.
+
+Here's our base workflow:
+
+```python
+from llama_index.core.workflow import (
+    StartEvent,
+    StopEvent,
+    Workflow,
+    step,
+    Event,
+    Context,
+)
+
+
+class Step2Event(Event):
+    query: str
+
+
+class Step3Event(Event):
+    query: str
+
+
+class MainWorkflow(Workflow):
+    @step
+    async def start(self, ev: StartEvent) -> Step2Event:
+        print("Starting up")
+        return Step2Event(query=ev.query)
+
+    @step
+    async def step_two(self, ev: Step2Event) -> Step3Event:
+        print("Sending an email")
+        return Step3Event(query=ev.query)
+
+    @step
+    async def step_three(self, ev: Step3Event) -> StopEvent:
+        print("Finishing up")
+        return StopEvent(result=ev.query)
+```
+
+If we run this:
+
+```python
+w = MainWorkflow(timeout=10, verbose=False)
+result = await w.run(query="Initial query")
+print(result)
+```
+
+We get:
+
+```
+Starting up
+Sending an email
+Finishing up
+Initial query
+```
+
+Now let's subclass this workflow to send a text message as well:
+
+```python
+class Step2BEvent(Event):
+    query: str
+
+
+class CustomWorkflow(MainWorkflow):
+    @step
+    async def step_two(self, ev: Step2Event) -> Step2BEvent:
+        print("Sending an email")
+        return Step2BEvent(query=ev.query)
+
+    @step
+    async def step_two_b(self, ev: Step2BEvent) -> Step3Event:
+        print("Also sending a text message")
+        return Step3Event(query=ev.query)
+```
+
+Which will instead give us
+
+```
+Starting up
+Sending an email
+Also sending a text message
+Finishing up
+Initial query
+```
+
+We can visualize the subclassed workflow and it will show all the steps, like this:
+
+```python
+draw_all_possible_flows(CustomWorkflow, "custom_workflow.html")
+```
+
+![Custom workflow](subclass.png)
+
+Next, let's look at [nested workflows](nested.md).

docs/mkdocs.yml

@@ -53,6 +53,8 @@ nav:
           - Branches and loops: ./understanding/workflows/branches_and_loops.md
           - Maintaining state: ./understanding/workflows/state.md
           - Concurrent execution: ./understanding/workflows/concurrent_execution.md
+          - Subclassing workflows: ./understanding/workflows/subclass.md
+          - Nested workflows: ./understanding/workflows/nested.md
           - Observability: ./understanding/workflows/observability.md
       - Tracing and Debugging: ./understanding/tracing_and_debugging/tracing_and_debugging.md
       - Evaluating: