ruby_llm 1.16.0 → 2.0.0.rc2

data/lib/ruby_llm/protocols/chat_completions.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    # The OpenAI Chat Completions API — the lingua franca of LLM APIs.
+    class ChatCompletions < Protocol
+      include ChatCompletions::Chat
+      include ChatCompletions::Embeddings
+      include ChatCompletions::Models
+      include ChatCompletions::Moderation
+      include ChatCompletions::Streaming
+      include ChatCompletions::Tools
+      include ChatCompletions::Images
+      include ChatCompletions::Media
+      include ChatCompletions::Speech
+      include ChatCompletions::Transcription
+
+      public :render_transcription_options
+    end
+  end
+end

data/lib/ruby_llm/protocols/cohere/batch_requests.rb

@@ -0,0 +1,75 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Cohere
+      module BatchRequests # :nodoc: all
+        CHAT_FIELDS = %w[
+          messages tools temperature p frequency_penalty presence_penalty reasoning thinking_budget
+          return_prompt logprobs max_tokens max_input_tokens k seed
+        ].freeze
+        EMBEDDING_FIELDS = %w[texts images input_type inputs max_tokens output_dimension embedding_types
+                              truncate].freeze
+        private_constant :CHAT_FIELDS, :EMBEDDING_FIELDS
+
+        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
+
+        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
+
+        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
+
+        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
+
+        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
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/cohere/batches.rb

@@ -0,0 +1,98 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Cohere
+      module Batches # :nodoc: all
+        include RubyLLM::Batch::Helpers
+        include Cohere::BatchRequests
+
+        TERMINAL_STATUSES = %w[BATCH_STATUS_COMPLETED BATCH_STATUS_FAILED BATCH_STATUS_CANCELED].freeze
+        Response = Struct.new(:body)
+        private_constant :TERMINAL_STATUSES, :Response
+
+        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
+
+        def find_batch(id)
+          parse_batch_response(batch_data(id))
+        end
+
+        def cancel_batch(id)
+          @connection.post("v2/batches/#{id}/cancel", {})
+          find_batch(id)
+        end
+
+        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
+
+        private
+
+        def datasets
+          @datasets ||= Cohere::Datasets.new(@provider)
+        end
+
+        def batch_data(id)
+          @connection.get("v2/batches/#{id}").body.fetch('batch')
+        end
+
+        def parse_batch_response(data)
+          {
+            id: data.fetch('id'), raw_status: data.fetch('status'),
+            completed: TERMINAL_STATUSES.include?(data['status']), request_count: data['num_records'],
+            request_counts: {
+              'total' => data['num_records'], 'succeeded' => data['num_successful_records'],
+              'failed' => data['num_failed_records']
+            }.compact
+          }
+        end
+
+        def parse_batch_status(raw_status, completed:)
+          return :pending unless completed
+          return :succeeded if raw_status == 'BATCH_STATUS_COMPLETED'
+          return :cancelled if raw_status == 'BATCH_STATUS_CANCELED'
+
+          :failed
+        end
+
+        def parse_batch_result(row, parser:, model:)
+          custom_id, shape = row.fetch('custom_id').split(':', 2)
+          index = batch_result_index(custom_id)
+          return [index, nil, batch_failure(custom_id, row['error'])] if !row['error'].to_s.empty? || !row['body']
+
+          body = row.fetch('body')
+          result = if body['embeddings']
+                     parser.send(:parse_embedding_response, Response.new(body), model:,
+                                                                                text: shape == 'array' ? [] : nil)
+                   else
+                     parser.send(:parse_completion_body, body, raw: body)
+                   end
+          [index, result]
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/cohere/chat.rb

