batch

package
v0.11.44 Latest Latest
Warning

This package is not in the latest version of its module.

 Go to latest Published: Oct 19, 2026 License: Apache-2.0 Imports: 30 Imported by: 0

Documentation

Overview

Package batch runs Claude and Gemini requests through Vertex AI batch prediction, which bills at about half the online rate in exchange for results that arrive within hours rather than seconds. It is a single-shot client, not a conversation executor: a job is one provider, one model, and a set of independent requests, each answered once.

Nothing in the package waits. Client.SubmitClaude and Client.SubmitGemini stage the input in Cloud Storage and create the job; Client.Status reports progress; Client.Results streams the outcomes written so far, and once the job is terminal also reports the keys that have none; the caller decides when to check again.

Submission identity

Vertex job creation carries no request id, so a lost response would otherwise submit the rows twice. Every submission names a caller-chosen id, used as the job's display name and its storage directory. Before each create attempt the client lists the id's jobs, so resubmitting an id, or retrying after a lost response, returns the existing job. A job is adopted only when its display name, input URI, and output destination all match; one that reads the id's input but writes elsewhere fails the submission with ErrInvalidRequest. A caller that means to run the same rows again uses a new id.

The manifest is written create-only, so it claims the id for one set of rows: staging different rows under an id that has no job yet fails with ErrInvalidRequest and leaves the staged objects alone, while the same rows (a retry, or recovery after a crash before create) carry on.

Once a job exists, the submitter records its name in job.json, also create-only. Concurrent submissions of one id and the same rows may each create a job, but only one record is written: every caller returns the job it names and cancels the others it lists, whatever their creation times or the order the job list shows them in. The recording caller lists again after its record, so a twin that created and then stopped before its own record is cancelled too. Later submissions return the recorded job and cancel any other listed duplicates; Client.Find returns the recorded job and never cancels. A record naming a job that does not exist or was not created for the id fails with ErrInvalidRequest.

A submitter that stops between create and its record leaves jobs with no record. The next submission records the earliest listed job and cancels the rest; until then, Find returns that same job. This recovery, and the lookup before each create, rely on the job list showing a job soon after it is created, which is not yet verified against Vertex: a job the list misses can be created twice, and a create that lands after the recording caller's second listing is cancelled only when a later submission lists it. Callers should still serialize submissions per id.

Models

Callers pass bare Vertex model IDs ("claude-opus-5-5@default", "gemini-3.8-flash"). The publisher path is added only at the Vertex boundary, so the prefixed form never reaches model routing.

Correlation

Each request carries a key that comes back on its Result. Claude lines carry a digest of the key as custom_id. Gemini lines have no such field, so the digest goes in a request label, and a manifest of canonical request digests matches the echoed request when Vertex does not echo the label. Keys with no output once the job is terminal yield ErrMissing.

Outcomes

Classification matches the online clients and fails closed. A Claude stop_reason of "refusal", a Gemini prompt blockReason, or any Gemini finishReason other than STOP or MAX_TOKENS (including reasons this package does not know) yields ErrRefused with the provider's category and detail in Result.Refusal; the package does not interpret the category. A refused result carries no text, matching the online clients, which never accept partial content. A max-tokens stop yields ErrTruncated and keeps the partial text. A finished response with no text, and a Gemini response with no candidates, are transient per-line errors. Per-line provider failures are returned as Result.Err, and IsRetryable tells transient ones apart. The package never resubmits on its own.

Telemetry

With WithMetrics, each yielded result records its tokens and outcome in the same genai.* series as the online executors, under gen_ai.operation.name "batch_predict" and gen_ai.provider.name "gcp.vertex_ai".

Request content is staged in the bucket in plain JSONL; callers own its retention and access.

Example

Example submits two Claude requests under a stable id, then, on a later pass, checks the job and reads its results. It needs Vertex AI and a bucket the Vertex AI service agent can read and write, so it is not run. 

