Automated LLM Pipeline Optimization with DSPy for Multi-Stage Agent Development

jetblue · Other · 2024 · source Agentic: yes · promotional bias: med · source-backed · extracted with prompt 0.5.4 (schema v2)

Generated from extraction.json — do not edit by hand. Every claim cites verbatim source evidence; stated means the quote directly asserts the claim, implied means it supports it indirectly.

Problem

Domain: Airline LLM application development

JetBlue needed a more reliable way to build and optimize multi-stage LLM applications such as customer-feedback classification and RAG predictive-maintenance chatbots without hand-tuning fragile prompt chains.

Constraints:

multi-stage LLM pipelines contain several promptable components
single prompt-word changes can affect deployment quality
Databricks Model Serving expects DataFrame-oriented interfaces while DSPy works with strings
the airline needed production deployment through internal serving endpoints and application APIs
retrieval quality and answer quality needed objective optimization metrics

Architecture

JetBlue adopted DSPy on Databricks to express LLM applications as modular signatures and PyTorch-like modules rather than opaque prompt chains. A DSPy program can reformulate a question, choose whether to use a vector store, retrieve context, and generate an answer. The module is wrapped as an MLflow PyFunc model, translated between DSPy's string interface and Databricks Model Serving's DataFrame expectations, and deployed to internal serving endpoints backed by Databricks Marketplace models such as Llama 2 70B. DSPy optimizers then treat natural-language pipeline components as tunable parameters and optimize them against metrics such as retrieval quality, answer quality, toxicity, or other LLM-as-judge metrics.

Diagram confidence: explicit (topology described by the source).

flowchart LR
  user-application-layer[User application layer] --> databricks-serving-endpoint[Databricks serving endpoint]
  databricks-serving-endpoint --> mlflow-pyfunc-wrapper[MLflow PyFunc wrapper]
  mlflow-pyfunc-wrapper --> dspy-module[DSPy module]
  dspy-module --> question-reformulation-signature[Question reformulation signature]
  dspy-module --> vector-store[Vector store]
  dspy-module --> marketplace-foundation-model[Marketplace foundation model]
  dspy-optimizer[DSPy optimizer] --> dspy-module
  quality-metrics[Quality metrics] --> dspy-optimizer

Drag to pan · scroll or use +/− to zoom · reconstructed from the source's described component topology

Components

Component	Role	Interacts with
`User application layer`	Sends chatbot or classification requests to the deployed Databricks endpoint.	`Databricks serving endpoint`
`Databricks serving endpoint`	Hosts the logged MLflow PyFunc wrapper for the DSPy pipeline.	`User application layer`, `MLflow PyFunc wrapper`
`MLflow PyFunc wrapper`	Adapts DSPy's string-oriented pipeline to Databricks Model Serving's DataFrame interface.	`Databricks serving endpoint`, `DSPy module`
`DSPy module`	Composes signatures, non-LLM logic, retrieval choices, and answer generation into a transparent pipeline.	`MLflow PyFunc wrapper`, `Question reformulation signature`, `Vector store`, `Marketplace foundation model`
`Question reformulation signature`	Turns an initial question and context into a query for downstream retrieval or answering.	`DSPy module`
`Vector store`	Provides retrieved context for the RAG chatbot when the DSPy module chooses retrieval.	`DSPy module`
`Marketplace foundation model`	Executes the LLM calls inside the DSPy signatures, including Llama 2 70B in the described deployment.	`DSPy module`
`DSPy optimizer`	Optimizes natural-language pipeline parameters against retrieval, answer, or custom LLM-judge metrics.	`DSPy module`, `Quality metrics`
`Quality metrics`	Define task objectives such as retrieval quality, answer quality, toxicity, or LLM-as-judge assessments.	`DSPy optimizer`

Data flow

A user question enters a DSPy module where signatures transform the question into an intermediate query, control logic decides whether to use a vector store, retrieved context is passed to a final answer generator, and the whole module is deployed behind a Databricks endpoint through an MLflow PyFunc wrapper. Evaluation metrics and labeled or unlabeled data feed DSPy optimizers, which choose in-context examples and tune prompts to improve downstream quality.

4 supporting quotes

“For context: think of a signature as a customized, single LLM call in a pipeline. A common, first signature would be to reformat an initial user question into a query using some pre-defined context.”article, stated

“Each signature is accessed within the model’s forward method, sequentially passing an input from one step to the next.”article, stated

“we take a generated query from user input, choose to use a vector store if appropriate, and then generate an answer from our retrieved context.”article, stated

“The model is created using the mlflow.pyfunc.log_model function”article, stated

Why this architecture

JetBlue chose DSPy because hand-maintaining multi-stage prompt chains was fragile, while DSPy exposes signatures and modules as tunable, inspectable pipeline components. Databricks and MLflow provide the production serving path JetBlue already uses, while DSPy optimizers turn quality metrics into an automated prompt/program improvement loop. The PyFunc wrapper is a pragmatic adaptation that lets a string-based DSPy program run on Databricks Model Serving.

4 supporting quotes

“In a complex, multi-stage LLM pipeline, there are often multiple prompts along the way that require tuning. Most industry LLM developers are all too familiar with single words within their prompts that can make or break a deployment (Figure 1).”article, stated

“Most industry LLM developers are all too familiar with single words within their prompts that can make or break a deployment (Figure 1).”article, stated

“Instead of relying on an opaque agent, the internal components are modularized so that each step is clear and able to be assessed and modified.”article, stated

“It should be noted that Databricks Model Serving expects DataFrame formatting, whereas DSPy is working with strings.”article, stated

