Module: RubyLLM::Protocols::Bedrock::Batches

Includes:
Batch::Helpers
Included in:
Converse::Batches, InvokeModel::EmbeddingBatches
Defined in:
lib/ruby_llm/protocols/bedrock/batches.rb

Overview

Shared Bedrock Model Invocation Job and S3 result lifecycle.

Instance Method Summary collapse

Instance Method Details

#batch_results(id) ⇒ Object 



60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
 
# File 'lib/ruby_llm/protocols/bedrock/batches.rb', line 60

def batch_results(id)
  job = @provider.signed_get(@provider.control_api_base, bedrock_job_url(id)).body
  output_uri = job.dig('outputDataConfig', 's3OutputDataConfig', 's3Uri')
  unless output_uri
    status = parse_batch_status(job['status'], completed: TERMINAL.include?(job['status']))
    return [] if %i[failed cancelled].include?(status)

    raise Error, 'bedrock batch has no S3 output URI yet'
  end

  outputs = @provider.list_file_uris(output_uri)
                     .grep(/\.jsonl\.out\z/)
                     .reject { |uri| uri.end_with?('/manifest.json.out') }
                     .map { |uri| @provider.download_file(uri) }
  parse_bedrock_outputs(outputs, model: job['modelId'])
end

#cancel_batch(id) ⇒ Object 



55
56
57
58
 
# File 'lib/ruby_llm/protocols/bedrock/batches.rb', line 55

def cancel_batch(id)
  @provider.signed_post(@provider.control_api_base, "#{bedrock_job_url(id)}/stop", {})
  find_batch(id)
end

#create_batch(requests) ⇒ Object 



17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
 
# File 'lib/ruby_llm/protocols/bedrock/batches.rb', line 17

def create_batch(requests)
  model = single_batch_model!(requests, 'bedrock')
  validate_bedrock_batch_requests!(requests)
  role_arn = bedrock_batch_role_arn
  input_uri, output_uri = bedrock_batch_storage_uris
  @provider.upload_file(
    StringIO.new(bedrock_batch_jsonl(requests)),
    filename: 'input.jsonl',
    uri: input_uri,
    content_type: 'application/jsonl'
  )

  response = @provider.signed_post(@provider.control_api_base, '/model-invocation-job', {
                                     clientRequestToken: Digest::SHA256.hexdigest(input_uri),
                                     jobName: bedrock_job_name(input_uri, requests:),
                                     roleArn: role_arn,
                                     modelId: model,
                                     modelInvocationType: bedrock_invocation_type,
                                     inputDataConfig: {
                                       s3InputDataConfig: { s3Uri: input_uri }
                                     },
                                     outputDataConfig: {
                                       s3OutputDataConfig: { s3Uri: output_uri }
                                     }
                                   })

  find_batch(response.body['jobArn'])
end

#find_batch(id) ⇒ Object 



46
47
48
49
50
51
52
53
 
# File 'lib/ruby_llm/protocols/bedrock/batches.rb', line 46

def find_batch(id)
  data = @provider.signed_get(@provider.control_api_base, bedrock_job_url(id)).body
  protocol = @provider.embedding_batch_protocol(data['modelId']) if data['modelInvocationType'] == 'InvokeModel'
  parser = protocol ? protocol.new(@provider) : self
  parser.send(:parse_batch_response, data).tap do |result|
    result[:batch_protocol] = protocol if protocol
  end
end