# Ax a New Way to Build Complex Workflows with LLMs

Vikram Rangnekar, Stealth | MLOps Podcast | Episode 259 | 50:51
Hosted by Demetrios Brinkmann

Source: https://www.youtube.com/watch?v=9AI_aeof8-4
Channel: MLOps Community, now AAIF Live (https://www.youtube.com/@AAIFLive-x1r). Summarised by MLOps Talks.
Page: https://mlopstalks.com/talks/ax-a-new-way-to-build-complex-workflows-with-llms
Published: 2024-09-11
Tags: observability, prompt-engineering, structured-outputs, workflows

## TL;DR
- Ax turns LLM prompts into typed functions with defined inputs and outputs, which makes them easier to compose in code.
- Vikram Rangnekar argues that complex workflows should split work into small tasks, since current models struggle with long, multi-step tasks even when they support long context windows.
- DSP-style examples capture useful intermediate traces across a workflow, helping prompts work more consistently and allowing the same approach to support smaller models and model tuning.

## Summary
Vikram Rangnekar explains why he built Ax, a TypeScript framework for composing LLM workflows. He began with a project called 42 Papers, where small models showed him that language models could summarize research papers. The need to manage model APIs, prompts, loops, and output formats led him to the DSP research from Stanford. Ax applies prompt signatures as typed functions, so a workflow can pass structured inputs and outputs between small LLM tasks. Vikram also discusses examples, bootstrapping, and compiler-style optimization. A good top-level example can produce useful traces for every intermediate step, giving each subtask its own context and demonstrations. He recommends avoiding prompts that combine too many jobs, including unnecessary JSON generation. Ax adds streaming assertions, multimodal inputs, agent composition, semantic routing, and built-in OpenTelemetry. The conversation is candid about the immaturity of current LLM frameworks and the work still needed before teams can build useful production systems easily.

## Key ideas
### Small language models first showed Vikram that paper summarization could be automated
[01:25](https://www.youtube.com/watch?v=9AI_aeof8-4&t=85s)
Vikram Rangnekar describes 42 Papers, an early project for discovering papers on arXiv and extracting their main points. He initially considered using people with PhDs to read papers and summarize them, but worried that the process would be difficult to scale. Early models from AI2, trained on arXiv, could summarize papers despite being much smaller than current models. This changed his view of what language models could do. His previous experience at LinkedIn involved traditional machine learning for areas such as ad relevance and fraud, so the ability to turn a blob of text into useful points felt qualitatively different to him.

### Ax began by hiding model API differences and then grew into a workflow framework
[07:07](https://www.youtube.com/watch?v=9AI_aeof8-4&t=427s)
Vikram says Ax started as a way to abstract over LLM client libraries, which each had their own bugs and API behavior. He then focused on prompting because maintaining large instruction blobs felt unpleasant. Prompts with instructions such as 'do this', 'do not do this', and 'then do that' often needed extra loops and error handling. The framework, called Ax, is written for the JavaScript and TypeScript ecosystem. Vikram likes that environment because it lets him iterate quickly. He presents Ax as a framework for building LLM-powered programs rather than as a collection of manually maintained prompt strings.

### Prompt signatures make an LLM call look like a typed function
[09:05](https://www.youtube.com/watch?v=9AI_aeof8-4&t=545s)
The DSP research gave Vikram a cleaner abstraction through signatures. A signature describes inputs and outputs, much like a function. Ax extends that idea with types, so a developer can declare two string inputs and expect a Boolean and a number as outputs. The framework can enforce those types instead of making developers parse blobs of generated text. Vikram says this structure turns a prompt into something closer to an API. Once each operation has defined inputs and outputs, extraction, classification, augmentation, and function-calling steps can be connected as ordinary workflow components.

### Complex LLM work should be split into focused tasks
[10:50](https://www.youtube.com/watch?v=9AI_aeof8-4&t=650s)
Vikram argues that a single prompt should not summarize text, classify it, call functions, and augment it with more data at the same time. The model may perform worse when too many tasks compete for its context and output. A workflow can instead summarize first, pass that result to a classification step, then call functions or retrieve more data. He connects this to the limits of current models on long-horizon tasks. A larger context window allows more text to be supplied, but it does not mean the model has learned to perform a long sequence of different operations reliably.

### Examples can teach each step of a workflow what good behavior looks like
[14:34](https://www.youtube.com/watch?v=9AI_aeof8-4&t=874s)
Vikram explains that DSP uses examples containing inputs and outputs, rather than relying only on written instructions. Examples can expose patterns that are difficult to describe in English, much as a Shakespeare example communicates more than simply asking for Shakespeare-like writing. In a multi-step workflow, the top-level input and final output can be supplied, while bootstrapping records the intermediate inputs and outputs whenever the result is good. Those traces become demonstrations for the individual steps. A final marketing example may be useful for the final result, while an intermediate classification or extraction operation needs its own examples.

### DSP-style traces can make smaller models more useful
[22:09](https://www.youtube.com/watch?v=9AI_aeof8-4&t=1329s)
Vikram says the example and trace process improved the reliability of his workflows enough that he could move some work to smaller models. He also sees a path from captured information to model tuning APIs. The input tokens used for examples can be cheaper than producing incorrect output and then correcting it, so adding useful examples may reduce wasted work. The main benefit is not only a better final response. Each subtask receives examples that match its own inputs and outputs, which helps the whole chain behave more consistently.

### Structured output can add unnecessary work to an LLM
[23:54](https://www.youtube.com/watch?v=9AI_aeof8-4&t=1434s)
Vikram points out that asking for JSON adds another task to an LLM call, even when developers think of JSON as a simple output format. The model must generate the structure and then produce the requested content. He prefers simpler key-value output when nested JSON is not needed. Key-value output can also be processed while it streams. Ax supports streaming assertions that can inspect incoming output, stop early when the format is wrong, and save latency and tokens. Assertions can check declared types or run custom validation functions, including another LLM call when needed.

### Ax composes agents, routes tasks, and keeps production concerns inside the framework
[13:12](https://www.youtube.com/watch?v=9AI_aeof8-4&t=792s)
Ax extends the input-output abstraction to agents. An agent can call functions, reason across turns, and call other agents, allowing developers to build trees of specialized components. For broader inputs, Ax includes a semantic router that embeds the request, compares it with route embeddings, and sends it to a matching workflow. Vikram says this avoids invoking an LLM until a route has been selected. He also describes built-in multimodal inputs and OpenTelemetry tracing. His criticism of many existing frameworks is that they have too many dependencies, immature production behavior, and insufficient visibility into failures such as output limits.

## Notable quotes
- Vikram Rangnekar: "The tool I built, it's basically a framework, it's called Ax." (07:07)
- Vikram Rangnekar: "You don't want to do too much in a prompt." (10:50)
- Vikram Rangnekar: "The examples that DSP captures are basically putting something in there based on what you want like output." (20:17)
- Vikram Rangnekar: "You don't really need that nested structure out from this thing, like mostly not." (24:44)
- Vikram Rangnekar: "Each subtask is like its own clean context." (49:02)

## Tools & references mentioned
- Ax
- 42 Papers
- arXiv
- AI2
- LinkedIn
- DSP
- Omar
- Stanford
- Google
- Gemini
- OpenAI
- Honeycomb
- PyTorch
- JavaScript
- TypeScript
- OpenTelemetry

## Who should watch
- You are building LLM features in TypeScript and want model calls to look more like typed functions than prompt strings.
- Your workflow has several extraction, classification, retrieval, or function-calling steps and failures are hard to diagnose.
- You want to understand DSP-style signatures, demonstrations, bootstrapping, and agent composition before choosing a framework.

## Related talks

- [Using LLMs to Punch Above Your Weight!](https://mlopstalks.com/talks/using-llms-to-punch-above-your-weight) (Cameron Feenstra, Anzen, 35:49)
- [DSPy: Transforming Language Model Calls into Smart Pipelines](https://mlopstalks.com/talks/dspy-transforming-language-model-calls-into-smart-pipelines) (Omar Khattab, Stanford, 1:05:40)
- [Why You Should Care About Observability in LLM Workflows](https://mlopstalks.com/talks/why-you-should-care-about-observability-in-llm-workflows) (Colin McNamera, AlwaysCool.ai, 15:03)
- [AI Squared: Breaking LLMs out of the Chat Application](https://mlopstalks.com/talks/ai-squared-breaking-llms-out-of-the-chat-application) (Benjamin Harvey, AI Squared, 52:20)
- [Alignment is Real](https://mlopstalks.com/talks/alignment-is-real) (Shiva Bhattacharjee, TrueLaw Inc, 38:36)