package main

import (
	"context"
	"errors"

	"cloud.google.com/go/storage"
	"github.com/anthropics/anthropic-sdk-go"
	"github.com/chainguard-dev/clog"
	"google.golang.org/genai"

	"chainguard.dev/driftlessaf/agents/executor/batch"
	"chainguard.dev/driftlessaf/agents/metrics"
)

func main() {
	ctx := context.Background()
	gc, err := genai.NewClient(ctx, &genai.ClientConfig{
		Backend:  genai.BackendVertexAI,
		Project:  "my-project",
		Location: "us",
	})
	if err != nil {
		panic(err)
	}
	sc, err := storage.NewClient(ctx)
	if err != nil {
		panic(err)
	}
	defer sc.Close()

	client, err := batch.NewClient(gc, sc.Bucket("my-project-batch"), "jobs",
		batch.WithMetrics(metrics.NewGenAI("chainguard.ai.agents")))
	if err != nil {
		panic(err)
	}

	params := func(prompt string) anthropic.MessageNewParams {
		return anthropic.MessageNewParams{
			MaxTokens: 4096,
			Messages:  []anthropic.MessageParam{anthropic.NewUserMessage(anthropic.NewTextBlock(prompt))},
		}
	}
	// The id is the caller's stable identity for this set of rows: running
	// this again after a crash returns the same job.
	job, err := client.SubmitClaude(ctx, "review-0a1b2c", "claude-opus-5-5@default", []batch.ClaudeRequest{
		{Key: "skill-a", Params: params("Review skill A.")},
		{Key: "skill-b", Params: params("Review skill B.")},
	})
	if err != nil {
		panic(err)
	}

	// Later: check the job, and collect once it is terminal.
	st, err := client.Status(ctx, job)
	if err != nil {
		panic(err)
	}
	if !st.Terminal {
		return
	}
	for res, err := range client.Results(ctx, job) {
		if err != nil {
			panic(err)
		}
		switch {
		case res.Err == nil:
			clog.InfoContextf(ctx, "%s: %d output tokens", res.Key, res.Usage.Out)
		case errors.Is(res.Err, batch.ErrRefused):
			clog.InfoContextf(ctx, "%s: refused (%s)", res.Key, res.Refusal.Category)
		case errors.Is(res.Err, batch.ErrTruncated), errors.Is(res.Err, batch.ErrMissing), batch.IsRetryable(res.Err):
			clog.InfoContextf(ctx, "%s: resubmit under a new id", res.Key)
		default:
			clog.InfoContextf(ctx, "%s: %v", res.Key, res.Err)
		}
	}
}

Index

Examples

Constants

This section is empty.

Variables

View Source 
var (
	// ErrInvalidRequest identifies a request the client rejects: invalid
	// input, or an id whose existing job does not match this client.
	ErrInvalidRequest = errors.New("invalid batch request")
	// ErrRefused marks a result the provider's safety system declined to
	// answer. [Result.Refusal] carries the provider's category and detail.
	ErrRefused = errors.New("batch request refused")
	// ErrTruncated marks a result that stopped at its output token limit.
	ErrTruncated = errors.New("batch response truncated at max tokens")
	// ErrMissing marks a submitted key with no output line once the job
	// reached a terminal state.
	ErrMissing = errors.New("batch result missing")
)

Functions

func IsRetryable 

func IsRetryable(err error) bool

IsRetryable reports whether err is a transient failure worth another attempt: a per-line provider error or an API error with a rate-limit, overload, timeout, or server status, or a network timeout. ErrMissing is not covered; a caller decides separately whether to resubmit a key with no output. ErrInvalidRequest never is.

Types

type ClaudeRequest 

type ClaudeRequest struct {
	// Key identifies the request in [Result.Key]; it must be unique in the
	// job.
	Key string
	// Params is the request. Its Model is ignored, since the job carries the
	// model, and streaming does not apply.
	Params anthropic.MessageNewParams
}

