ruby_llm 1.16.0 → 2.0.0.rc2

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

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

@@ -0,0 +1,55 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Protocols
+    class Converse
+      # Bedrock Model Invocation Jobs for Converse payloads.
+      module Batches
+        include Protocols::Bedrock::Batches
+
+        private
+
+        def bedrock_invocation_type
+          'Converse'
+        end
+
+        def bedrock_batch_jsonl(requests)
+          requests.map do |request|
+            JSON.generate(recordId: request[:custom_id], modelInput: batch_payload(request))
+          end.join("\n")
+        end
+
+        def validate_bedrock_batch_requests!(requests)
+          unsupported = requests.any? do |request|
+            payload = request.fetch(:payload)
+            payload.key?(:toolConfig) || payload.key?('toolConfig') ||
+              payload.key?(:outputConfig) || payload.key?('outputConfig')
+          end
+          return unless unsupported
+
+          raise Error, 'bedrock batch requests do not support tools or structured output'
+        end
+
+        def parse_bedrock_outputs(outputs, **)
+          outputs.flat_map { |body| parse_bedrock_output(body) }
+        end
+
+        def parse_bedrock_output(body)
+          body.to_s.each_line.filter_map { |line| parse_bedrock_record(JSON.parse(line)) }
+        end
+
+        def parse_bedrock_record(line)
+          record_id = line['recordId'] || line['record_id']
+          index = batch_result_index(record_id)
+
+          if line['modelOutput']
+            body = line['modelOutput']
+            [index, parse_completion_body(body, raw: body)]
+          else
+            [index, nil, batch_failure(record_id, line.dig('error', 'message') || line['errorMessage'])]
+          end
+        end
+      end
+    end
+  end
+end