Baton.LLMStep behaviour (Baton v0.27.4)

Copy Markdown View Source

A structured contract for LLM workflow steps: you describe the request and how to handle the decoded response — the engine owns everything in between.

Baton.LLMWorker gives LLM steps the right Oban posture (queue, backoff, idempotency, usage stats). But every step body still hand-writes the same transport loop: time the call, invoke the client, map max_tokens truncation, map 429 to a snooze, decode the (possibly fence-wrapped) JSON, and remember to merge llm_usage into the result. That loop is identical across steps and easy to get subtly wrong — a missed error clause either burns tokens retrying a non-retryable request or drops a retryable one.

Baton.LLMStep moves the loop into the library. A step implements:

  • request/1 (required) — build the messages + client options from the job (read upstream deps, render prompts). Return {:ok, messages, opts} or {:ok, messages, opts, context} to thread a value through to handle_response/3. Return {:done, result} to short-circuit with a stored result and no model call (e.g. a fan-out step whose claim has no limitations to check) — no usage is attached, since nothing was spent.
  • handle_response/3 (usually) — turn the decoded payload into the step's stored result. Defaults to {:ok, data}.
  • decode/1 (optional) — defaults to decode_json/1, a tolerant JSON decoder (raw → fenced block → brace slice). Override with {:ok, text} passthrough for free-text steps.
  • output_schema/0 (optional) — a JSON Schema map. When defined it is injected into the client options as :output_schema automatically, so the schema is declared once and both the API call and any introspection tooling read the same one.
  • classify_error/1 (optional) — override the transport-error taxonomy.

What the engine does per call

  1. request/1 → messages + opts (schema injected, context captured)
  2. Gates on the configured Baton.RateLimiter: estimates the call's input tokens from the prompt and calls acquire/3. {:snooze, n} parks the job attempt-free (nothing is sent), tagged :rate_budget so the "snoozed" broadcast names the wait; {:ok, opts} proceeds with the possibly rewritten opts. After a successful response it calls reconcile/3 with the actual input and output tokens (paired with the input estimate and the call's requested :max_tokens as the output estimate). The default limiter is a no-op, so this is invisible until a host configures one.
  3. Times the call and invokes Baton.Debug.call_llm/3 (so context-window capture keeps working unchanged)
  4. stop_reason: "max_tokens"{:error, :max_tokens_truncated} (retryable — an adaptive-thinking retry may fit; override classify_error/1 to discard instead)
  5. Transport errors → classify_error/1, whose snooze verdicts the engine tags :provider_limit for the broadcast. Default taxonomy:
    • HTTP 429 (rate limit) and 529 (overloaded) → {:snooze, n} — no retry attempt consumed
    • other HTTP 4xx (except 408) → {:cancel, reason} — the request itself is invalid; retrying can never succeed
    • anything else → {:error, reason} — retried per max_attempts
  6. decode/1 on response.text; a decode error is a job error (retryable)
  7. handle_response/3 with the decoded data, your context, and the job
  8. If the result is {:ok, map} and carries no llm_usage, the engine attaches one built from the client's normalized response.usage plus measured latency — so stats/cost recording needs nothing from the step. A rejected sample — truncation, a decode failure, a handle_response that returns an error to draw a fresh one — is recorded directly into stats at the point of rejection instead: the call was paid for whether or not the answer was kept, so every attempt's spend is visible in workflow_step_stats.

Example

defmodule MyApp.Steps.AssessQuality do
  use Baton.LLMStep

  @impl true
  def output_schema, do: %{"type" => "object", ...}

  @impl true
  def request(%Oban.Job{} = job) do
    {:ok, %{"parsed" => parsed}} = Baton.Results.get_result(job, :parse_claims)

    messages = [%{role: "user", content: prompt(parsed)}]
    {:ok, messages, [model: "claude-sonnet-4-20250514", system: @system]}
  end

  @impl true
  def handle_response(%{"score" => _} = quality, _ctx, _job) do
    {:ok, %{"quality" => quality}}
  end
end

request/1 may also return {:error, _}, {:snooze, _}, or {:cancel, _} directly (e.g. a dependency read failed) — they pass through untouched.

The client contract

Baton.Debug.call_llm/3 calls the configured :llm_client's complete/2, which must return {:ok, response} with text, model, stop_reason, and a usage map, or {:error, reason} (an HTTP failure as %{status: status} so the taxonomy can read it). Usage keys follow Anthropic naming (cache_read_input_tokens / cache_creation_input_tokens); the engine renames them to the workflow_step_stats columns and passes any extra keys (e.g. web_search_requests) through to the configured pricing module.

Batch mode

A step can trade latency for cost by switching transports:

use Baton.LLMStep, mode: :batch