ClaudeRequest is one Claude Messages request in a batch.

type Client 

type Client struct {
	// contains filtered or unexported fields
}

Client submits and collects Vertex AI batch prediction jobs for Claude and Gemini. Construct one with NewClient; a Client is safe for concurrent use.

func NewClient 

func NewClient(gc *genai.Client, bucket *storage.BucketHandle, prefix string, opts ...Option) (*Client, error)

NewClient returns a Client that runs jobs through gc, which must use the Vertex AI backend, and stages input and output under prefix in bucket. The Vertex AI service agent needs read and create access to the bucket.

func (*Client) Cancel 

func (c *Client) Cancel(ctx context.Context, job Job) error

Cancel asks Vertex to cancel the job. Cancellation is asynchronous: the job reaches CANCELLED later, and rows already finished keep their output.

func (*Client) Find 

func (c *Client) Find(ctx context.Context, id string) (Job, bool, error)

Find returns the job submitted under id, if any. It filters the job list on the display name and re-checks the name and input URI on each job, so a filter Vertex ignores, or a job with the same name from another client, cannot produce a false match. A job with this id's name and input but another output destination, or none, was not created by this client: Find returns an error wrapping ErrInvalidRequest rather than adopt it.

Find returns the job Submit recorded for the id, whatever its state and whatever duplicates are listed. A recorded job that does not exist or fails the checks above is also ErrInvalidRequest. With no record yet (a submitter stopped between creating the job and recording it), Find returns the earliest listed job, the one a recovering Submit records. Find is read-only: duplicates are cancelled only by Submit.

Example

ExampleClient_Find looks up a job by its submission id without submitting. 

package main

import (
	"context"

	"github.com/chainguard-dev/clog"

	"chainguard.dev/driftlessaf/agents/executor/batch"
)

func main() {
	ctx := context.Background()
	var client *batch.Client // from batch.NewClient
	job, ok, err := client.Find(ctx, "review-0a1b2c")
	if err != nil {
		panic(err)
	}
	if ok {
		clog.InfoContextf(ctx, "job %s has %d requests", job.Name, job.Keys)
	}
}

func (*Client) Results 

func (c *Client) Results(ctx context.Context, job Job) iter.Seq2[Result, error]

Results streams the job's results. Each output line yields one Result keyed by the caller's key; once the job is terminal, every key with no line then yields ErrMissing, in submission order. Before the job is terminal it yields only the lines already written. A non-nil error ends the sequence: the job id is invalid, or the job or its output could not be read.

Lines that match no submitted key are skipped and logged, so they surface as ErrMissing. With WithMetrics, usage is recorded as each result is yielded, so iterating twice records twice.

func (*Client) Status 

func (c *Client) Status(ctx context.Context, job Job) (Status, error)

Status returns the job's current state and completion counts.

func (*Client) SubmitClaude 

func (c *Client) SubmitClaude(ctx context.Context, id, model string, reqs []ClaudeRequest, opts ...SubmitOption) (Job, error)

SubmitClaude submits reqs to model (a bare Vertex model ID such as "claude-opus-5-5@default") as one batch job under the submission id. Submitting an id that already has a job returns that job without writing input or creating another; a new id is needed to resubmit the same rows. Different rows under an id with no job yet are rejected with ErrInvalidRequest. Concurrent calls with the same id and rows return the job recorded first, though a duplicate may be created and then cancelled; see the package documentation.

func (*Client) SubmitGemini 

func (c *Client) SubmitGemini(ctx context.Context, id, model string, reqs []GeminiRequest, opts ...SubmitOption) (Job, error)

SubmitGemini submits reqs to model (a bare Vertex model ID such as "gemini-3.8-flash") as one batch job under the submission id. It shares Client.SubmitClaude's submission identity contract.

type GeminiRequest 

