Module: RubyLLM::Protocols::Cohere::Batches

Includes:
Batch::Helpers, BatchRequests
Defined in:
lib/ruby_llm/protocols/cohere/batches.rb

Overview

:nodoc: all

Instance Method Summary collapse

Methods included from BatchRequests

#batch_dataset_type, #render_batch_chat, #render_batch_content, #render_batch_message, #render_batch_request

Instance Method Details

#batch_results(id) ⇒ Object 



36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
 
# File 'lib/ruby_llm/protocols/cohere/batches.rb', line 36

def batch_results(id)
  data = batch_data(id)
  return [] if data['output_dataset_id'].to_s.empty?

  file = datasets.wait_for_validation(data.fetch('output_dataset_id'))
  model = RubyLLM.models.find(data.fetch('model'), provider: @provider.slug, config: @config)
  parser = self.class.new(@provider, model)
  results = datasets.records(file).map { |row| parse_batch_result(row, parser:, model: model.id) }
  unless results.map(&:first).uniq.size == results.size
    raise Error,
          'Cohere returned duplicate batch request IDs'
  end

  results
end

#cancel_batch(id) ⇒ Object 



31
32
33
34
 
# File 'lib/ruby_llm/protocols/cohere/batches.rb', line 31

def cancel_batch(id)
  @connection.post("v2/batches/#{id}/cancel", {})
  find_batch(id)
end

#create_batch(requests) ⇒ Object 



14
15
16
17
18
19
20
21
22
23
24
25
 
# File 'lib/ruby_llm/protocols/cohere/batches.rb', line 14

def create_batch(requests)
  model = single_batch_model!(requests, 'Cohere')
  type = batch_dataset_type(requests)
  rows = requests.map { |request| render_batch_request(request, type:) }
  file = datasets.upload(StringIO.new(rows.map { |row| "#{JSON.generate(row)}\n" }.join),
                         filename: 'ruby-llm-batch.jsonl', purpose: type)
  datasets.wait_for_validation(file.id)
  response = @connection.post('v2/batches', {
                                name: 'ruby-llm-batch', input_dataset_id: file.id, model: model
                              }, idempotent: false)
  parse_batch_response(response.body.fetch('batch'))
end

#find_batch(id) ⇒ Object 



27
28
29
 
# File 'lib/ruby_llm/protocols/cohere/batches.rb', line 27

def find_batch(id)
  parse_batch_response(batch_data(id))
end