Class: RubyLLM::Protocols::Cohere::Datasets

Inherits:
Files show all
Defined in:
lib/ruby_llm/protocols/cohere/datasets.rb

Overview

Cohere datasets for batch input and generated results.

Instance Attribute Summary

Attributes inherited from RubyLLM::Protocol

#config, #connection, #model, #provider

Instance Method Summary collapse

Methods inherited from Files

#find, #list_uris

Methods inherited from RubyLLM::Protocol

abstract, #animate_later, #apply_compaction, #apply_compaction_headers, #apply_end_user, #cache_content, #compact, #complete, #count_tokens, #delete_cache, #embed, #extend_cache, #find_cache, #initialize, #list_models, #moderate, #ocr, #paint, #parse_error, #parse_image_responses, #post_image, #post_video, #preprocess_message, #raise_transcription_streaming_unsupported, #refresh_video_job, #render, #render_embedding, #render_transcription_options, #render_video_extension_payload, #rerank, #server_tool_aliases, #speak, #stream_speech, #stream_speech_response, #stream_transcription, #supports_embedding_media?, #tokenize, #tool_approval_response, #transcribe, #video_extension_attachment, #video_extension_url, #video_request_url

Methods included from RubyLLM::Protocol::BinaryStreaming

#stream_binary

Methods included from RubyLLM::Protocol::Streaming

build_on_data_handler, build_stream_error_response, error_chunk?, failed_http_status, faraday_1?, handle_data, handle_error_chunk, handle_error_event, handle_failed_response, handle_json_error_chunk, handle_sse, handle_stream, json_error_payload?, parse_error_from_json, parse_streaming_error, process_stream_chunk, raise_stream_error, stream_events, stream_response, stream_state

Constructor Details

This class inherits a constructor from RubyLLM::Protocol

Instance Method Details

#download(file_id) ⇒ Object 



24
25
26
27
28
29
30
31
32
33
34
35
 
# File 'lib/ruby_llm/protocols/cohere/datasets.rb', line 24

def download(file_id)
  file = wait_for_validation(file_id)
  parts = dataset_parts(file.metadata)
  originals = parts.filter_map { |part| part['original_url'] }.uniq
  if parts.any? && parts.all? { |part| part['original_url'] }
    return originals.map do |url|
      download_part(url)
    end.join
  end

  records(file).map { |row| "#{JSON.generate(row)}\n" }.join
end

#records(file) ⇒ Object 



37
38
39
40
41
42
43
44
45
46
47
 
# File 'lib/ruby_llm/protocols/cohere/datasets.rb', line 37

def records(file)
  load_avro
  dataset_parts(file.metadata).flat_map do |part|
    reader = nil
    content = StringIO.new(download_part(part.fetch('url')))
    reader = Avro::DataFile::Reader.new(content, Avro::IO::DatumReader.new)
    reader.to_a
  ensure
    reader&.close
  end
end

#upload(file, filename: nil, purpose: nil, expires_in: nil, uri: nil, content_type: nil, provider_options: {}) ⇒ Object

Raises:

  • (ArgumentError) 


8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
 
# File 'lib/ruby_llm/protocols/cohere/datasets.rb', line 8

def upload(file, filename: nil, purpose: nil, expires_in: nil, uri: nil, content_type: nil,
           provider_options: {})
  raise ArgumentError, 'Cohere datasets require purpose: with a dataset type' unless purpose
  raise ArgumentError, 'Cohere datasets do not accept expires_in or uri' if expires_in || uri

  attachment = file_attachment(file, filename:)
  options = { name: attachment.filename, type: purpose, keep_original_file: true }
            .merge(provider_options.transform_keys(&:to_sym))
  response = @connection.post('v1/datasets', { data: file_part(attachment, content_type:) },
                              idempotent: false) do |request|
    request.headers.delete('Content-Type')
    request.params.update(options)
  end
  find(response.body.fetch('id'))
end

#wait_for_validation(id) ⇒ Object 



49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
 
# File 'lib/ruby_llm/protocols/cohere/datasets.rb', line 49

def wait_for_validation(id)
  deadline = Process.clock_gettime(Process::CLOCK_MONOTONIC) + @config.request_timeout
  loop do
    file = find(id)
    return file if file.status == 'validated'
    if file.status == 'failed'
      raise Error, "Cohere dataset #{id} failed validation: #{file.metadata['validation_error']}"
    end
    if Process.clock_gettime(Process::CLOCK_MONOTONIC) >= deadline
      raise Error, "Cohere dataset validation timed out: #{id}"
    end

    sleep 1
  end
end