Module: RubyLLM::Protocols::Cohere::BatchRequests

Included in:
Batches
Defined in:
lib/ruby_llm/protocols/cohere/batch_requests.rb

Overview

:nodoc: all

Instance Method Summary collapse

Instance Method Details

#batch_dataset_type(requests) ⇒ Object

Raises:



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

def batch_dataset_type(requests)
  types = requests.map do |request|
    body = request.fetch(:payload)
    body.key?(:messages) || body.key?('messages') ? 'batch-chat-v2-input' : 'batch-embed-v2-input'
  end.uniq
  raise ArgumentError, 'Cohere batches cannot mix chat and embeddings' unless types.one?

  types.first
end

#render_batch_chat(body) ⇒ Object 



44
45
46
47
48
49
50
51
52
53
54
55
 
# File 'lib/ruby_llm/protocols/cohere/batch_requests.rb', line 44

def render_batch_chat(body)
  if (thinking = body.delete('thinking'))
    body['reasoning'] = thinking['type'] != 'disabled'
    body['thinking_budget'] = thinking['token_budget'] if thinking['token_budget']
  end
  Array(body['messages']).each { |message| render_batch_message(message) }
  Array(body['tools']).each do |tool|
    function = tool.fetch('function')
    parameters = function['parameters']
    function['parameters'] = JSON.generate(parameters) unless parameters.is_a?(String)
  end
end

#render_batch_content(part) ⇒ Object 



63
64
65
66
67
68
69
70
71
 
# File 'lib/ruby_llm/protocols/cohere/batch_requests.rb', line 63

def render_batch_content(part)
  unless %w[text thinking image_url].include?(part['type'])
    raise ArgumentError, "Cohere batches do not support #{part['type']} content"
  end

  part = part.dup
  part['image_url'] = part['image_url'].fetch('url') if part['image_url'].is_a?(Hash)
  part
end

#render_batch_message(message) ⇒ Object 



57
58
59
60
61
 
# File 'lib/ruby_llm/protocols/cohere/batch_requests.rb', line 57

def render_batch_message(message)
  content = message['content']
  content = [{ 'type' => 'text', 'text' => content }] if content.is_a?(String)
  message['content'] = content&.map { |part| render_batch_content(part) }
end

#render_batch_request(request, type:) ⇒ Object 



25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
 
# File 'lib/ruby_llm/protocols/cohere/batch_requests.rb', line 25

def render_batch_request(request, type:)
  body = JSON.parse(JSON.generate(batch_payload(request, except: :model)))
  if type == 'batch-embed-v2-input' && body['output_dimension']
    raise ArgumentError,
          'Cohere batch datasets currently reject dimensions; omit dimensions to use the model default'
  end

  render_batch_chat(body) if type == 'batch-chat-v2-input'
  allowed = type == 'batch-chat-v2-input' ? CHAT_FIELDS : EMBEDDING_FIELDS
  unsupported = body.keys - allowed
  unless unsupported.empty?
    raise ArgumentError, "Cohere batches do not support these request options: #{unsupported.join(', ')}"
  end

  custom_id = request.fetch(:custom_id)
  custom_id = "#{custom_id}:array" if request[:text].is_a?(Array)
  { custom_id:, body: }
end