ruby_llm 2.1.0 → 2.2.0

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

@@ -44,12 +44,12 @@ module RubyLLM
           RubyLLM::Support::Utils.deep_merge(definition, tool.provider_options)
         end
 
-        def format_tool_calls(tool_calls)
+        def render_tool_calls(tool_calls)
           return nil unless tool_calls&.any?
 
-          tool_calls.map do |_, tc|
+          tool_calls.map do |tc|
             call = {
-              id: tc.id,
+              id: tc.provider_call_id,
               type: 'function',
               function: {
                 name: tc.name,
@@ -65,39 +65,26 @@ module RubyLLM
           end
         end
 
-        def parse_tool_call_arguments(tool_call, response: nil, finish_reason: nil)
-          arguments = tool_call.dig('function', 'arguments')
-
-          if arguments.nil? || arguments.empty?
-            {}
-          else
-            JSON.parse(arguments)
-          end
-        rescue JSON::ParserError => e
-          raise ToolCallParseError.new(response: response, finish_reason: finish_reason), cause: e
-        end
-
         # Streaming deltas key by index so id-less argument fragments find
         # their call even when a delta batches or interleaves several calls;
         # complete responses key by id, which tool results join on.
         def parse_tool_calls(tool_calls, parse_arguments: true, response: nil, finish_reason: nil, stream_keys: false)
           return nil unless tool_calls&.any?
 
-          tool_calls.to_h do |tc|
+          calls = tool_calls.map do |tc|
             [
-              stream_keys ? tc['index'] || tc['id'] : tc['id'],
-              ToolCall.new(
-                id: tc['id'],
-                name: tc.dig('function', 'name'),
-                arguments: if parse_arguments
-                             parse_tool_call_arguments(tc, response: response, finish_reason: finish_reason)
-                           else
-                             tc.dig('function', 'arguments')
-                           end,
-                thought_signature: extract_tool_call_thought_signature(tc)
-              )
+              tc['index'] || tc['id'],
+              parse_tool_call(tc, parse_arguments:, response:, finish_reason:)
             ]
           end
+          stream_keys ? calls.to_h : ToolCall.index(calls.map(&:last), by: :provider_call_id)
+        end
+
+        def parse_tool_call(call, parse_arguments:, response:, finish_reason:)
+          arguments = call.dig('function', 'arguments')
+          arguments = parse_arguments ? ToolCall.parse_arguments(arguments, response:, finish_reason:) : { arguments: }
+          ToolCall.new(provider_call_id: call['id'], name: call.dig('function', 'name'), **arguments,
+                       thought_signature: parse_tool_call_thought_signature(call))
         end
 
         def build_tool_choice(tool_choice)
@@ -114,7 +101,7 @@ module RubyLLM
           end
         end
 
-        def extract_tool_call_thought_signature(tool_call)
+        def parse_tool_call_thought_signature(tool_call)
           tool_call.dig('extra_content', 'google', 'thought_signature')
         end
       end

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

@@ -5,6 +5,13 @@ module RubyLLM
     class ChatCompletions
       # Audio transcription methods for the OpenAI API integration
       module Transcription
+        STREAM_TYPES = {
+          'transcript.text.delta' => RubyLLM::TranscriptionChunk::DELTA,
+          'transcript.text.partial' => RubyLLM::TranscriptionChunk::PARTIAL,
+          'transcript.text.segment' => RubyLLM::TranscriptionChunk::SEGMENT,
+          'transcript.text.done' => RubyLLM::TranscriptionChunk::DONE
+        }.freeze
+
         module_function
 
         def transcription_url
@@ -44,7 +51,7 @@ module RubyLLM
           return nil unless references
 
           references.map do |ref|
-            Attachment.new(ref, config: @config).for_llm
+            (ref.is_a?(Attachment) ? ref : Attachment.new(ref, config: @config)).for_llm
           end
         end
 
@@ -74,7 +81,7 @@ module RubyLLM
         end
 
         def build_transcription_chunk(data)
-          type = data['type']
+          type = STREAM_TYPES[data['type']]
 
           RubyLLM::TranscriptionChunk.new(
             type: type,
@@ -96,6 +103,7 @@ module RubyLLM
             language: data['language'],
             duration: transcription_duration(usage),
             segments: streamed_transcription_segments(chunks, data),
+            audio_seconds: billed_audio_seconds(usage),
             reported_cost: reported_cost(usage),
             **transcription_tokens(usage)
           )
@@ -129,6 +137,7 @@ module RubyLLM
             duration: data['duration'] || transcription_duration(usage),
             segments: data['segments'],
             words: data['words'],
+            audio_seconds: billed_audio_seconds(usage, data['duration']),
             reported_cost: reported_cost(usage),
             **transcription_tokens(usage)
           )
@@ -144,6 +153,12 @@ module RubyLLM
         def transcription_duration(usage)
           usage['seconds']
         end
+
+        def billed_audio_seconds(usage, duration = nil)
+          return usage['seconds'] if usage['type'] == 'duration'
+
+          duration if transcription_tokens(usage).values.all?(&:nil?)
+        end
       end
     end
   end

data/lib/ruby_llm/protocols/chat_completions.rb

@@ -4,6 +4,10 @@ module RubyLLM
   module Protocols
     # The OpenAI Chat Completions API — the lingua franca of LLM APIs.
     class ChatCompletions < Protocol
+      # The endpoint paths a base URL of an OpenAI-compatible API must not
+      # end in, since RubyLLM adds them.
+      ENDPOINT_PATHS = %r{/(?:chat/completions|completions|responses|embeddings)\z}
+
       include ChatCompletions::Chat
       include ChatCompletions::Embeddings
       include ChatCompletions::Models

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

@@ -38,7 +38,7 @@ module RubyLLM
           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)
