Need for Workflow Management in AI SDK Library

[hashiromer]

So, I came across this and it sounds perfect for the application I am building except managing workflows. I am copying the agent example from documentation to show what am I looking for and what seems to be currently missing.

import { openai } from '@ai-sdk/openai';
import { generateText, tool } from 'ai';
import * as mathjs from 'mathjs';
import { z } from 'zod';

const { text: answer } = await generateText({
  model: openai('gpt-4o-2024-08-06', { structuredOutputs: true }),
  tools: {
    calculate: tool({
      description:
        'A tool for evaluating mathematical expressions. ' +
        'Example expressions: ' +
        "'1.2 * (2 + 4.5)', '12.7 cm to inch', 'sin(45 deg) ^ 2'.",
      parameters: z.object({ expression: z.string() }),
      execute: async ({ expression }) => mathjs.evaluate(expression),
    }),
  },
  maxSteps: 10,
  system:
    'You are solving math problems. ' +
    'Reason step by step. ' +
    'Use the calculator when necessary. ' +
    'When you give the final answer, ' +
    'provide an explanation for how you arrived at it.',
  prompt:
    'A taxi driver earns $9461 per 1-hour of work. ' +
    'If he works 12 hours a day and in 1 hour ' +
    'he uses 12 liters of petrol with a price  of $134 for 1 liter. ' +
    'How much money does he earn in one day?',
});

console.log(`ANSWER: ${answer}`);

In this example, we are providing tools to LLM and it can choose to use any of them at its discretion based on chat history, we do not have any control when does the LLM call each tool nor can we control the context provided to LLM for each tool call which can contain a lot of noise if the tool generates a lot of output irrelevant to LLM.

It is a huge problem for any moderate complex AI tool. A lot of times, we need to use the provided tools in a specific order and at other times rely on LLMs. For example, different variations of RAG will start with converting the text into embedding, check some sort of similarity against stored embeddings in a vector DB and rank them followed by retrieval and then generating the final response. All of these tools need to be called in a fixed sequence. If it is left to LLMs, even if it makes one request out of order, it will lead to unnecessary errors which can be avoided. This is also a very simple example, a lot of times, we need to use 10 or 20 tools in a fixed order for complex tasks.

These workflows also frequently require conditional branches, so a RAG tool for example can use an LLM judge or a human in the loop to determine if the response is complete or accurate, if it thinks its not accurate, we can run the above sequence of actions again and again until the judge determines the output is acceptable. So, this sort of conditional branches is also necessary.

There are lots of things like these which are necessary to automate multi step workflows for businesses and this is a solved problem but unfortunately only in Python. I was researching a JavaScript based library which could integrate with my existing NextJs application and runs entirely in browser. Actually, that's what brought me here but from reading the docs, I think it also does not support creating such workflows.

Let me share an example of such workflow from LlamaIndex which allows us to create such workflows to show what sort of features am I looking for.

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

# `pip install llama-index-llms-openai` if you don't already have it
from llama_index.llms.openai import OpenAI


class JokeEvent(Event):
    joke: str


class JokeFlow(Workflow):
    llm = OpenAI()

    @step
    async def generate_joke(self, ev: StartEvent) -> JokeEvent:
        topic = ev.topic

        prompt = f"Write your best joke about {topic}."
        response = await self.llm.acomplete(prompt)
        return JokeEvent(joke=str(response))

    @step
    async def critique_joke(self, ev: JokeEvent) -> StopEvent:
        joke = ev.joke

        prompt = f"Give a thorough analysis and critique of the following joke: {joke}"
        response = await self.llm.acomplete(prompt)
        return StopEvent(result=str(response))


w = JokeFlow(timeout=60, verbose=False)
result = await w.run(topic="pirates")
print(str(result))

We define multiple steps in a workflow by sub classing workflow class and use a decorator to describe each step as a method on this class. Each step is triggered when any step emits an event corresponding to its input event and it returns an event itself. There are other frameworks like Langgraph which do it differently.

