跳到内容

工作流子类化#

工作流的另一个很棒的特性是它们的可扩展性。你可以使用其他人编写的工作流或 LlamaIndex 的内置工作流,并对其进行扩展以根据你的需求进行自定义。我们将介绍两种实现方法。

第一种是子类化:工作流只是普通的 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)

如果我们运行它

w = MainWorkflow(timeout=10, verbose=False)
result = await w.run(query="Initial query")
print(result)

我们得到

Starting up
Sending an email
Finishing up
Initial query

现在,让我们对这个工作流进行子类化,以便也能发送短信

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)

这将转而给我们

Starting up
Sending an email
Also sending a text message
Finishing up
Initial query

我们可以可视化子类化的工作流,它将显示所有步骤,如下所示

draw_all_possible_flows(CustomWorkflow, "custom_workflow.html")

Custom workflow

接下来,让我们看看扩展工作流的另一种方法:嵌套工作流