[Question]: step function outside workflow class

[AnnPe]

Question Validation

• I have searched both the documentation and discord for an answer.

Question

is it possible to define step functions outside workflow class? so that they can be used as lego pieces - to easily create different flows for experiments?

[dosubot]

Hey @AnnPe! How can I assist you today? I'm here to help you with any bugs, questions, or contributions while you wait for a human maintainer.

Yes, it is possible to define step functions outside the workflow class so that they can be used as modular components to easily create different flows for experiments. You can define these step functions as free functions and use the @step decorator to associate them with a specific workflow class. This allows for modular and reusable step functions that can be added to any existing workflow class without needing to subclass it ^[1][2].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[logan-markewich]

Yes! However, this is undocumented for now since I feel the UX isn't great, but we weren't really sure of a solid use case or how people would use this.

llama_index/llama-index-core/llama_index/core/workflow/workflow.py

Line 159 in c80992f

def add_step(cls, func: Callable) -> None:

Basically the UX is something like

w = Workflow()

@step(workflow=w)
def my step(ev: StartEvent) -> StopEvent:
  return StopEvent()

w.add_step(my_step)

Im realizing already how janky that is haha, but it's there

[AnnPe]

hi, thanks for your quick reply!
im sorry for not being clear enough, my use case is the following:
i want to have predefined step functions with some parameters, and to be able to combine them as lego pieces - to create workflows for experiments. your example still allocates step function to a certain class (@step(workflow = w)) - but what if i want to create 2 different workflows and compare them. is there a built-in way or recommended workaround?

thanks!

[AnnPe]

@dosu

i need to create them w/o associating them to a specific class - so that i can use them as building blocks for different workflows. is that possible?

[dosubot]

Yes, you can define step functions without immediately associating them with a specific workflow instance. You can use the @step decorator to mark functions as workflow steps and later associate them with different workflows. When using the decorator on free functions, you can specify the workflow parameter to associate the step with a particular workflow class later. This allows you to define step functions independently and use them as building blocks for different workflows.

Here's an example of how you can define a step function without associating it with a specific workflow initially:

@step
def my_step_function(event: Event) -> Event:
    # Define your step logic here
    return some_event

Later, you can associate this step function with different workflows as needed. If you don't specify a workflow class when using the @step decorator on a free function, it will raise a WorkflowValidationError, so make sure to associate it with a workflow when you're ready to use it ^[1][2].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[logan-markewich]

@AnnPe it's tied to the instance w, not the class, so if you want more, just make another variable

w1 = Workflow()
w2 = Workflow()

[AnnPe]

yeah, to the instance of course - but this means that i have to define another function to tie it to another workflow, don't i?
or do i misunderstand smth?

[AnnPe]

@dosu
that's great, but afaik, i have to use specific Event-inherited class to link those steps together inside a workflow don't I?
how to link these defined-outside-class step functions together?

[dosubot]

To link step functions defined outside of a class together within a workflow using specific Event-inherited classes, you can define a workflow and use the @step decorator to attach each step function to the workflow. Each step function should specify the type of event it consumes and the type of event it produces. Here's an example:

from llama_index.core.workflow import (
    Event,
    StartEvent,
    StopEvent,
    Workflow,
    step,
)

class MyStart(StartEvent):
    query: str

class MyStop(StopEvent):
    outcome: str

# Define step functions without associating them with a specific workflow initially
@step
async def start_step(ev: MyStart) -> Event:
    # Your logic here
    return Event()  # Replace with actual event

@step
async def middle_step(ev: Event) -> Event:
    # Your logic here
    return Event()  # Replace with actual event

@step
async def end_step(ev: Event) -> MyStop:
    # Your logic here
    return MyStop(outcome="Workflow completed")

# Create workflow instances and associate steps
w1 = Workflow()
w1.add_step(start_step)
w1.add_step(middle_step)
w1.add_step(end_step)

w2 = Workflow()
w2.add_step(start_step)
w2.add_step(middle_step)
w2.add_step(end_step)

# Example of running a workflow
handler = w1.run(stepwise=True)
while ev := await handler.run_step():
    handler.ctx.send_event(ev)

result = await handler

In this setup, MyStart and MyStop are custom events that inherit from StartEvent and StopEvent, respectively. The step functions are defined independently and can be added to any workflow instance. This allows you to reuse the same step functions across different workflows ^[1][2][3].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[AnnPe]

@dosu, great
and how can i add Context parameter to track variables?

[dosubot]

To add a Context parameter to track variables in a workflow, define a tool with the Context parameter as the first parameter. Here's an example of how to add a name to the state:

from llama_index.core.workflow import Context

async def set_name(ctx: Context, name: str) -> str:
    state = await ctx.get("state")
    state["name"] = name
    await ctx.set("state", state)
    return f"Name set to {name}"

Create an AgentWorkflow using this tool and optionally provide an initial state:

workflow = AgentWorkflow.from_tools_or_functions(
    [set_name],
    llm=llm,
    system_prompt="You are a helpful assistant that can set a name.",
    initial_state={"name": "unset"},
)

Create a Context and use it to run the workflow:

