ruby_llm 1.15.0 → 2.0.0.rc1

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

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

@@ -0,0 +1,63 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Cohere
+      # Cohere Parse accepts one document image and returns Markdown or blocks.
+      module OCR
+        module_function
+
+        def ocr_url
+          'v2/parse'
+        end
+
+        def render_ocr_payload(file, model:, pages: nil, provider_options: {})
+          attachment = file.is_a?(Attachment) ? file : Attachment.new(file, config: @config)
+          raise UnsupportedAttachmentError, attachment.mime_type unless attachment.image?
+          raise ArgumentError, 'Cohere Parse accepts one image; pages must be [0]' unless pages.nil? || pages == [0]
+
+          reference = attachment.url? ? attachment.source.to_s : attachment.for_llm
+          {
+            model: model,
+            document: { type: 'image_url', image_url: reference },
+            output_format: 'markdown'
+          }.merge(provider_options)
+        end
+
+        def parse_ocr_response(response, model:)
+          data = response.body
+          RubyLLM::OCR.new(
+            pages: data.fetch('pages').map { |page| parse_ocr_page(page) },
+            model: model,
+            usage: data.dig('meta', 'billed_units'),
+            raw: data
+          )
+        end
+
+        def parse_ocr_page(page)
+          content = page['markdown'] || parse_ocr_blocks(page.fetch('blocks'))
+          RubyLLM::OCR::Page.new(
+            index: page['index'], markdown: content['content'],
+            images: content['images'], tables: content['tables'], raw: page
+          )
+        end
+
+        def parse_ocr_blocks(blocks)
+          {
+            'content' => blocks.filter_map { |block| parse_ocr_block_text(block) }.join("\n\n"),
+            'images' => blocks.filter_map { |block| block['image'] },
+            'tables' => blocks.filter_map { |block| block['table'] }
+          }
+        end
+
+        def parse_ocr_block_text(block)
+          case block['type']
+          when 'text' then block.dig('text', 'content')
+          when 'table' then block.dig('table', 'html')
+          when 'image' then "![#{block.dig('image', 'description')}](#{block.dig('image', 'id')})"
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,52 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Cohere
+      # Rerank methods for the Cohere v2 API integration
+      module Rerank
+        module_function
+
+        def rerank_url
+          'v2/rerank'
+        end
+
+        def render_rerank_payload(query, documents, model:, top_n: nil, provider_options: {})
+          {
+            model: model,
+            query: query,
+            documents: documents,
+            top_n: top_n
+          }.compact.merge(provider_options)
+        end
+
+        # Cohere returns positions and scores only, so the ranked text comes
+        # back from the documents that were sent.
+        def parse_rerank_response(response, model:, documents: [])
+          data = response.body
+          billed = data.dig('meta', 'billed_units') || {}
+          tokens = data.dig('meta', 'tokens') || {}
+
+          RubyLLM::Rerank.new(
+            results: parse_rerank_results(data, documents),
+            model: model,
+            raw: data,
+            input_tokens: tokens['input_tokens'] || billed['input_tokens']
+          )
+        end
+
+        def parse_rerank_results(data, documents)
+          Array(data['results']).map do |result|
+            index = result['index']
+
+            RubyLLM::Rerank::Result.new(
+              index: index,
+              document: result.dig('document', 'text') || (index && documents[index]),
+              score: result['relevance_score']
+            )
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,105 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Cohere
+      # Streaming methods of the Cohere v2 API integration
+      module Streaming
+        private
+
+        def stream_url
+          completion_url
+        end
+
+        def build_chunk(data)
+          type = data['type']
+          track_content_lengths(data, type)
+
+          Chunk.new(
+            role: :assistant,
+            model: model&.id,
+            content: extract_content_delta(data, type),
+            citations: extract_citations_delta(data, type),
+            thinking: Thinking.build(text: extract_thinking_delta(data, type)),
+            tool_calls: extract_tool_calls(data, type),
+            finish_reason: normalize_finish_reason(data.dig('delta', 'finish_reason')),
+            **stream_usage(data.dig('delta', 'usage'))
+          )
+        end
+
+        # Citation spans are offsets into their own content block, so the
+        # running length of every earlier text block is what turns them into
+        # offsets into the accumulated response content.
+        def track_content_lengths(data, type)
+          return unless type == 'content-delta'
+
+          text = data.dig('delta', 'message', 'content', 'text')
+          return unless text
+
+          @stream_content_lengths ||= Hash.new(0)
+          @stream_content_lengths[data['index']] += text.length
+        end
+
+        def content_offset(index)
+          return 0 unless @stream_content_lengths
+
+          @stream_content_lengths.sum { |block_index, length| block_index < index ? length : 0 }
+        end
+
+        def extract_content_delta(data, type)
+          data.dig('delta', 'message', 'content', 'text') if type == 'content-delta'
+        end
+
+        def extract_thinking_delta(data, type)
+          case type
+          when 'content-delta' then data.dig('delta', 'message', 'content', 'thinking')
+          when 'tool-plan-delta' then data.dig('delta', 'message', 'tool_plan')
+          end
+        end
+
+        def extract_citations_delta(data, type)
+          return nil unless type == 'citation-start'
+
+          citation = data.dig('delta', 'message', 'citations')
+          return nil unless citation
+
+          index = citation['content_index'] || 0
+          [Chat.parse_citation(citation, { index => content_offset(index) })]
+        end
+
+        def extract_tool_calls(data, type)
+          case type
+          when 'tool-call-start' then stream_tool_call_start(data)
+          when 'tool-call-delta' then stream_tool_call_delta(data)
+          end
+        end
+
+        def stream_tool_call_start(data)
+          tool_call = data.dig('delta', 'message', 'tool_calls')
+          return nil unless tool_call
+
+          {
+            data['index'] => ToolCall.new(
+              id: tool_call['id'],
+              name: tool_call.dig('function', 'name'),
+              arguments: tool_call.dig('function', 'arguments').to_s
+            )
+          }
+        end
+
+        def stream_tool_call_delta(data)
+          arguments = data.dig('delta', 'message', 'tool_calls', 'function', 'arguments')
+          return nil unless arguments
+
+          { data['index'] => ToolCall.new(id: nil, name: nil, arguments: arguments) }
+        end
+
+        def stream_usage(usage)
+          return {} unless usage
+
+          Chat.usage_tokens(usage)
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Cohere
+      # Cohere's plain-text tokenizer is served by the v1 API.
+      module Tokenization
+        def tokenization_url
+          'v1/tokenize'
+        end
+
+        def render_tokenization_payload(text, model:)
+          { model:, text: }
+        end
+
+        def parse_tokenization_response(response, model:)
+          RubyLLM::Tokenization.new(ids: response.body.fetch('tokens'), model:, raw: response.body)
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,132 @@
+# frozen_string_literal: true
+
+require 'json'
+
+module RubyLLM
+  module Protocols
+    class Cohere
+      # Tools methods of the Cohere v2 API integration
+      module Tools
+        EMPTY_PARAMETERS_SCHEMA = {
+          'type' => 'object',
+          'properties' => {},
+          'required' => []
+        }.freeze
+
+        module_function
+
+        def function_for(tool)
+          definition = {
+            type: 'function',
+            function: {
+              name: tool.name,
+              description: tool.description,
+              parameters: parameters_schema_for(tool)
+            }.compact
+          }
+
+          return definition if tool.provider_options.empty?
+
+          RubyLLM::Support::Utils.deep_merge(definition, tool.provider_options)
+        end
+
+        # Cohere takes plain JSON Schema here and controls strictness with the
+        # request's own strict_tools flag, so the OpenAI strict keyword goes.
+        def parameters_schema_for(tool)
+          schema = tool.parameters_schema ||
+                   RubyLLM::Tool::SchemaDefinition.from_parameters(tool.declared_parameters)&.json_schema
+          return EMPTY_PARAMETERS_SCHEMA unless schema
+
+          schema = RubyLLM::Support::Utils.deep_dup(schema)
+          schema.delete(:strict)
+          schema.delete('strict')
+          schema
+        end
+
+        # Cohere only forces tool use on or off; it cannot pin a named tool.
+        def build_tool_choice(choice)
+          case choice
+          when nil, :auto then nil
+          when :none then 'NONE'
+          when :required then 'REQUIRED'
+          else
+            raise ArgumentError,
+                  "Cohere tool choice accepts :auto, :required, or :none, got #{choice.inspect}"
+          end
+        end
+
+        def format_tool_calls(tool_calls)
+          tool_calls.map do |_, tool_call|
+            {
+              id: tool_call.id,
+              type: 'function',
+              function: {
+                name: tool_call.name,
+                arguments: JSON.generate(tool_call.arguments)
+              }
+            }
+          end
+        end
+
+        def format_tool_result(msg)
+          {
+            role: 'tool',
+            tool_call_id: msg.tool_call_id,
+            content: format_tool_result_content(msg)
+          }
+        end
+
+        # Search results become document content blocks, the shape Cohere
+        # cites tool output against. A tool that attaches files sends its
+        # text and each attachment as blocks, so the model reads what the
+        # tool fetched rather than only the sentence describing it.
+        def format_tool_result_content(msg)
+          search_results = RubyLLM::SearchResults.from_content(msg.content)
+          return document_blocks(search_results) if search_results
+          return msg.content.to_s if msg.attachments.empty?
+
+          [Media.format_text(msg.content.to_s)] + msg.attachments.map do |attachment|
+            raise UnsupportedAttachmentError, attachment.mime_type unless attachment.type == :text
+
+            Media.format_text(attachment.for_llm)
+          end
+        end
+
+        def document_blocks(search_results)
+          search_results.results.map.with_index do |result, index|
+            {
+              type: 'document',
+              document: {
+                id: result[:url] || "tool:#{index}",
+                data: { title: result[:title], text: result[:text], url: result[:url] }.compact
+              }
+            }
+          end
+        end
+
+        def parse_tool_calls(tool_calls, response: nil, finish_reason: nil)
+          return nil unless tool_calls&.any?
+
+          tool_calls.to_h do |tool_call|
+            [
+              tool_call['id'],
+              ToolCall.new(
+                id: tool_call['id'],
+                name: tool_call.dig('function', 'name'),
+                arguments: parse_arguments(tool_call.dig('function', 'arguments'), response, finish_reason)
+              )
+            ]
+          end
+        end
+
+        def parse_arguments(arguments, response, finish_reason)
+          return {} if arguments.nil? || arguments.empty?
+
+          JSON.parse(arguments)
+        rescue JSON::ParserError => e
+          raise ToolCallParseError.new(response: response, finish_reason: finish_reason), cause: e
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,41 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Cohere
+      # Audio transcription methods for the Cohere v2 API integration
+      module Transcription
+        DEFAULT_LANGUAGE = 'en'
+
+        module_function
+
+        def transcription_url
+          'v2/audio/transcriptions'
+        end
+
+        # Cohere parses the multipart body as it streams, so every form field
+        # has to be written before the file part or the request is rejected as
+        # missing them. The file goes in last, after provider options.
+        # rubocop:disable-next Lint/UnusedMethodArgument
+        def render_transcription_payload(file_part, model:, language:, format: nil, speaker_names: nil,
+                                         speaker_references: nil, provider_options: {}, prompt: nil,
+                                         temperature: nil)
+          {
+            model: model,
+            language: language || DEFAULT_LANGUAGE,
+            temperature: temperature
+          }.compact.merge(provider_options).merge(file: file_part)
+        end
+
+        # Cohere Transcribe returns the transcript alone: no timestamps,
+        # segments, or speaker labels.
+        def parse_transcription_response(response, model:)
+          data = response.body
+          return RubyLLM::Transcription.new(text: data, model: model) if data.is_a?(String)
+
+          RubyLLM::Transcription.new(text: data['text'], model: model)
+        end
+      end
+    end
+  end
+end

data/lib/ruby_llm/protocols/cohere.rb

@@ -0,0 +1,21 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    # The Cohere v2 API, a native wire format shared by Cohere's own host and
+    # the platforms that serve Command, Embed, Rerank, and Transcribe models.
+    class Cohere < Protocol
+      include Cohere::Batches
+      include Cohere::Chat
+      include Cohere::Embeddings
+      include Cohere::Media
+      include Cohere::Models
+      include Cohere::OCR
+      include Cohere::Tokenization
+      include Cohere::Rerank
+      include Cohere::Streaming
+      include Cohere::Tools
+      include Cohere::Transcription
+    end
+  end
+end