@@ -0,0 +1,227 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Cohere
+      # Chat methods for the Cohere v2 API implementation
+      module Chat
+        FINISH_REASONS = {
+          'COMPLETE' => :stop, 'STOP_SEQUENCE' => :stop, 'MAX_TOKENS' => :max_tokens, 'TOOL_CALL' => :tool_calls
+        }.freeze
+
+        module_function
+
+        def finish_reasons = FINISH_REASONS
+
+        def normalize_finish_reason(reason)
+          return nil if reason.nil?
+
+          finish_reasons.fetch(reason.to_s) { reason.to_s.to_sym }
+        end
+
+        def completion_url
+          'v2/chat'
+        end
+
+        # rubocop:disable-next Lint/UnusedMethodArgument
+        def render_payload(messages, tools:, temperature:, model:, stream: false, max_output_tokens: nil,
+                           schema: nil, thinking: nil, citations: false, caching: nil, tool_prefs: nil)
+          warn_unsupported_citations(model) if citations && !model.supports?(:citations)
+
+          payload = {
+            model: model.id,
+            messages: format_messages(messages, citations: citations),
+            stream: stream
+          }
+
+          add_optional_fields(payload, messages, temperature:, max_output_tokens:, citations:, schema:)
+          add_tools(payload, tools, tool_prefs || {})
+          add_thinking(payload, thinking)
+          payload
+        end
+
+        def add_optional_fields(payload, messages, temperature:, max_output_tokens:, citations:, schema:)
+          payload[:temperature] = temperature unless temperature.nil?
+          payload[:max_tokens] = max_output_tokens unless max_output_tokens.nil?
+          payload[:documents] = Media.format_documents(messages) if citations && Media.documents?(messages)
+          payload[:response_format] = build_response_format(schema) if schema
+        end
+
+        def warn_unsupported_citations(model)
+          RubyLLM.logger.warn(
+            "#{model.id} does not support citations according to the model registry. " \
+            'with_citations may have no effect.'
+          )
+        end
+
+        def add_tools(payload, tools, tool_prefs)
+          return if tools.empty?
+
+          payload[:tools] = tools.values.map { |tool| Tools.function_for(tool) }
+          tool_choice = Tools.build_tool_choice(tool_prefs[:choice])
+          payload[:tool_choice] = tool_choice if tool_choice
+        end
+
+        # Cohere takes a bare JSON Schema under json_schema, with no name or
+        # strict wrapper.
+        def build_response_format(schema)
+          normalized = RubyLLM::Support::Utils.deep_dup(schema[:schema])
+          normalized.delete(:strict)
+          normalized.delete('strict')
+
+          { type: 'json_object', json_schema: normalized }
+        end
+
+        # Reasoning is on by default for models that support it, so an
+        # explicit disable is as meaningful as an explicit enable.
+        def add_thinking(payload, thinking)
+          return unless thinking&.enabled?
+          return payload[:thinking] = { type: 'disabled' } if thinking.disabled?
+
+          payload[:thinking] = { type: 'enabled', token_budget: thinking.budget }.compact
+        end
+
+        def format_messages(messages, citations: false)
+          messages.map { |msg| format_message(msg, citations: citations) }
+        end
+
+        def format_message(msg, citations: false)
+          return Tools.format_tool_result(msg) if msg.tool_result?
+          return format_assistant_message(msg) if msg.role == :assistant
+
+          {
+            role: msg.role.to_s,
+            content: Media.format_content(msg.content, msg.attachments, citations: citations)
+          }
+        end
+
+        # Cohere returns tool_plan alongside tool calls, and RubyLLM surfaces
+        # it as thinking, but the newer models reject it on the way back in.
+        # It is optional in a request, so the plan stays out of the history.
+        def format_assistant_message(msg)
+          message = { role: 'assistant' }
+          content = Media.format_content(msg.content, msg.attachments)
+          message[:content] = content unless content.empty?
+          message[:tool_calls] = Tools.format_tool_calls(msg.tool_calls) if msg.tool_call?
+          message
+        end
+
+        def parse_completion_body(data, raw:)
+          message_data = data['message'] || {}
+          blocks = Array(message_data['content'])
+          content, offsets = extract_text(blocks)
+          finish_reason = normalize_finish_reason(data['finish_reason'])
+
+          Message.new(
+            role: :assistant,
+            content: content,
+            citations: parse_citations(message_data['citations'], offsets),
+            thinking: Thinking.build(text: extract_thinking(blocks, message_data)),
+            tool_calls: Tools.parse_tool_calls(message_data['tool_calls'], response: raw, finish_reason:),
+            finish_reason: finish_reason,
+            model: model&.id,
+            raw: raw,
+            **usage_tokens(data['usage'] || {})
+          )
+        end
+
+        # Cohere reports what the model processed under tokens and what it
+        # charges for under billed_units; the two differ because Cohere does
+        # not bill its own preamble.
+        def usage_tokens(usage)
+          tokens = usage['tokens'] || {}
+          billed = usage['billed_units'] || {}
+
+          {
+            input_tokens: tokens['input_tokens'] || billed['input_tokens'],
+            output_tokens: tokens['output_tokens'] || billed['output_tokens'],
+            cache_read_tokens: usage['cached_tokens']
+          }
+        end
+
+        # Returns the joined text of the response along with the offset each
+        # content block starts at, so citation spans resolve against the
+        # content string RubyLLM exposes.
+        def extract_text(blocks)
+          text = +''
+          offsets = {}
+
+          blocks.each_with_index do |block, index|
+            next unless block['type'] == 'text'
+
+            offsets[index] = text.length
+            text << block['text'].to_s
+          end
+
+          [text, offsets]
+        end
+
+        # The tool plan is the model's reasoning for a tool-calling turn, and
+        # is the only reasoning Cohere returns when thinking blocks are absent.
+        def extract_thinking(blocks, message_data)
+          thoughts = blocks.select { |block| block['type'] == 'thinking' }
+                           .map { |block| block['thinking'] }.join
+          thoughts.empty? ? message_data['tool_plan'] : thoughts
+        end
+
+        def parse_citations(citations, offsets)
+          Array(citations).map { |citation| parse_citation(citation, offsets) }
+        end
+
+        # Only citations of the response text carry offsets into #content.
+        # Citations of the thinking blocks or the tool plan point into text
+        # RubyLLM exposes elsewhere, so they keep their snippet without a span.
+        def parse_citation(data, offsets = {})
+          source = Array(data['sources']).first
+          document = source_document(source)
+          start_index, end_index = citation_span(data, offsets)
+
+          Citation.new(
+            url: citation_url(document),
+            title: document['title'] || document['id'],
+            cited_text: document['text'] || document['snippet'],
+            text: data['text'],
+            start_index: start_index,
+            end_index: end_index,
+            source_index: source_index(source)
+          )
+        end
+
+        def citation_span(data, offsets)
+          return [nil, nil] unless text_citation?(data)
+
+          offset = offsets[data['content_index'] || 0]
+          return [nil, nil] unless offset
+
+          [data['start'] && (offset + data['start']), data['end'] && (offset + data['end'])]
+        end
+
+        def text_citation?(data)
+          type = data['type']
+          type.nil? || type == 'TEXT_CONTENT'
+        end
+
+        def source_document(source)
+          return {} unless source
+
+          source['document'] || source['tool_output'] || {}
+        end
+
+        def citation_url(document)
+          url = document['url']
+          url if url.is_a?(String) && url.match?(%r{\Ahttps?://}i)
+        end
+
+        # Documents RubyLLM sends, and those Cohere numbers itself, are
+        # identified as doc:N where N is the document's position.
+        def source_index(source)
+          id = source && source['id']
+          return unless id.is_a?(String)
+
+          match = id.match(/\Adoc:(\d+)\z/)
+          match && match[1].to_i
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/cohere/datasets.rb

@@ -0,0 +1,102 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Cohere
+      # Cohere datasets for batch input and generated results.
+      class Datasets < Protocols::Files
+        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
+
+        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
+
+        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
+
+        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
+
+        private
+
+        def files_url
+          'v1/datasets'
+        end
+
+        def parse_file_response(response)
+          data = response.fetch('dataset')
+          parts = dataset_parts(data)
+          original = parts.first&.fetch('original_url', nil)
+          filename = original ? File.basename(URI.parse(original).path) : "#{data.fetch('name')}.jsonl"
+          mime_type = if File.extname(filename) == '.jsonl'
+                        'application/jsonl'
+                      else
+                        RubyLLM::Files::MimeType.for(name: filename)
+                      end
+          uploaded_file(data, id: data.fetch('id'), filename:, mime_type:,
+                              created_at: timestamp(data['created_at']), status: data['validation_status'],
+                              purpose: data['dataset_type'], downloadable: !parts.empty?)
+        end
+
+        def dataset_parts(data)
+          Array(data['dataset_parts']).sort_by.with_index { |part, index| part.fetch('index', index) }
+        end
+
+        def download_part(url)
+          Transport::Connection.basic(@config).get(url).body
+        end
+
+        def load_avro
+          require 'avro'
+        rescue LoadError
+          raise LoadError, 'Add gem "avro" to your Gemfile to read Cohere batch results and processed datasets'
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/cohere/embeddings.rb

@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Cohere
+      # Embeddings methods for the Cohere v2 API integration
+      module Embeddings
+        DEFAULT_INPUT_TYPE = 'search_document'
+
+        module_function
+
+        def embedding_url(...)
+          'v2/embed'
+        end
+
+        # rubocop:disable-next Lint/UnusedMethodArgument
+        def render_embedding_payload(text, model:, dimensions:, task_type: nil, title: nil, with: [],
+                                     provider_options: {})
+          image_only = with.any? && separate_image_embeddings?(model)
+          payload = {
+            model: model,
+            input_type: task_type || (image_only ? 'image' : DEFAULT_INPUT_TYPE),
+            embedding_types: ['float'],
+            output_dimension: dimensions
+          }.compact
+
+          payload.merge!(image_only ? image_embedding_inputs(text, with) : embedding_inputs(text, with))
+          Support::Utils.deep_merge(payload, provider_options)
+        end
+
+        def supports_embedding_media?
+          true
+        end
+
+        def separate_image_embeddings?(model) # :nodoc:
+          %w[embed-english-v3.0 embed-multilingual-v3.0].include?(model)
+        end
+
+        def image_embedding_inputs(text, attachments) # :nodoc:
+          raise ArgumentError, 'Cohere Embed v3 accepts text or an image, not both' unless text.nil? || text == ''
+          raise ArgumentError, 'Cohere Embed v3 accepts one image per request' unless attachments.one?
+          raise UnsupportedAttachmentError, attachments.first.mime_type unless attachments.first.image?
+
+          { images: ["data:#{attachments.first.mime_type};base64,#{attachments.first.encoded}"] }
+        end
+
+        def parse_embedding_response(response, model:, text:)
+          data = response.body
+          vectors = data.dig('embeddings', 'float')
+          vectors = vectors.first if vectors&.length == 1 && !text.is_a?(Array)
+          billed = data.dig('meta', 'billed_units') || {}
+
+          Embedding.new(vectors:, model:, input_tokens: billed['input_tokens'])
+        end
+
+        # embed-v4 takes mixed text and images as `inputs`; text-only requests
+        # keep using the simpler `texts` array every Embed model accepts.
+        def embedding_inputs(text, attachments)
+          return { texts: Support::Utils.to_safe_array(text).map(&:to_s) } if attachments.empty?
+
+          raise ArgumentError, 'embed one text at a time when embedding attachments' if text.is_a?(Array)
+
+          { inputs: [{ content: Media.format_content(text, attachments) }] }
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/cohere/media.rb

@@ -0,0 +1,77 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Cohere
+      # Handles formatting of media content for the Cohere v2 API
+      module Media
+        # Cohere resolves a remote image by its extension and rejects a URL
+        # without one, so those are inlined from the bytes already fetched.
+        URL_IMAGE_EXTENSIONS = %w[png jpeg jpg gif webp].freeze
+
+        module_function
+
+        # Cohere user messages accept only text and image_url blocks. Text
+        # attachments become citable documents when citations are on, so they
+        # are pulled out of the message and left out of the content blocks.
+        def format_content(content, attachments = [], citations: false)
+          parts = []
+          parts << format_text(content) if content
+
+          attachments.each do |attachment|
+            case attachment.type
+            when :image
+              parts << format_image(attachment)
+            when :text
+              parts << format_text(attachment.for_llm) unless citations
+            else
+              raise UnsupportedAttachmentError, attachment.mime_type
+            end
+          end
+
+          parts
+        end
+
+        def format_text(text)
+          {
+            type: 'text',
+            text: text
+          }
+        end
+
+        def format_image(image)
+          {
+            type: 'image_url',
+            image_url: { url: image_source(image) }
+          }
+        end
+
+        def image_source(image)
+          return image.source.to_s if image.url? && URL_IMAGE_EXTENSIONS.include?(image.extension)
+
+          "data:#{image.mime_type};base64,#{image.encoded}"
+        end
+
+        # Text attachments across the conversation become the request's
+        # top-level documents array, which is what Cohere cites against.
+        def format_documents(messages)
+          index = -1
+
+          messages.flat_map do |message|
+            message.attachments.select(&:text?).map do |attachment|
+              index += 1
+              {
+                id: "doc:#{index}",
+                data: { title: attachment.filename, text: attachment.content }.compact
+              }
+            end
+          end
+        end
+
+        def documents?(messages)
+          messages.any? { |message| message.attachments.any?(&:text?) }
+        end
+      end
+    end
+  end
+end