Module: RubyLLM::Protocols::VertexAI::BatchPrediction

Includes:
Batch::Helpers
Included in:
EmbeddingPrediction, RubyLLM::Providers::VertexAI::Anthropic::Batches, RubyLLM::Providers::VertexAI::ChatCompletions::Batches, RubyLLM::Providers::VertexAI::Gemini::Batches
Defined in:
lib/ruby_llm/protocols/vertexai/batch_prediction.rb

Overview

Shared Vertex AI batchPredictionJobs plumbing. The input row and output result shapes belong to each Vertex protocol.

Instance Method Summary collapse

Instance Method Details

#batch_results(id) ⇒ Object 



48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
 
# File 'lib/ruby_llm/protocols/vertexai/batch_prediction.rb', line 48

def batch_results(id)
  job = @connection.get(vertex_batch_name(id)).body
  output_uri = vertex_output_uri(job)
  unless output_uri
    status = parse_batch_status(job['state'], completed: TERMINAL.include?(job['state']))
    return [] if %i[failed cancelled].include?(status)

    raise Error, 'vertexai batch has no GCS output URI yet'
  end

  rows = @provider.list_file_uris(output_uri).grep(/\.jsonl\z/).flat_map do |uri|
    @provider.download_file(uri).to_s.each_line.filter_map do |line|
      next if line.strip.empty?

      JSON.parse(line)
    end
  end
  parse_vertex_batch_results(rows, job:)
end

#cancel_batch(id) ⇒ Object 



43
44
45
46
 
# File 'lib/ruby_llm/protocols/vertexai/batch_prediction.rb', line 43

def cancel_batch(id)
  @connection.post("#{vertex_batch_name(id)}:cancel", {})
  find_batch(id)
end

#create_batch(requests) ⇒ Object 



22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
 
# File 'lib/ruby_llm/protocols/vertexai/batch_prediction.rb', line 22

def create_batch(requests)
  model = single_batch_model!(requests, 'vertexai')
  validate_batch_requests!(requests)
  input_uri, output_uri = vertex_batch_storage_uris
  @provider.upload_file(
    StringIO.new(vertex_batch_jsonl(requests)),
    filename: 'input.jsonl',
    uri: input_uri,
    content_type: 'application/jsonl'
  )

  response = @connection.post("#{@provider.location_path}/batchPredictionJobs",
                              vertex_batch_job(model, input_uri, output_uri), idempotent: false)

  parse_batch_response(response.body)
end

#find_batch(id) ⇒ Object 



39
40
41
 
# File 'lib/ruby_llm/protocols/vertexai/batch_prediction.rb', line 39

def find_batch(id)
  parse_batch_response @connection.get(vertex_batch_name(id)).body
end