+          model = batch_model(data.fetch('model'))
           parser = self.class.new(@provider, model)
           datasets.records(file).map { |row| parse_batch_result(row, parser:, model: model.id) }
         end
@@ -49,6 +49,12 @@ module RubyLLM
           @datasets ||= Cohere::Datasets.new(@provider)
         end
 
+        def batch_model(id)
+          RubyLLM.models.find(id, provider: @provider.slug, config: @config)
+        rescue ModelNotFoundError
+          Model.default(id, @provider.slug)
+        end
+
         def batch_data(id)
           @connection.get("v2/batches/#{id}").body.fetch('batch')
         end
@@ -80,13 +86,14 @@ module RubyLLM
           body = row.fetch('body')
           return [index, nil, batch_failure(custom_id, 'generation failed')] if body['finish_reason'] == 'ERROR'
 
-          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]
+          billed = body.merge('model' => model) unless body['embeddings']
+          parsed_batch_result(index, custom_id, billed, parser:) do
+            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
+          end
         end
       end
     end

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

@@ -16,21 +16,21 @@ module RubyLLM
         def normalize_finish_reason(reason)
           return nil if reason.nil?
 
-          finish_reasons.fetch(reason.to_s) { reason.to_s.to_sym }
+          finish_reasons.fetch(reason.to_s) { reason.to_s.downcase.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)
+          reject_caching_options(caching)
 
           payload = {
             model: model.id,
-            messages: format_messages(messages, citations: citations),
+            messages: render_messages(messages, citations: citations),
             stream: stream
           }
 
@@ -40,10 +40,18 @@ module RubyLLM
           payload
         end
 
+        def cache_option_names = []
+
+        def reject_caching_options(caching)
+          return unless caching.is_a?(Hash) && caching.any?
+
+          raise ArgumentError, "Cohere caches on its own and does not accept #{caching.keys.join(', ')} for caching"
+        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[:documents] = Media.render_documents(messages) if citations && Media.documents?(messages)
           payload[:response_format] = build_response_format(schema) if schema
         end
 
@@ -81,28 +89,28 @@ module RubyLLM
           payload[:thinking] = { type: 'enabled', token_budget: thinking.budget }.compact
         end
 
-        def format_messages(messages, citations: false)
-          messages.map { |msg| format_message(msg, citations: citations) }
+        def render_messages(messages, citations: false)
+          messages.map { |msg| render_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
+        def render_message(msg, citations: false)
+          return Tools.render_tool_result(msg) if msg.tool_result?
+          return render_assistant_message(msg) if msg.role == :assistant
 
           {
             role: msg.role.to_s,
-            content: Media.format_content(msg.content, msg.attachments, citations: citations)
+            content: Media.render_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)
+        def render_assistant_message(msg)
           message = { role: 'assistant' }
-          content = Media.format_content(msg.content, msg.attachments)
+          content = Media.render_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[:tool_calls] = Tools.render_tool_calls(msg.tool_calls) if msg.tool_call?
           message
         end
 
@@ -111,46 +119,53 @@ module RubyLLM
 
           message_data = data['message'] || {}
           blocks = Array(message_data['content'])
-          content, offsets = extract_text(blocks)
+          content, offsets = parse_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)),
+            thinking: Thinking.build(text: parse_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'] || {})
+            **parse_completion_usage(data)
           )
         end
 
+        def parse_completion_usage(data)
+          usage_tokens(data['usage'] || {})
+        end
+
         # Cohere answers a generation that failed on its side with an ERROR
         # finish reason instead of an error status.
         def raise_failed_generation(message, response: nil)