ctx = Context(workflow)
response = await workflow.run(user_msg="My name is Logan", ctx=ctx)
print(str(response))

state = await ctx.get("state")
print(state["name"])

The Context object stores data that needs to be available across iterations during a workflow execution and across multiple workflow runs. It offers both global and private data storage, making it easier to maintain state across different steps without needing to pass data explicitly between them ^[1][2][3][4].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}

[AnnPe]

To link step functions defined outside of a class together within a workflow using specific Event-inherited classes, you can define a workflow and use the @step decorator to attach each step function to the workflow. Each step function should specify the type of event it consumes and the type of event it produces. Here's an example:

from llama_index.core.workflow import (
Event,
StartEvent,
StopEvent,
Workflow,
step,
)

class MyStart(StartEvent):
query: str

class MyStop(StopEvent):
outcome: str

Define step functions without associating them with a specific workflow initially

@step
async def start_step(ev: MyStart) -> Event:
# Your logic here
return Event() # Replace with actual event

@step
async def middle_step(ev: Event) -> Event:
# Your logic here
return Event() # Replace with actual event

@step
async def end_step(ev: Event) -> MyStop:
# Your logic here
return MyStop(outcome="Workflow completed")

Create workflow instances and associate steps

w1 = Workflow()
w1.add_step(start_step)
w1.add_step(middle_step)
w1.add_step(end_step)

w2 = Workflow()
w2.add_step(start_step)
w2.add_step(middle_step)
w2.add_step(end_step)

Example of running a workflow

handler = w1.run(stepwise=True)
while ev := await handler.run_step():
handler.ctx.send_event(ev)

result = await handler
In this setup, MyStart and MyStop are custom events that inherit from StartEvent and StopEvent, respectively. The step functions are defined independently and can be added to any workflow instance. This allows you to reuse the same step functions across different workflows [1][2][3].

To continue talking to Dosu, mention @dosu.

Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other

@dosu,

you created two identical workflows, didn't you?
i want to create different workflows (probably with loops and conditions)
how shall i deal with input/output Events for each step defined outside?

[dosubot]

To create different workflows with loops and conditions using input/output events for each step defined outside a class, you can define custom event types and use them in your workflow steps. Here's how you can do it:

1. Define Custom Event Types: Create event classes for each step and condition in your workflow.

class LoopEvent(Event):
    loop_output: str

class BranchA1Event(Event):
    payload: str

class BranchA2Event(Event):
    payload: str

class BranchB1Event(Event):
    payload: str

class BranchB2Event(Event):
    payload: str

2. Implement Workflow Steps: Use these event types in your workflow steps to handle loops and branches. Define these steps as free functions and add them to your workflow using the add_step method.

@step
async def step_one(ev: StartEvent | LoopEvent) -> FirstEvent | LoopEvent:
    if random.randint(0, 1) == 0:
        print("Bad thing happened")
        return LoopEvent(loop_output="Back to step one.")
    else:
        print("Good thing happened")
        return FirstEvent(first_output="First step complete.")

@step
async def start(ev: StartEvent) -> BranchA1Event | BranchB1Event:
    if random.randint(0, 1) == 0:
        print("Go to branch A")
        return BranchA1Event(payload="Branch A")
    else:
        print("Go to branch B")
        return BranchB1Event(payload="Branch B")

@step
async def step_a1(ev: BranchA1Event) -> BranchA2Event:
    print(ev.payload)
    return BranchA2Event(payload=ev.payload)

@step
async def step_b1(ev: BranchB1Event) -> BranchB2Event:
    print(ev.payload)
    return BranchB2Event(payload=ev.payload)

@step
async def step_a2(ev: BranchA2Event) -> StopEvent:
    print(ev.payload)
    return StopEvent(result="Branch A complete.")

@step
async def step_b2(ev: BranchB2Event) -> StopEvent:
    print(ev.payload)
    return StopEvent(result="Branch B complete.")

3. Add Steps to Workflow: Use the add_step method to add these free functions as steps to your workflow instance.

workflow = Workflow()
workflow.add_step(step_one)
workflow.add_step(start)
workflow.add_step(step_a1)
workflow.add_step(step_b1)
workflow.add_step(step_a2)
workflow.add_step(step_b2)

4. Running the Workflow: Since workflows are async by default, use await to get the result of the run command. In a vanilla Python script, wrap your code in an async function and use asyncio.run() to execute it.

async def main():
    result = await workflow.run()
    print(result)

if __name__ == "__main__":
    import asyncio
    asyncio.run(main())

5. Visualizing the Workflow: You can visualize the workflow using the built-in visualizer to see all possible flows.

from llama_index.utils.workflow import draw_all_possible_flows

draw_all_possible_flows(workflow, filename="workflow_visualization.html")

This setup allows you to create workflows with loops and branches by defining the appropriate event types and using them to control the flow of your application ^[1].

To continue talking to Dosu, mention @dosu.

---

^{Help Dosu learn! Give it feedback: Great Response | Irrelevant Answer | Incorrect Sources | Too Verbose | Hallucination | Bug Report | Other}