type GeminiRequest struct {
	// Key identifies the request in [Result.Key]; it must be unique in the
	// job.
	Key string
	// Contents is the conversation sent to the model.
	Contents []*genai.Content
	// Config holds the generation settings and request-level fields such as
	// the system instruction, safety settings, and labels. Client-side
	// fields (HTTPOptions, automatic continuation) are not sent.
	Config *genai.GenerateContentConfig
}

GeminiRequest is one Gemini generateContent request in a batch.

type Job 

type Job struct {
	// ID is the caller's submission id, also the job's display name.
	ID string
	// Name is the Vertex resource name of the BatchPredictionJob.
	Name string
	// Provider is "anthropic" for Claude jobs and "google" for Gemini jobs.
	Provider string
	// Model is the bare model ID as submitted, without a publisher prefix.
	Model string
	// InputURI is the gs:// URI of the job's JSONL input.
	InputURI string
	// OutputPrefix is the gs:// prefix Vertex writes predictions under.
	OutputPrefix string
	// Keys is the number of requests in the job.
	Keys int
}

Job identifies one submitted batch job. It holds everything Client.Status, Client.Results, and Client.Cancel need, so a caller may persist it.

type Option 

type Option func(*Client) error

Option configures a Client.

func WithMetrics 

func WithMetrics(genai *metrics.GenAI) Option

WithMetrics records the token usage and outcome of every result Client.Results yields. Without it the client records nothing.

func WithResourceLabels 

func WithResourceLabels(labels map[string]string) Option

WithResourceLabels adds attributes to every metric the client records, for example service_name.

func WithRetryConfig 

func WithRetryConfig(cfg retry.RetryConfig) Option

WithRetryConfig overrides the retry policy for Vertex and Cloud Storage calls.

type Refusal 

type Refusal struct {
	// Category is Claude's stop_details category, Gemini's prompt
	// blockReason, or any Gemini finishReason other than STOP or MAX_TOKENS.
	Category string
	// Detail is Claude's explanation, Gemini's blockReasonMessage, or
	// Gemini's finishMessage.
	Detail string
}

Refusal is the provider's reason for refusing a request.

type Result 

type Result struct {
	// Key is the request's key as submitted.
	Key string
	// Text is the response text: Claude's text blocks, or the non-thought
	// text parts of Gemini's first candidate, concatenated in order. It is
	// empty when Err is [ErrRefused]; a truncated response keeps its partial text.
	Text string
	// Usage is the token accounting for the request.
	Usage Usage
	// Stop is the provider's stop or finish reason, verbatim.
	Stop string
	// Refusal is set with [ErrRefused].
	Refusal *Refusal
	// Err is nil for a usable response, [ErrRefused], [ErrTruncated],
	// [ErrMissing], or a per-line provider error; [IsRetryable] reports
	// whether that last kind is transient.
	Err error
}

Result is the outcome of one request in a job.

type Status 

type Status struct {
	// State is the Vertex job state.
	State genai.JobState
	// Terminal reports that the job will make no further progress: it
	// succeeded, partially succeeded, failed, was cancelled, expired, or was
	// paused.
	Terminal bool
	// Succeeded, Failed, and Incomplete are the Vertex completion counts.
	// Incomplete is -1 when Vertex does not know it.
	Succeeded, Failed, Incomplete int
}

Status is a job's progress as Vertex reports it.

type SubmitOption 

type SubmitOption func(*submitConfig) error

SubmitOption configures one submission. None are defined yet; the parameter keeps later options from changing the Submit signatures.

type Usage 

type Usage struct {
	In, Out, CacheRead, CacheWrite int64
}

Usage is a request's token accounting, with the same meaning as the online executors record: In includes cached input for Gemini, Out includes Gemini's thinking tokens.

Source Files

View all Source files

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL
go.dev uses cookies from Google to deliver and enhance the quality of its services and to analyze traffic. Learn more.