-          raise ServerError.new(message || 'The generation failed due to an internal error', response:)
+          raise ServerError.new('The generation failed due to an internal error', response:) unless message
+
+          raise ServerError.new(provider_message: message, response:, redact: @config&.redact_error_messages)
         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.
+        # Cohere reports what it charges for under billed_units and what the
+        # model processed under tokens. The processed counts and cached_tokens
+        # include Cohere's own preamble, which it does not bill, so a reply's
+        # usage is its billed units.
         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']
+            input_tokens: billed['input_tokens'] || tokens['input_tokens'],
+            output_tokens: billed['output_tokens'] || tokens['output_tokens'],
+            thinking_tokens: tokens['reasoning_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)
+        def parse_text(blocks)
           text = +''
           offsets = {}
 
@@ -166,7 +181,7 @@ module RubyLLM
 
         # 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)
+        def parse_thinking(blocks, message_data)
           thoughts = blocks.select { |block| block['type'] == 'thinking' }
                            .map { |block| block['thinking'] }.join
           thoughts.empty? ? message_data['tool_plan'] : thoughts
@@ -191,6 +206,7 @@ module RubyLLM
             text: data['text'],
             start_index: start_index,
             end_index: end_index,
+            source_id: tool_output_id(source),
             source_index: source_index(source)
           )
         end
@@ -220,6 +236,12 @@ module RubyLLM
           url if url.is_a?(String) && url.match?(%r{\Ahttps?://}i)
         end
 
+        # Cohere cites a tool output by the id its document was sent with,
+        # so a search result comes back with its source_id or URL.
+        def tool_output_id(source)
+          source['id'] if source && source['type'] == 'tool'
+        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)

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

@@ -5,10 +5,14 @@ module RubyLLM
     class Cohere
       # Cohere datasets for batch input and generated results.
       class Datasets < Protocols::Files
+        def rejected_options(operation)
+          operation == :upload ? %i[expires_in uri] : super
+        end
+
+        # rubocop:disable-next Lint/UnusedMethodArgument
         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 }
@@ -88,7 +92,7 @@ module RubyLLM
         end
 
         def download_part(url)
-          Transport::Connection.basic(@config).get(url).body
+          Transport::Connection.download(url, @config)
         end
 
         def load_avro

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

@@ -13,9 +13,11 @@ module RubyLLM
           'v2/embed'
         end
 
-        # rubocop:disable-next Lint/UnusedMethodArgument
-        def render_embedding_payload(text, model:, dimensions:, task_type: nil, title: nil, with: [],
-                                     provider_options: {})
+        def rejected_options(operation)
+          operation == :embed ? %i[title] : super
+        end
+
+        def render_embedding_payload(text, model:, dimensions:, task_type: nil, with: [], provider_options: {}, **)
           image_only = with.any? && separate_image_embeddings?(model)
           payload = {
             model: model,
@@ -28,7 +30,7 @@ module RubyLLM
           Support::Utils.deep_merge(payload, provider_options)
         end
 
-        def supports_embedding_media?
+        def embedding_media?
           true
         end
 
@@ -63,7 +65,7 @@ module RubyLLM
 
           raise ArgumentError, 'embed one text at a time when embedding attachments' if text.is_a?(Array)
 
-          { inputs: [{ content: Media.format_content(text, attachments) }] }
+          { inputs: [{ content: Media.render_content(text, attachments) }] }
         end
       end
     end

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

@@ -18,32 +18,32 @@ module RubyLLM
         # 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)
+        def render_content(content, attachments = [], citations: false)
           parts = []
-          parts << format_text(content) if content
+          parts << render_text(content) if content
 
           attachments.each do |attachment|
             raise UnsupportedAttachmentError, attachment.mime_type unless supported_attachment?(attachment)
 
             case attachment.type
             when :image
-              parts << format_image(attachment)
+              parts << render_image(attachment)
             when :text
-              parts << format_text(attachment.for_llm) unless citations
+              parts << render_text(attachment.for_llm) unless citations
             end
           end
 
           parts
         end
 
-        def format_text(text)
+        def render_text(text)
           {
             type: 'text',
             text: text
           }
         end
 
-        def format_image(image)
+        def render_image(image)
           {
             type: 'image_url',
             image_url: { url: image_source(image) }
@@ -58,7 +58,7 @@ module RubyLLM
 
         # Text attachments across the conversation become the request's
         # top-level documents array, which is what Cohere cites against.
-        def format_documents(messages)
+        def render_documents(messages)
           index = -1
 
           messages.flat_map do |message|

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

@@ -31,7 +31,8 @@ module RubyLLM
             results: parse_rerank_results(data, documents),
             model: model,
             raw: data,
-            input_tokens: tokens['input_tokens'] || billed['input_tokens']
+            input_tokens: tokens['input_tokens'] || billed['input_tokens'],
+            search_units: billed['search_units']
           )
         end
 

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

@@ -11,23 +11,32 @@ module RubyLLM
           completion_url
         end
 
+        def stream_response(...)
+          @stream_thinking = false
+          super
+        end
+
         def build_chunk(data)
           type = data['type']
-          raise_failed_generation(data.dig('delta', 'error')) if data.dig('delta', 'finish_reason') == 'ERROR'
+          raise_failed_stream(data) if data.dig('delta', 'finish_reason') == 'ERROR'
           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)),
