Module: Raif::Concerns::Llms::XAi::BatchInference

Extended by:
ActiveSupport::Concern
Includes:
SupportsBatchInference
Included in:
Llms::XAi
Defined in:
app/models/raif/concerns/llms/x_ai/batch_inference.rb

Overview

xAI Batch API support for Raif::Llms::XAi. Implements the Raif::Concerns::Llms::SupportsBatchInference contract on top of /v1/files (multipart JSONL upload) and /v1/batches.

xAI's Batch API supports two submission flows; we use the JSONL file-upload flow because it lets us send each per-request body in /v1/chat/completions shape (with messages), matching the sync path verbatim. The alternate REST append flow (POST /v1/batches/id/requests) only accepts the responses wrapper, which targets xAI's Responses API and would force a body-shape conversion (messages -> input, response_format -> text.format, ...).

Submission flow:

POST /v1/files                        -> multipart upload of the JSONL,
                                       returns { id: "file-..." }
POST /v1/batches { input_file_id }    -> creates the batch, returns
                                       { batch_id, state: {...} }
GET  /v1/batches/{id}                 -> status (counts: pending/success/error/cancelled)
GET  /v1/batches/{id}/results         -> paginated per-entry results
POST /v1/batches/{id}:cancel          -> cancels pending entries

Result envelope per entry (same for both submission flows):

{
"batch_request_id": "...",
"batch_result": {
  "response": {
    "chat_get_completion": { "id": ..., "choices": [...], "usage": {...} }
  }
}
}

The chat_get_completion payload matches the synchronous /v1/chat/completions response shape, so we feed it through the host class's #update_model_completion verbatim (the same parser used by the sync and streaming paths). Failures carry an "error" / "error_message" field instead of "response".

xAI has no batch-level state enum: terminal is derived locally from num_pending hitting zero against a nonzero num_requests (or an explicit cancel acknowledgement, a cancel_by_xai_message rejection, or expire_time elapsing). Note the counts are all zero both before xAI has parsed the input file and after it has rejected it, so num_pending alone cannot be trusted.

Not every model a provider serves synchronously can be batched: xAI rejects grok-4.5 at JSONL validation, which is why Raif::Llm#supports_batch_inference? is overridable per model via model_provider_settings.

Instance Method Summary collapse

Instance Method Details

#apply_batch_result(mc, raw_result) ⇒ Object

Applies one per-entry xAI batch result envelope to a Raif::ModelCompletion. The success path feeds batch_result.response.chat_get_completion through the host class's #update_model_completion (same parser used by the sync path), so token counts, tool calls, and response shape are populated identically. The 50% batch discount is applied automatically by Raif::ModelCompletion#calculate_costs.



185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
 
# File 'app/models/raif/concerns/llms/x_ai/batch_inference.rb', line 185

def apply_batch_result(mc, raw_result)
  batch_result = raw_result["batch_result"] || {}
  response_envelope = batch_result["response"] || raw_result["response"]
  error_envelope = batch_result["error"] || raw_result["error"] || raw_result["error_message"]

  mc.started_at ||= mc.raif_model_completion_batch&.started_at || Time.current

  chat_payload = response_envelope.is_a?(Hash) ? response_envelope["chat_get_completion"] : nil

  if chat_payload.is_a?(Hash)
    update_model_completion(mc, chat_payload)
    mc.completed!
  else
    err_message = if error_envelope.is_a?(Hash)
      error_envelope["message"] || error_envelope["error_message"]
    else
      error_envelope
    end

    mc.failure_error = "xAI batch entry failed"
    mc.failure_reason = (err_message.presence || "unknown xAI batch failure").to_s.truncate(255)
    mc.failed!
  end

  mc
end

#batch_class ⇒ Object 



57
58
59
 
# File 'app/models/raif/concerns/llms/x_ai/batch_inference.rb', line 57

def batch_class
  Raif::ModelCompletionBatches::XAi
end

#cancel_batch!(batch) ⇒ Object

Cancellation is fire-and-forget for pending entries -- xAI continues to serve already-processed results, but no further entries are processed. The provider response after a cancel may still report num_pending > 0 if the in-flight cancel hasn't propagated yet; treat as transitional in_progress until counts settle.

Raises:



122
123
124
125
126
127
128
129
130
131
132
133
 
# File 'app/models/raif/concerns/llms/x_ai/batch_inference.rb', line 122

