ruby_llm 1.16.0 → 2.0.0.rc4

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

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

@@ -0,0 +1,92 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Cohere
+      # Model information for the Cohere API
+      module Models
+        module_function
+
+        FEATURE_CAPABILITIES = {
+          'tools' => 'function_calling',
+          'tool_choice' => 'tool_choice',
+          'json_schema' => 'structured_output',
+          'json_mode' => 'json_mode',
+          'citations' => 'citations'
+        }.freeze
+
+        # The model catalog is the one endpoint Cohere still serves from v1.
+        def models_url
+          'v1/models?page_size=1000'
+        end
+
+        def parse_list_models_response(response, slug)
+          Array(response.body['models']).reject { |model| model['is_deprecated'] }.map do |model_data|
+            model_id = model_data['name']
+            endpoints = Array(model_data['endpoints'])
+
+            Model.new(
+              id: model_id,
+              name: model_id,
+              provider: slug,
+              context_window: model_data['context_length']&.to_i,
+              modalities: modalities_from(endpoints),
+              capabilities: capabilities_from(endpoints, model_data['features']),
+              metadata: {
+                endpoints: endpoints,
+                default_endpoints: Array(model_data['default_endpoints']),
+                features: Array(model_data['features']),
+                finetuned: model_data['finetuned'],
+                tokenizer_url: model_data['tokenizer_url']
+              }.compact
+            )
+          end
+        end
+
+        def modalities_from(endpoints)
+          return { input: ['image'], output: ['text'] } if endpoints.include?('parse')
+          return { input: ['audio'], output: ['text'] } if transcription_endpoint?(endpoints)
+
+          return embedding_modalities(endpoints) if embedding_endpoint?(endpoints)
+          return { input: ['text'], output: ['rerank'] } if rerank_endpoint?(endpoints)
+
+          { input: ['text'], output: ['text'] }
+        end
+
+        def capabilities_from(endpoints, features)
+          return ['ocr'] if endpoints.include?('parse')
+          return ['transcription'] if transcription_endpoint?(endpoints)
+          return [] if embedding_endpoint?(endpoints) || rerank_endpoint?(endpoints)
+
+          features = Array(features)
+          capabilities = []
+          capabilities << 'streaming' if endpoints.intersect?(%w[chat generate])
+          reported = FEATURE_CAPABILITIES.filter_map do |feature, capability|
+            capability if features.include?(feature)
+          end
+          capabilities.concat(reported)
+          capabilities
+        end
+
+        def embedding_endpoint?(endpoints)
+          endpoints.any? { |endpoint| endpoint.start_with?('embed') }
+        end
+
+        def embedding_modalities(endpoints)
+          input = []
+          input << 'text' if endpoints.include?('embed')
+          input << 'image' if endpoints.any? { |endpoint| endpoint.start_with?('embed_image') }
+          { input: input.empty? ? ['text'] : input, output: ['embeddings'] }
+        end
+
+        def rerank_endpoint?(endpoints)
+          endpoints.any? { |endpoint| endpoint.start_with?('rerank') }
+        end
+
+        def transcription_endpoint?(endpoints)
+          endpoints.any? { |endpoint| endpoint.start_with?('transcri') }
+        end
+      end
+    end
+  end
+end