+            content: parse_content_delta(data, type),
+            citations: parse_citations_delta(data, type),
+            thinking: Thinking.build(text: parse_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
 
+        def raise_failed_stream(data)
+          raise_failed_generation(data.dig('delta', 'error'), response: build_stream_error_response(data, nil, 500))
+        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.
@@ -47,18 +56,24 @@ module RubyLLM
           @stream_content_lengths.sum { |block_index, length| block_index < index ? length : 0 }
         end
 
-        def extract_content_delta(data, type)
+        def parse_content_delta(data, type)
           data.dig('delta', 'message', 'content', 'text') if type == 'content-delta'
         end
 
-        def extract_thinking_delta(data, type)
+        def parse_thinking_block_delta(data)
+          thinking = data.dig('delta', 'message', 'content', 'thinking')
+          @stream_thinking = true if thinking
+          thinking
+        end
+
+        def parse_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')
+          when 'content-delta' then parse_thinking_block_delta(data)
+          when 'tool-plan-delta' then data.dig('delta', 'message', 'tool_plan') unless @stream_thinking
           end
         end
 
-        def extract_citations_delta(data, type)
+        def parse_citations_delta(data, type)
           return nil unless type == 'citation-start'
 
           citation = data.dig('delta', 'message', 'citations')
@@ -81,7 +96,7 @@ module RubyLLM
 
           {
             data['index'] => ToolCall.new(
-              id: tool_call['id'],
+              provider_call_id: tool_call['id'],
               name: tool_call.dig('function', 'name'),
               arguments: tool_call.dig('function', 'arguments').to_s
             )
@@ -92,7 +107,7 @@ module RubyLLM
           arguments = data.dig('delta', 'message', 'tool_calls', 'function', 'arguments')
           return nil unless arguments
 
-          { data['index'] => ToolCall.new(id: nil, name: nil, arguments: arguments) }
+          { data['index'] => ToolCall.new(name: nil, arguments: arguments) }
         end
 
         def stream_usage(usage)

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

@@ -55,10 +55,10 @@ module RubyLLM
           end
         end
 
-        def format_tool_calls(tool_calls)
-          tool_calls.map do |_, tool_call|
+        def render_tool_calls(tool_calls)
+          tool_calls.map do |tool_call|
             {
-              id: tool_call.id,
+              id: tool_call.provider_call_id,
               type: 'function',
               function: {
                 name: tool_call.name,
@@ -68,11 +68,11 @@ module RubyLLM
           end
         end
 
-        def format_tool_result(msg)
+        def render_tool_result(msg)
           {
             role: 'tool',
-            tool_call_id: msg.tool_call_id,
-            content: format_tool_result_content(msg)
+            tool_call_id: msg.provider_call_id,
+            content: render_tool_result_content(msg)
           }
         end
 
@@ -84,15 +84,15 @@ module RubyLLM
         # 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)
+        def render_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|
+          [Media.render_text(msg.content.to_s)] + msg.attachments.map do |attachment|
             raise UnsupportedAttachmentError, attachment.mime_type unless supported_message_attachment?(msg, attachment)
 
-            Media.format_text(attachment.for_llm)
+            Media.render_text(attachment.for_llm)
           end
         end
 
@@ -101,7 +101,7 @@ module RubyLLM
             {
               type: 'document',
               document: {
-                id: result[:url] || "tool:#{index}",
+                id: result[:source_id] || result[:url] || "tool:#{index}",
                 data: { title: result[:title], text: result[:text], url: result[:url] }.compact
               }
             }
@@ -111,24 +111,13 @@ module RubyLLM
         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
+          ToolCall.index(tool_calls.map do |tool_call|
+            ToolCall.new(
+              provider_call_id: tool_call['id'],
+              name: tool_call.dig('function', 'name'),
+              **ToolCall.parse_arguments(tool_call.dig('function', 'arguments'), response:, finish_reason:)
+            )
+          end, by: :provider_call_id)
         end
       end
     end