def cancel_batch!(batch)
  raise Raif::Errors::InvalidBatchError, "Batch ##{batch.id} has no provider_batch_id" if batch.provider_batch_id.blank?
  raise Raif::Errors::InvalidBatchError, "Batch ##{batch.id} is already terminal (status=#{batch.status})" if batch.terminal?

  response = with_batch_transient_retry(:cancel, batch_id: batch.id) do
    batch_connection.post("batches/#{batch.provider_batch_id}:cancel")
  end
  body = response.body
  new_status = derive_batch_status(body, batch, post_cancel: true)

  apply_batch_state!(batch, body, new_status)
end

#fetch_batch_results!(batch) ⇒ Object 



135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
 
# File 'app/models/raif/concerns/llms/x_ai/batch_inference.rb', line 135

def fetch_batch_results!(batch)
  completions_by_id = batch.raif_model_completions.index_by(&:batch_custom_id)

  pagination_token = nil
  loop do
    params = {}
    params[:pagination_token] = pagination_token if pagination_token.present?

    response = with_batch_transient_retry(:fetch_results_page, batch_id: batch.id) do
      batch_connection.get("batches/#{batch.provider_batch_id}/results", params)
    end
    body = response.body || {}

    Array(body["results"]).each do |raw|
      custom_id = raw["batch_request_id"]
      mc = completions_by_id[custom_id]
      if mc.nil?
        Raif.logger.warn(
          "xAI batch results: batch_request_id #{custom_id.inspect} did not match any child completion in batch ##{batch.id}"
        )
        next
      end

      apply_batch_result(mc, raw)
    end

    pagination_token = body["pagination_token"].presence
    break if pagination_token.nil?
  end

  completions_by_id.each_value do |mc|
    mc.reload
    next if mc.completed? || mc.failed?

    mc.started_at ||= batch.started_at
    mc.failure_error = "xAI batch entry missing"
    mc.failure_reason = "Result not present in /results stream (batch ##{batch.id})"
    mc.failed!
  end

  batch.recalculate_costs!
  batch
end

#fetch_batch_status!(batch) ⇒ Object 



107
108
109
110
111
112
113
114
115
 
# File 'app/models/raif/concerns/llms/x_ai/batch_inference.rb', line 107

def fetch_batch_status!(batch)
  response = with_batch_transient_retry(:fetch_status, batch_id: batch.id) do
    batch_connection.get("batches/#{batch.provider_batch_id}")
  end
  body = response.body
  new_status = derive_batch_status(body, batch)

  apply_batch_state!(batch, body, new_status)
end

#submit_batch!(batch) ⇒ Object

Raises:



61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
 
# File 'app/models/raif/concerns/llms/x_ai/batch_inference.rb', line 61

def submit_batch!(batch)
  batch.assert_submittable!

  completions = batch.raif_model_completions.to_a
  raise Raif::Errors::InvalidBatchError, "Batch ##{batch.id} has no child completions" if completions.empty?

  jsonl = build_batch_jsonl(completions)
  input_file_id = with_batch_transient_retry(:submit_upload_input, batch_id: batch.id) do
    upload_batch_input_file!(jsonl)
  end

  response = with_batch_transient_retry(:submit_create, batch_id: batch.id) do
    batch_connection.post("batches") do |req|
      req.body = { name: "raif_batch_#{batch.id}", input_file_id: input_file_id }
    end
  end
  body = response.body.is_a?(Hash) ? response.body : {}
  provider_batch_id = body["batch_id"] || body["id"]

  if provider_batch_id.blank?
    raise Raif::Errors::InvalidBatchError,
      "xAI batch create returned no batch id (body=#{response.body.inspect})"
  end

  submitted_at = Time.current

  Raif::ModelCompletionBatch.transaction do
    batch.update!(
      provider_batch_id: provider_batch_id,
      status: "submitted",
      submitted_at: submitted_at,
      started_at: submitted_at,
      provider_response: (batch.provider_response || {}).merge(
        "input_file_id" => input_file_id,
        "expire_time" => body["expire_time"],
        "cost_breakdown" => body["cost_breakdown"]
      ).compact,
      request_counts: derive_request_counts(body) || batch.request_counts
    )

    batch.raif_model_completions.where(started_at: nil).update_all(started_at: submitted_at)
  end

  batch
end