So, is something like this currently possible in AI SDK or on the roadmap? I am also open to third party libraries which integrate well with AI SDK providing such functionality.

[lgrammel]

What would be the downside of using regular programming to achieve what you want?

[hashiromer]

Thanks for your response, I think these are some important features it should have:

• Workflow should handle all the plumbing automatically, for example, go to next step based on current state and the result of execution. Branches and looping in the workflow should be seamless. I like how Llama does it by emitting events . There should also be a simple way to pass the results of one step to the next and there should be a notion of workflow scoped state each step can mutate and read.

• Have some way to pause workflow at some predefined step to gather human feedback and manually override data and/or next step for critical operates like mutating important external state.

• Have some way to compose different workflows together based on clear interfaces to create larger workflows, Llama index use hierarchal workflows, but I think it could also be done differently. Let me give one concrete example, lets say one step of a workflow needs to write an SQL query, this step itself could be a RAG based workflow consumers should not need to know or worry about. This will also open the door for a plugin system where users could publish their workflows in a registry others could consume without knowing or caring about the internal details of this step. This is recursive of course just like third party libraries in programming languages.

• Monitoring, tracing and debugging functionality. For example, a workflow visualizer see this (https://studio.langchain.com/) for an example of workflow visualization. Debugging features could include some web app to run the workflow step by step, we could skip over nested workflows or go into them similar to stepping over or step into functionality for common debuggers. This can be part of workflow visualizer.

• There should be a clear interface defined in the workflow each step can use to publish notifications. This will not trigger the next step. This could be used for streaming responses, errors, logging etc. There can be other components in the workflow which could subscribe to the stream of these updates via clear interfaces and do whatever they want with them. There are lots of things we can do here too. Anthropic recently created a whole protocol just to describe how external systems should communicate with LLMs (https://modelcontextprotocol.io/introduction). I think its a good start and they describe a lot of problems production systems need to deal with but it also does not define a workflow abstraction and UI specific concerns AI SDK can solve being tightly integrated with UI.

• Workflow description and running configuration should be separate, a workflow could run on a client or on a server or even a combination where some steps run on client and some steps run on the server. Of course, it will need some clear lower level transport layer interface but it will allow plugin creators to use propriety systems in parts of their published workflow and do things like rate limiting, access control etc.

• All API calls should have proxy support including but not limited to LLM API calls.

Just sharing my wish list from the top of my head :)

If you find this intriguing and if its something AI SDK might be interested in, I could write an initial RFC for this feature in more detail.

[lgrammel]

Thanks for the writeup, this is very helpful. We are looking at approaches in this space.

[pattobrien]

Would also be interested in this, but I also recognize that this may be best achieved by using a dedicated workflow/pipeline framework. I'm personally considering using something like Effect.js.

[OmarKhattab]

Interested in this as well

cc @lgrammel

[Und3rf10w]

• Branches and loops are huge, and can be a bit of an issue
• Ideally input properties and output properties would be exposed for each node, and easily accessible
• Workflow definitions should be serializable (e.g. json, ideal for LLM ingestion)
• I agree on separating the workflow definition and the run itself. In our setup, that posted an issue where the composer and executor are in the same "space", so it was a bit of a pain when we wanted to trigger a workflow.
• A big part of a workflow abstraction should simplify context management. E.g. On this branch I might only pass in certain messages of a conversation, process those, and output that and merge that back in at a later point in the workflow.
• The concept of subgraphs is important (importing a workflow into a workflow?)
• Ability to define your own node in a workflow (e.g. I need to retrieve a prompt from a database, then use that as an input to a "system message" node)

Here's a very basic example of a simple workflow from a prototype I built leveraging the AI-SDK:

This has debugging (forwards and backwards debugging), concurrency, breakpoints, serialization, loops, and some other stuff I'm forgetting.

Happy to discuss more on this topic. I've done a lot of exploration on this and have opinions

Mastra's workflow library may be worth looking at, but I haven't used it personally.

[luebken]

An AI SDK based workflow system:
https://github.com/callstackincubator/flows-ai
Possibly helpful.