Add initial Cohere Chat and Embeddings module implementations

Author: infinityrobot

lib/ruby_llm/providers/cohere.rb

@@ -13,7 +13,14 @@ module Providers
     # See https://docs.cohere.com/docs/compatibility-api for more information.
     module Cohere
       extend Provider
+      extend Cohere::Chat
+      extend Cohere::Embeddings
+      extend Cohere::Reranking
       extend Cohere::Models
+      extend Cohere::Streaming
+      extend Cohere::Tools
+      extend Cohere::Media
+
       module_function
 
       def api_base(_config)

lib/ruby_llm/providers/cohere/chat.rb

@@ -0,0 +1,83 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    module Cohere
+      # Chat methods of the Cohere API integration
+      # - https://docs.cohere.com/reference/chat
+      # - https://docs.cohere.com/docs/chat-api
+      module Chat
+        def completion_url
+          'v2/chat'
+        end
+
+        module_function
+
+        def render_payload(messages, tools:, temperature:, model:, stream: false)
+          @model_id = model
+
+          {
+            model: model,
+            messages: format_messages(messages),
+            temperature: temperature,
+            stream: stream,
+            tools: tools.any? ? tools.map { |_, tool| Tools.tool_for(tool) } : nil
+          }.compact
+        end
+
+        def parse_completion_response(response)
+          data = response.body
+          return if data.empty?
+
+          raise Error.new(response, data['message']) if data['message'] && response.status != 200
+
+          message_data = data['message']
+          return unless message_data
+
+          Message.new(
+            role: message_data['role'].to_sym,
+            content: message_data.dig('content', 0, 'text'),
+            tool_calls: Tools.parse_tool_calls(message_data['tool_calls']),
+            input_tokens: data.dig('usage', 'tokens', 'input_tokens'),
+            output_tokens: data.dig('usage', 'tokens', 'output_tokens'),
+            model_id: @model_id
+          )
+        end
+
+        def format_messages(messages)
+          messages.map { |msg| format_message(msg) }
+        end
+
+        def format_message(msg)
+          if msg.tool_call?
+            Tools.format_tool_call(msg)
+          elsif msg.tool_result?
+            Tools.format_tool_result(msg)
+          else
+            format_basic_message(msg)
+          end
+        end
+
+        def format_basic_message(msg)
+          {
+            role: format_role(msg.role),
+            content: Media.format_content(msg.content)
+          }.compact
+        end
+
+        def format_role(role)
+          case role
+          when :system
+            'system'
+          when :user, :tool
+            'user'
+          when :assistant
+            'assistant'
+          else
+            role.to_s
+          end
+        end
+      end
+    end
+  end
+end

lib/ruby_llm/providers/cohere/embeddings.rb

@@ -0,0 +1,42 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    module Cohere
+      # Embeddings methods of the Cohere API integration
+      # - https://docs.cohere.com/reference/embed
+      # - https://docs.cohere.com/docs/embeddings
+      module Embeddings
+        module_function
+
+        def embedding_url(...)
+          'v2/embed'
+        end
+
+        def render_embedding_payload(input, model:, dimensions: nil)
+          {
+            model: model,
+            embedding_types: ['float'],
+            texts: Array(input),
+            input_type: 'search_document',
+            output_dimension: dimensions,
+            truncate: 'END' # Handle long texts by truncating at the end
+          }
+        end
+
+        def parse_embedding_response(response, model:)
+          data = response.body
+          raise Error.new(response, data['message']) if data['message'] && response.status != 200
+
+          vectors = data.dig('embeddings', 'float') || []
+          input_tokens = data.dig('meta', 'billed_units', 'input_tokens') || 0
+
+          # If we only got one embedding, return it as a single vector
+          vectors = vectors.first if vectors.length == 1
+
+          Embedding.new(vectors:, model:, input_tokens:)
+        end
+      end
+    end
+  end
+end

lib/ruby_llm/providers/cohere/media.rb