It then goes through the provider's Message Batches API — around half the token price, hours instead of seconds — while every callback above stays exactly as written. The engine submits a one-request batch, parks the job on Oban snoozes until the batch ends (snoozes cost no retry attempts, so a step can wait a day without touching its budget), then feeds the result through the same decode → handle_response → attach-usage pipeline. Downstream steps, completion, retries, and stats can't tell the difference; the attached usage carries "service_tier" => "batch" so pricing can.

This requires a client implementing Baton.LLMClient's batch callbacks — a step whose client doesn't cancels with {:batch_unsupported, client} on its first attempt. Chaining batch steps stacks their latency, so prefer them for DAG tails or where the downstream work is cheap and local.

Options

All Baton.LLMWorker options are accepted and forwarded, plus:

  • :rate_limit_snooze — seconds to snooze on a rate limit when the provider doesn't advertise its own Retry-After (default: 30; see default_classify/2)
  • :mode:live (default) or :batch
  • :poll_interval — seconds between batch polls (default: 300)
  • :batch_deadline — seconds before a batch is abandoned (default: 90_000, i.e. 25h — a backstop above the provider's own 24h expiry)

Summary

Types

A chat message as accepted by the configured LLM client.

What perform_workflow/1 may return — see Baton.Worker. Engine-originated snoozes carry their reason as a third element (:rate_budget, :provider_limit, :batch_slot, :awaiting); Baton.Worker broadcasts it and strips the tuple back to the 2-tuple Oban accepts.

Callbacks

Map a transport error from the client to a worker outcome. Overrides the default taxonomy (default_classify/2) wholesale — call it yourself for the cases you don't handle.

Decode the model's reply text. Defaults to decode_json/1. Return {:ok, text} verbatim for free-text steps.

Turn the decoded payload into the step's stored result. context is whatever request/1 returned in the 4-tuple form (nil otherwise). Return any Baton.Worker result; llm_usage is attached automatically to {:ok, map} results that don't already carry one.

JSON Schema for the model's output. When defined, injected into the client options as :output_schema (unless the request already set one).

Build the LLM request from the job: read upstream results, render prompts, and return the messages plus client options. The 4-tuple form threads context (any term) through to handle_response/3 — use it for values computed here that the response handler needs (deterministic inputs, the loaded record, …) so nothing is loaded twice.

Functions

Decode a model reply as JSON, tolerating the shapes models actually produce. Tries, in order: the raw text, the content of the first Markdown code fence, and a first-{-to-last-} slice — returning the first that parses, or {:error, {:invalid_json, text}}.

The default transport-error taxonomy

The engine: request → call → classify/decode → handle → attach usage. Invoked by the generated perform_workflow/1; public so a step that overrides perform_workflow/1 (e.g. to branch between two requests) can still delegate to it.

The batch engine: submit → poll → ingest, one Oban attempt per state.

Types

message()

@type message() :: %{role: String.t() | atom(), content: term()}

A chat message as accepted by the configured LLM client.

step_result()

@type step_result() ::
  {:ok, map()}
  | {:error, term()}
  | {:snooze, pos_integer()}
  | {:snooze, pos_integer(), atom()}
  | {:cancel, term()}

What perform_workflow/1 may return — see Baton.Worker. Engine-originated snoozes carry their reason as a third element (:rate_budget, :provider_limit, :batch_slot, :awaiting); Baton.Worker broadcasts it and strips the tuple back to the 2-tuple Oban accepts.

Callbacks

classify_error(reason)

(optional)
@callback classify_error(reason :: term()) ::
  {:error, term()} | {:snooze, pos_integer()} | {:cancel, term()}

Map a transport error from the client to a worker outcome. Overrides the default taxonomy (default_classify/2) wholesale — call it yourself for the cases you don't handle.

decode(t)

(optional)
@callback decode(String.t()) :: {:ok, term()} | {:error, term()}

Decode the model's reply text. Defaults to decode_json/1. Return {:ok, text} verbatim for free-text steps.

decode(t, t)

(optional)
@callback decode(String.t(), Oban.Job.t()) :: {:ok, term()} | {:error, term()}

handle_response(data, context, t)

(optional)
@callback handle_response(data :: term(), context :: term(), Oban.Job.t()) ::
  step_result()

Turn the decoded payload into the step's stored result. context is whatever request/1 returned in the 4-tuple form (nil otherwise). Return any Baton.Worker result; llm_usage is attached automatically to {:ok, map} results that don't already carry one.

output_schema()

(optional)
@callback output_schema() :: map()

JSON Schema for the model's output. When defined, injected into the client options as :output_schema (unless the request already set one).

request(t)

@callback request(Oban.Job.t()) ::
  {:ok, [message()], keyword()}
  | {:ok, [message()], keyword(), context :: term()}
  | {:done, map()}
  | {:error, term()}
  | {:snooze, pos_integer()}
  | {:cancel, term()}

Build the LLM request from the job: read upstream results, render prompts, and return the messages plus client options. The 4-tuple form threads context (any term) through to handle_response/3 — use it for values computed here that the response handler needs (deterministic inputs, the loaded record, …) so nothing is loaded twice.

{:done, result} short-circuits: the step stores result (a {:ok, map} is what perform_workflow/1 returns) and never calls the model — for guards like "this claim has no limitations, so the finding list is empty".

Functions

decode_json(text)

@spec decode_json(String.t()) :: {:ok, term()} | {:error, {:invalid_json, String.t()}}

Decode a model reply as JSON, tolerating the shapes models actually produce. Tries, in order: the raw text, the content of the first Markdown code fence, and a first-{-to-last-} slice — returning the first that parses, or {:error, {:invalid_json, text}}.

With structured outputs the raw text parses directly; the fallbacks cover steps that can't use a schema (e.g. server tools + output_config don't combine), where models tend to wrap JSON in prose and a ```json fence.

default_classify(reason, snooze)

@spec default_classify(term(), pos_integer()) ::
  {:snooze, pos_integer()} | {:error, term()} | {:cancel, term()}

The default transport-error taxonomy:

  • %{status: 429} (rate limited) and %{status: 529} (overloaded) → {:snooze, snooze} — waiting, not failing; no retry attempt consumed. When the error term carries the provider's own advice — a :retry_after key in seconds, or a retry-after header under :headers (a map, a Req-style map of lists, or a list of pairs; name matched case-insensitively) — that advertised delay is used instead of the configured snooze, clamped to 3600s so a malformed value can't park a job for a week. An HTTP-date Retry-After (the other legal form) is not parsed and falls back to the configured snooze.
  • any other %{status: 4xx} except 408 → {:cancel, reason} — the request itself is invalid (bad params, auth, model capability mismatch such as a thinking config the model rejects); retrying resends the same request, so it can never succeed and would only burn attempts and tokens
  • {:cancel, reason}{:cancel, reason} — the client declaring, in its own words, that this can never succeed (see below)
  • everything else (5xx, 408, transport exceptions) → {:error, reason} — retried per max_attempts with the LLMWorker's jittered backoff, discarded when attempts run out

Declaring a permanent failure

The 4xx → cancel rule reads HTTP, which only helps for failures that actually reached the provider. A client can also fail before the request goes out — a model routed to a provider whose batch API it doesn't implement, a missing credential, an endpoint it can't speak — and no retry fixes any of those. Rather than have hosts fabricate a plausible status code to get the cancel they want, a client may return the verdict directly:

{:error, {:cancel, {:batch_unsupported_provider, "openai"}}}

which this taxonomy honours as {:cancel, {:batch_unsupported_provider, "openai"}}. The step is cancelled on its first attempt, its dependents are cancelled with it, and the reason survives intact into oban_jobs.errors.

This matters most on the batch path, where a step that merely errors burns its whole retry budget on a misconfiguration before discarding — and takes the rest of the workflow down slowly rather than immediately.

run(module, job)

@spec run(module(), Oban.Job.t()) :: step_result()

The engine: request → call → classify/decode → handle → attach usage. Invoked by the generated perform_workflow/1; public so a step that overrides perform_workflow/1 (e.g. to branch between two requests) can still delegate to it.

run_batch(module, job, opts \\ [])

@spec run_batch(module(), Oban.Job.t(), keyword()) :: step_result()

The batch engine: submit → poll → ingest, one Oban attempt per state.

Invoked by the generated perform_workflow/1 of a use Baton.LLMStep, mode: :batch step. The step's own callbacks are the same ones a live step implements — only the transport differs, so to the rest of the DAG a batch step is an ordinary step that happens to take hours.

Each attempt looks at the node's checkpoint to decide where it is:

  • no checkpoint — build the request, gate on Baton.RateLimiter.acquire_batch/1, submit a one-request batch, save the batch id, and snooze.
  • checkpoint, batch still running — snooze again. No dependency reads, no prompt rendering; just a poll.
  • checkpoint, batch ended — collect the result and run it through the same decode → handle_response → attach-usage pipeline as a live call.

Snoozing is what makes this affordable: Oban raises max_attempts alongside attempt, so a step can poll for a day without spending its retry budget, and a scheduled job holds the workflow open and its dependents parked.

The checkpoint is kept across snoozes and dropped as soon as the batch's verdict for this request is final — so an Oban retry after a failure submits a fresh batch, exactly as a retried live step makes a fresh call.