Models used: Llama 2 70B · DBRX · Databricks Marketplace models · foundation models · LLM-as-a-judge

Scale

DSPy was released by Stanford researchers in October 2023.
“In October 2023, researchers working in Databricks co-founder Matei Zaharia’s Stanford research lab released [DSPy]”article, stated
JetBlue reported a RAG chatbot deployment twice as fast as the earlier Langchain deployment.
“Our RAG chatbot deployment was 2x faster than our Langchain deployment! ![Figure 5. Running a query through our deployed endpoint, returning results from a vector database](https://www.”article, stated
The deployment used Llama 2 70B through Databricks Marketplace.
“we can easily set DSPy to use a Databricks Marketplace model like Llama 2 70B.”article, stated

Alternatives considered (roads not taken)

Manual prompt optimization

Rejected because: JetBlue wanted to optimize retrieval and answer metrics directly instead of repeatedly hand-editing prompt wording.

“Before DSPy we manually optimized our prompts to improve these metrics; now we can use DSPy to directly optimize these metrics”article, stated

Opaque off-the-shelf chatbot interfaces

Rejected because: JetBlue needed components administrators could control, inspect, assess, and modify.

“As more and more companies leverage LLMs, the limitations of a generic chatbot interface are increasingly clear.”article, stated

“These off-the-shelf platforms are highly dependent on parameters that are outside the control of both end-users and administrators.”article, stated

Raw DSPy serving without a Databricks adapter

Rejected because: Model Serving expects DataFrames, so the team needed a wrapper that translated between serving contracts and DSPy's string interface.

“It should be noted that Databricks Model Serving expects DataFrame formatting, whereas DSPy is working with strings.”article, stated

“For that reason, we’ll modify the standard predict and run functions as follows: Figure 5: Modifications to PyFunc model serving definition needed for translating between DSPy and MLflow”article, stated

Patterns

`prompt-versioning` (prompt-operations)

DSPy treats prompt components as tunable parameters that can be optimized against task metrics instead of manually edited by hand.

“thinking of the natural language components of the pipeline as tunable parameters.”article, stated

`eval-driven-development` (evaluation)

JetBlue optimizes RAG retrieval quality, answer quality, and custom LLM-as-judge metrics through DSPy optimizers.

“DSPy optimizers tune these weights by maximizing toward a task objective, requiring just a defined metric”article, stated

“Custom metrics can be designed using LLM-as-a-judge and directly improved upon using DSPy’s optimizers.”article, stated

`rag-retrieval` (retrieval)

The example DSPy pipeline can generate a query, decide to use a vector store, retrieve context, and generate an answer.

“we take a generated query from user input, choose to use a vector store if appropriate, and then generate an answer from our retrieved context.”article, stated

`llm-as-judge` (evaluation)

JetBlue can define custom LLM-as-judge metrics and use DSPy to optimize directly against them.

“Custom metrics can be designed using LLM-as-a-judge and directly improved upon using DSPy’s optimizers.”article, stated

Tradeoffs

Represent LLM applications as compound systems

Chose modular DSPy signatures and traditional software logic over generic chatbot interfaces.

Rationale: JetBlue needed control, dynamic updates, cost optimization, and inspectable internal steps.

“By constructing compound systems that leverage a combination of LLM calls and traditional software development”article, stated

“The DSPy module enables us to optimize LLMOps for better control, dynamic updates, and cost.”article, stated

Optimize with metrics instead of manual prompt iteration

Chose DSPy optimizers with defined task objectives over manual prompt-tuning loops.

Rationale: The optimizer can select examples and tune pipeline behavior to improve retrieval and answer quality metrics.

“DSPy effectively chooses which examples to use in context to improve the reliability and quality of the LLM’s responses.”article, stated

“making the need to manually iterate on prompts unnecessary.”article, stated

Gotchas (what bit them)

Prompt chains become fragile as pipelines grow.

Root cause: Multi-stage systems contain multiple prompts, and small wording changes can affect production behavior.

Resolution: Model the pipeline as DSPy signatures/modules and optimize them against defined metrics.

“In a complex, multi-stage LLM pipeline, there are often multiple prompts along the way that require tuning.”article, stated

“single words within their prompts that can make or break a deployment”article, stated

Serving contracts can mismatch framework contracts.

Root cause: DSPy works with strings, while Databricks Model Serving expects DataFrame formatting.

Resolution: Wrap the DSPy module as an MLflow PyFunc model and adapt predict/run functions.

“Databricks Model Serving expects DataFrame formatting, whereas DSPy is working with strings.”article, stated

“we’ll modify the standard predict and run functions”article, stated

Evaluation

The DSPy setup evaluates and optimizes against retrieval quality, answer quality, custom metrics, and LLM-as-judge metrics such as toxicity.

Metrics: retrieval quality · answer quality · toxicity or other custom LLM-as-judge metrics · deployment speed compared with Langchain

“we have metrics related to retrieval quality and answer quality.”article, stated

“requiring just a defined metric (ie an LLM-as-a-judge assessing toxicity)”article, stated

Outcomes

JetBlue reported its RAG chatbot deployment was 2x faster than the Langchain deployment.
The approach reduced the need for manual prompt iteration.
JetBlue gained a path to optimize multi-stage customer feedback and RAG applications inside Databricks.

The case covers production deployment to internal serving endpoints and application-layer API calls for JetBlue chatbot use cases.

3 supporting quotes

“Our RAG chatbot deployment was 2x faster than our Langchain deployment!”article, stated

“This drastically simplifies the iterative development process of all our LLM applications”article, stated

“deployed to one of JetBlue’s internal serving endpoints”article, stated