@@ -0,0 +1,67 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    module Cohere
+      # Handles formatting of media content (images) for Cohere APIs
+      # Supports Aya Vision models with multimodal capabilities
+      module Media
+        module_function
+
+        def format_content(content)
+          return [format_text(content)] unless content.is_a?(RubyLLM::Content)
+
+          parts = []
+          parts << format_text(content.text) if content.text
+
+          content.attachments.each do |attachment|
+            case attachment.type
+            when :image
+              parts << format_image(attachment)
+            when :text
+              parts << format_text_file(attachment)
+            else
+              raise RubyLLM::UnsupportedAttachmentError, attachment.type
+            end
+          end
+
+          parts
+        end
+
+        def format_text(text)
+          {
+            type: 'text',
+            text: text
+          }
+        end
+
+        def format_image(image)
+          if image.url?
+            # Use URL directly for Cohere API
+            {
+              type: 'image_url',
+              image_url: {
+                url: image.source
+              }
+            }
+          else
+            # Use base64 encoding for local images
+            {
+              type: 'image_url',
+              image_url: {
+                url: "data:#{image.mime_type};base64,#{image.encoded}"
+              }
+            }
+          end
+        end
+
+        def format_text_file(text_file)
+          {
+            type: 'text',
+            text: Utils.format_text_file_for_llm(text_file)
+          }
+        end
+      end
+    end
+  end
+end

lib/ruby_llm/providers/cohere/streaming.rb

@@ -0,0 +1,101 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  module Providers
+    module Cohere
+      # Streaming methods of the Cohere API integration
+      # - https://docs.cohere.com/docs/streaming
+      module Streaming
+        private
+
+        def stream_url
+          completion_url
+        end
+
+        def build_chunk(data)
+          Chunk.new(
+            role: :assistant,
+            model_id: extract_model_id(data),
+            content: extract_content(data),
+            input_tokens: extract_input_tokens(data),
+            output_tokens: extract_output_tokens(data),
+            tool_calls: extract_tool_calls(data)
+          )
+        end
+
+        def extract_model_id(data)
+          data['response_id'] || data['id']
+        end
+
+        def extract_content(data)
+          case data['type']
+          when 'content-delta'
+            data.dig('delta', 'message', 'content', 'text')
+          when 'message-end'
+            # Final message content
+            data.dig('delta', 'message', 'content', 0, 'text')
+          end
+        end
+
+        def extract_input_tokens(data)
+          return unless data['type'] == 'message-end'
+
+          data.dig('delta', 'usage', 'tokens', 'input_tokens')
+        end
+
+        def extract_output_tokens(data)
+          return unless data['type'] == 'message-end'
+
+          data.dig('delta', 'usage', 'tokens', 'output_tokens')
+        end
+
+        def extract_tool_calls(data)
+          case data['type']
+          when 'tool-call-start'
+            tool_call_data = data.dig('delta', 'message', 'tool_calls')
+            return {} unless tool_call_data
+
+            tool_call = ToolCall.new(
+              id: tool_call_data['id'],
+              name: tool_call_data.dig('function', 'name'),
+              arguments: tool_call_data.dig('function', 'arguments') || ''
+            )
+            { tool_call.id => tool_call }
+          when 'tool-call-delta'
+            # Handle streaming tool call arguments
+            argument_delta = data.dig('delta', 'message', 'tool_calls', 'function', 'arguments')
+            return {} unless argument_delta
+
+            { nil => ToolCall.new(id: nil, name: nil, arguments: argument_delta) }
+          when 'message-end'
+            tool_calls = data.dig('delta', 'message', 'tool_calls')
+            return {} unless tool_calls
+
+            result = {}
+            tool_calls.each do |call|
+              tool_call = ToolCall.new(
+                id: call['id'],
+                name: call.dig('function', 'name'),
+                arguments: call.dig('function', 'parameters')
+              )
+              result[tool_call.id] = tool_call
+            end
+            result
+          else
+            {}
+          end
+        end
+
+        def parse_streaming_error(data)
+          error_data = JSON.parse(data)
+          return unless error_data['type'] == 'error'
+
+          message = error_data.dig('error', 'message') || 'Unknown error'
+          [500, message]
+        rescue JSON::ParserError
+          [500, 'Failed to parse error response']
+        end
+      end
+    end
+  end
+end