llm_gateway 0.5.0 → 0.7.0

data/lib/llm_gateway/adapters/anthropic/stream_mapper.rb

@@ -1,5 +1,7 @@
 # frozen_string_literal: true
 
+require "json"
+
 require_relative "../stream_mapper"
 
 module LlmGateway
@@ -14,25 +16,36 @@ module LlmGateway
               model: chunk.dig(:data, :message, :model),
               role: chunk.dig(:data, :message, :role)
             }
-            usage_increment = chunk.dig(:data, :message, :usage) || {}
-
-            accumulator.push({ type: :message_start, usage_increment:, delta: }, &block)
+            accumulator.push({ type: :message_start, delta: }, &block)
           when "content_block_start"
             content_block = chunk.dig(:data, :content_block) || {}
-            @current_content_block_type = content_block[:type]
+            @current_content_block_type = normalize_content_block_type(content_block[:type])
 
             case @current_content_block_type
             when "thinking"
               accumulator.push({ type: :reasoning_start, delta: content_block[:thinking], signature: "" }, &block)
             when "text"
               accumulator.push({ type: :text_start, delta: content_block[:text] }, &block)
-            when "tool_use"
+            when "tool_use", "server_tool_use"
               accumulator.push(
                 {
                   type: :tool_start,
                   delta: "",
                   id: content_block[:id],
-                  name: content_block[:name]
+                  name: content_block[:name],
+                  tool_type: @current_content_block_type
+                },
+                &block
+              )
+            when "server_tool_result"
+              content = content_block[:content]
+              result_delta = content.nil? ? "" : JSON.generate(content)
+              accumulator.push(
+                {
+                  type: :tool_result_start,
+                  delta: result_delta,
+                  tool_use_id: content_block[:tool_use_id],
+                  name: content_block[:type]
                 },
                 &block
               )
@@ -46,9 +59,13 @@ module LlmGateway
             when "text"
               delta = chunk.dig(:data, :delta, :text)
               accumulator.push({ type: :text_delta, delta: }, &block)
-            when "tool_use"
+            when "tool_use", "server_tool_use"
               delta = chunk.dig(:data, :delta, :partial_json)
               accumulator.push({ type: :tool_delta, delta: }, &block)
+            when "server_tool_result"
+              content = chunk.dig(:data, :delta, :content)
+              result_delta = content.nil? ? "" : JSON.generate(content)
+              accumulator.push({ type: :tool_result_delta, delta: result_delta }, &block)
             end
           when "content_block_stop"
             case @current_content_block_type
@@ -56,16 +73,21 @@ module LlmGateway
               accumulator.push({ type: :reasoning_end, delta: "", signature: "" }, &block)
             when "text"
               accumulator.push({ type: :text_end, delta: "" }, &block)
-            when "tool_use"
+            when "tool_use", "server_tool_use"
               accumulator.push({ type: :tool_end, delta: "" }, &block)
+            when "server_tool_result"
+              accumulator.push({ type: :tool_result_end, delta: "" }, &block)
             end
             @current_content_block_type = nil
           when "message_delta"
-            delta = normalize_message_delta(chunk.dig(:data, :delta) || {})
-            usage_increment = chunk.dig(:data, :usage) || {}
+            data = chunk[:data] || {}
+            delta = normalize_message_delta(data[:delta] || {})
+            patch = { type: :message_delta, delta: }
+            patch[:usage] = normalized_usage(data[:usage]) if data.key?(:usage)
 
-            accumulator.push({ type: :message_delta, usage_increment:, delta: }, &block)
+            accumulator.push(patch, &block)
           when "message_stop"
+
             accumulator.push({ type: :message_end }, &block)
           when "ping"
             nil
@@ -76,6 +98,34 @@ module LlmGateway
 
         private
 
+        def normalized_usage(usage)
+          usage = usage.to_h.symbolize_keys
+
+          input = token_count(usage[:input_tokens])
+          cache_write = token_count(usage[:cache_creation_input_tokens])
+          cache_read = token_count(usage[:cache_read_input_tokens])
+          output = token_count(usage[:output_tokens])
+
+          {
+            input:,
+            cache_write:,
+            cache_read:,
+            output:,
+            total: input + cache_write + cache_read + output,
+            raw: usage
+          }
+        end
+
+        def token_count(value)
+          value.to_i
+        end
+
+        def normalize_content_block_type(type)
+          return type unless type&.end_with?("_tool_result")
+
+          "server_tool_result"
+        end
+
         def normalize_message_delta(delta)
           return delta unless delta[:stop_reason] || delta["stop_reason"]
 

data/lib/llm_gateway/adapters/anthropic_option_mapper.rb

@@ -50,7 +50,7 @@ module LlmGateway
       module_function
 
       def map(options)
-        mapped_options = options.reject { |key, _| MANAGED_OPTIONS.include?(key) }
+        mapped_options = options.except(*MANAGED_OPTIONS)
         mapped_options[:max_tokens] = options[:max_completion_tokens] || DEFAULT_MAX_TOKENS
 
         response_format = options[:response_format]

data/lib/llm_gateway/adapters/groq/option_mapper.rb

@@ -68,7 +68,7 @@ module LlmGateway
         module_function
 
         def map(options)
-          mapped_options = options.reject { |key, _| MANAGED_OPTIONS.include?(key) }
+          mapped_options = options.except(*MANAGED_OPTIONS)
           mapped_options[:temperature] = options.key?(:temperature) ? options[:temperature] : DEFAULT_TEMPERATURE
           mapped_options[:max_completion_tokens] = options[:max_completion_tokens] || DEFAULT_MAX_COMPLETION_TOKENS
           mapped_options[:response_format] = normalize_response_format(options[:response_format] || "text")

data/lib/llm_gateway/adapters/input_message_sanitizer.rb

@@ -1,12 +1,14 @@
 # frozen_string_literal: true
 
+require "json"
+
 module LlmGateway
   module Adapters
     class InputMessageSanitizer
       def self.sanitize(messages, target_provider:, target_api:, target_model:)
         return messages unless messages.is_a?(Array)
 
-        messages.map do |message|
+        sanitized = messages.map do |message|
           sanitize_message(
             message,
             target_provider: target_provider,
@@ -14,6 +16,8 @@ module LlmGateway
             target_model: target_model
           )
         end
+
+        relocate_assistant_tool_results(sanitized)
       end
 
       def self.sanitize_message(message, target_provider:, target_api:, target_model:)
@@ -25,9 +29,14 @@ module LlmGateway
         return message unless message_metadata_present?(message)
 
         same_model_replay = same_model_replay?(message, target_provider:, target_api:, target_model:)
+        same_provider_api_replay = same_provider_api_replay?(message, target_provider:, target_api:)
 
         sanitized_content = content.each_with_object([]) do |block, acc|
-          sanitized = sanitize_content_block(block, same_model_replay: same_model_replay)
+          sanitized = sanitize_content_block(
+            block,
+            same_model_replay: same_model_replay,
+            same_provider_api_replay: same_provider_api_replay
+          )
           next if sanitized.nil?
 
           if sanitized.is_a?(Array)
@@ -40,19 +49,91 @@ module LlmGateway
         message.merge(content: sanitized_content)
       end
 
-      def self.sanitize_content_block(block, same_model_replay:)
+      def self.sanitize_content_block(block, same_model_replay:, same_provider_api_replay:)
         return block unless block.is_a?(Hash)
 
         type = block[:type] || block["type"]
+
+        if type == "server_tool_use"
+          return normalize_server_tool_use_for_replay(block) if same_provider_api_replay
+
+          return convert_server_tool_use_to_tool_use(block)
+        end
+
+        if type == "server_tool_result"
+          return block if same_provider_api_replay
+
+          return convert_server_tool_result_to_tool_result(block)
+        end
+
         return block unless %w[thinking reasoning].include?(type)
         return block if same_model_replay
 
         text = extract_reasoning_text(block)
-        return nil if text.nil? || text.strip.empty?
+        return nil if text.blank?
 
         { type: "text", text: text }
       end
 
+      def self.normalize_server_tool_use_for_replay(block)
+        input = block[:input] || block["input"]
+        return block unless input.is_a?(Hash)
+
+        outputs = input[:outputs] || input["outputs"]
+        return block unless outputs.is_a?(Hash)
+
+        normalized_input = input.merge(outputs: outputs.values)
+        normalized_input.delete(:outputs) if input.key?("outputs") && !input.key?(:outputs)
+        normalized_input["outputs"] = outputs.values if input.key?("outputs")
+
+        normalized = block.merge(input: normalized_input)
+        normalized.delete(:input) if block.key?("input") && !block.key?(:input)
+        normalized["input"] = normalized_input if block.key?("input")
+        normalized
+      end
+
+      def self.convert_server_tool_use_to_tool_use(block)
+        converted = block.merge(type: "tool_use")
+        converted.delete(:type) if block.key?("type") && !block.key?(:type)
+        converted["type"] = "tool_use" if block.key?("type")
+        converted
+      end
+
+      def self.convert_server_tool_result_to_tool_result(block)
+        converted = block.merge(type: "tool_result")
+        converted.delete(:type) if block.key?("type") && !block.key?(:type)
+        converted["type"] = "tool_result" if block.key?("type")
+
+        content = converted[:content] || converted["content"]
+        if content.is_a?(Hash)
+          converted = converted.merge(content: JSON.generate(content))
+          converted.delete(:content) if block.key?("content") && !block.key?(:content)
+          converted["content"] = JSON.generate(content) if block.key?("content")
+        end
+
+        converted
+      end
+
+      def self.relocate_assistant_tool_results(messages)
+        messages.flat_map do |message|
+          next message unless message.is_a?(Hash)
+
+          role = message[:role] || message["role"]
+          content = message[:content] || message["content"]
+          next message unless role == "assistant" && content.is_a?(Array)
+
+          tool_results, assistant_content = content.partition do |block|
+            block.is_a?(Hash) && (block[:type] || block["type"]) == "tool_result"
+          end
+          next message if tool_results.empty?
+
+          relocated = []
+          relocated << message.merge(content: assistant_content) unless assistant_content.empty?
+          relocated << { role: "user", content: tool_results }
+          relocated
+        end
+      end
+
       def self.extract_reasoning_text(block)
         return block[:thinking] if block[:thinking].is_a?(String)
         return block[:reasoning] if block[:reasoning].is_a?(String)
@@ -65,7 +146,7 @@ module LlmGateway
 
             item[:text] || item[:summary_text] || item[:reasoning]
           end.join("\n")
-          return text unless text.empty?
+          return text if text.present?
         end
 
         nil
@@ -79,15 +160,25 @@ module LlmGateway
         provider == target_provider && api == target_api && model == target_model
       end
 
+      def self.same_provider_api_replay?(message, target_provider:, target_api:)
+        provider = message[:provider] || message["provider"]
+        api = message[:api] || message["api"]
+
+        provider == target_provider && api == target_api
+      end
+
       def self.message_metadata_present?(message)
         provider = message[:provider] || message["provider"]
         api = message[:api] || message["api"]
         model = message[:model] || message["model"]
 
-        !provider.nil? && !api.nil? && !model.nil?
+        provider.present? && api.present? && model.present?
       end
 
-      private_class_method :sanitize_message, :sanitize_content_block, :extract_reasoning_text, :same_model_replay?, :message_metadata_present?
+      private_class_method :sanitize_message, :sanitize_content_block, :normalize_server_tool_use_for_replay,
+        :convert_server_tool_use_to_tool_use, :convert_server_tool_result_to_tool_result,
+        :relocate_assistant_tool_results, :extract_reasoning_text, :same_model_replay?,
+        :same_provider_api_replay?, :message_metadata_present?
     end
   end
 end

data/lib/llm_gateway/adapters/normalized_stream_accumulator.rb

@@ -22,8 +22,8 @@ module LlmGateway
       #
       # Accepted event shapes:
       #
-      #   { type: :message_start, delta: { id: "...", model: "...", role: "assistant" }, usage_increment: { ... } }
-      #   { type: :message_delta, delta: { stop_reason: "stop" }, usage_increment: { ... } }
+      #   { type: :message_start, delta: { id: "...", model: "...", role: "assistant", timestamp: 1716650000000 } }
+      #   { type: :message_delta, delta: { stop_reason: "stop" }, usage: { output: 2 } }
       #   { type: :message_end }
       #
       #   { type: :text_start, delta: "hi" }
@@ -34,10 +34,14 @@ module LlmGateway
       #   { type: :reasoning_delta, delta: "...", signature: "" }
       #   { type: :reasoning_end, delta: "", signature: "" }
       #
-      #   { type: :tool_start, id: "...", name: "tool_name", delta: "" }
+      #   { type: :tool_start, id: "...", name: "tool_name", tool_type: "tool_use", delta: "" }
       #   { type: :tool_delta, delta: "{\"a\":" }
       #   { type: :tool_end, delta: "" }
       #
+      #   { type: :tool_result_start, tool_use_id: "...", name: "server_tool_result", delta: "..." }
+      #   { type: :tool_result_delta, delta: "..." }
+      #   { type: :tool_result_end, delta: "" }
+      #
       # Mappers do not provide `content_index`. The accumulator assigns the next
       # public content index when a block starts and reuses the active content
       # index for that block's deltas and end event.
@@ -50,7 +54,16 @@ module LlmGateway
       # The accumulator creates the public Assistant* event structs, updates its
       # accumulated message state, then yields the created event to the callback.
       attr_accessor :blocks, :message_hash, :usage_hash
-      attr_reader :active_block_type
+      attr_reader :active_block_type, :final_message
+
+      DEFAULT_USAGE = {
+        input: 0,
+        cache_write: 0,
+        cache_read: 0,
+        output: 0,
+        total: 0,
+        raw: {}
+      }.freeze
 
       BLOCK_EVENT_TRANSITIONS = {
         text_start: { block_type: :text, phase: :start },
@@ -59,33 +72,40 @@ module LlmGateway
         tool_start: { block_type: :tool, phase: :start },
         tool_delta: { block_type: :tool, phase: :delta },
         tool_end: { block_type: :tool, phase: :end },
+        tool_result_start: { block_type: :tool_result, phase: :start },
+        tool_result_delta: { block_type: :tool_result, phase: :delta },
+        tool_result_end: { block_type: :tool_result, phase: :end },
         reasoning_start: { block_type: :reasoning, phase: :start },
         reasoning_delta: { block_type: :reasoning, phase: :delta },
         reasoning_end: { block_type: :reasoning, phase: :end }
       }.freeze
 
-      def initialize
+      def initialize(provider: nil, api: nil)
+        @provider = provider
+        @api = api
         @message_hash = {}
-        @usage_hash = {
-          input_tokens: 0,
-          cache_creation_input_tokens: 0,
-          cache_read_input_tokens: 0,
-          output_tokens: 0,
-          reasoning_tokens: 0
-        }
+        @usage_hash = default_usage
         @blocks = []
         @next_content_index = 0
         @active_block_type = nil
         @active_content_index = nil
+        @timestamp = nil
       end
 
       def result
+        ensure_timestamp!
+
         message_hash.merge(
+          timestamp: @timestamp,
           usage: usage_hash,
           content: serialized_blocks
         )
       end
 
+      def final_result
+        result.merge(provider: @provider, api: @api)
+      end
+
       def active_tool?
         active_block_type == :tool
       end
@@ -93,14 +113,22 @@ module LlmGateway
       def push(event_patch, &block)
         raise ArgumentError, "Normalized stream event patch must be a Hash" unless event_patch.is_a?(Hash)
 
-        event_patch = symbolize_keys(event_patch)
+        event_patch = event_patch.symbolize_keys
         type = event_patch.fetch(:type).to_sym
         event_patch = prepare_event_patch(event_patch.merge(type:), type)
+        ensure_timestamp!
 
-        event = build_event(event_patch)
+        if type == :message_end
+          @final_message = AssistantMessage.new(final_result)
+          block.call(AssistantStreamMessageEndEvent.new(type:, message: final_message)) if block
+          return nil
+        end
+
+        event = build_event(event_patch, partial: empty_partial)
         accumulate(event)
         content_index = event.content_index if event.respond_to?(:content_index)
         commit_block_transition(type, content_index)
+        event = build_event(event_patch, partial: partial_message)
         block.call(event) if block
 
         nil
@@ -166,16 +194,21 @@ module LlmGateway
         end
       end
 
-      def build_event(event_patch)
-        event_patch = symbolize_keys(event_patch)
+      def build_event(event_patch, partial:)
+        event_patch = event_patch.symbolize_keys
         type = event_patch.fetch(:type).to_sym
 
         case type
-        when :message_start, :message_delta, :message_end
+        when :message_start, :message_delta
+          delta = (event_patch[:delta] || {}).symbolize_keys
+          raw_usage = event_patch[:usage] || delta.delete(:usage) || {}
+          usage = raw_usage.empty? ? {} : normalized_usage(raw_usage)
+
           AssistantStreamMessageEvent.new(
             type:,
-            delta: symbolize_keys(event_patch[:delta] || {}),
-            usage_increment: symbolize_keys(event_patch[:usage_increment] || {})
+            delta:,
+            usage:,
+            partial:
           )
         when :tool_start
           AssistantToolStartEvent.new(
@@ -183,20 +216,33 @@ module LlmGateway
             content_index: event_patch.fetch(:content_index),
             delta: string_value(event_patch[:delta]),
             id: event_patch[:id],
-            name: event_patch[:name]
+            name: event_patch[:name],
+            partial:,
+            tool_type: event_patch[:tool_type] || "tool_use"
+          )
+        when :tool_result_start
+          AssistantToolResultStartEvent.new(
+            type:,
+            content_index: event_patch.fetch(:content_index),
+            delta: string_value(event_patch[:delta]),
+            tool_use_id: event_patch[:tool_use_id],
+            name: event_patch[:name],
+            partial:
           )
         when :reasoning_start, :reasoning_delta, :reasoning_end
           AssistantStreamReasoningEvent.new(
             type:,
             content_index: event_patch.fetch(:content_index),
             delta: string_value(event_patch[:delta]),
-            signature: string_value(event_patch[:signature])
+            signature: string_value(event_patch[:signature]),
+            partial:
           )
-        when :text_start, :text_delta, :text_end, :tool_delta, :tool_end
+        when :text_start, :text_delta, :text_end, :tool_delta, :tool_end, :tool_result_delta, :tool_result_end
           AssistantStreamEvent.new(
             type:,
             content_index: event_patch.fetch(:content_index),
-            delta: string_value(event_patch[:delta])
+            delta: string_value(event_patch[:delta]),
+            partial:
           )
         else
           raise ArgumentError, "Unsupported normalized stream event type: #{type.inspect}"
@@ -204,6 +250,8 @@ module LlmGateway
       end
 
       def accumulate(event)
+        @timestamp = event.delta[:timestamp] if event.respond_to?(:delta) && event.delta.is_a?(Hash) && event.delta[:timestamp]
+
         case event.type
         when :text_start
           blocks[event.content_index] = {
@@ -215,18 +263,23 @@ module LlmGateway
           blocks[event.content_index][:text] += event.delta
         when :tool_start
           blocks[event.content_index] = {
-            type: "tool_use",
+            type: event.tool_type,
             id: event.id,
             name: event.name,
             input: event.delta.to_s
           }
         when :tool_delta, :tool_end
           blocks[event.content_index][:input] += event.delta
+        when :tool_result_start
+          blocks[event.content_index] = {
+            type: event.name,
+            tool_use_id: event.tool_use_id,
+            content: event.delta.to_s
+          }
+        when :tool_result_delta, :tool_result_end
+          blocks[event.content_index][:content] += event.delta
         when :message_start
           message_hash.merge!(event.delta)
-          usage_hash.each_key do |key|
-            usage_hash[key] += event.usage_increment.fetch(key, 0)
-          end
         when :reasoning_start
           blocks[event.content_index] = {
             type: "reasoning",
@@ -240,36 +293,76 @@ module LlmGateway
           blocks[event.content_index][:signature] += event.signature
         when :message_delta
           message_hash.merge!(event.delta)
-          usage_hash.each_key do |key|
-            usage_hash[key] += event.usage_increment.fetch(key, 0)
-          end
-        when :message_end
+          assign_usage(event.usage) unless event.usage.empty?
         end
       end
 
+      def empty_partial
+        PartialAssistantMessage.new(timestamp: @timestamp)
+      end
+
+      def partial_message
+        PartialAssistantMessage.new(partial_result)
+      end
+
+      def partial_result
+        ensure_timestamp!
+
+        message_hash.merge(
+          timestamp: @timestamp,
+          content: serialized_blocks
+        )
+      end
+
+      def assign_usage(usage)
+        @usage_hash = normalized_usage(usage)
+      end
+
+      def normalized_usage(usage)
+        usage = default_usage.merge(usage.to_h.symbolize_keys.slice(*DEFAULT_USAGE.keys))
+        usage[:total] = usage[:input] + usage[:cache_write] + usage[:cache_read] + usage[:output]
+        usage[:raw] ||= {}
+        usage
+      end
+
+      def default_usage
+        DEFAULT_USAGE.merge(raw: {})
+      end
+
       def serialized_blocks
-        blocks.map do |content_block|
-          next content_block unless content_block[:type] == "tool_use"
+        blocks.compact.map do |content_block|
+          if [ "tool_use", "server_tool_use" ].include?(content_block[:type])
+            next content_block.merge(input: parse_tool_input(content_block[:input]).deep_symbolize_keys)
+          end
+
+          if content_block[:type]&.end_with?("_tool_result")
+            next {
+              type: "server_tool_result",
+              tool_use_id: content_block[:tool_use_id],
+              name: content_block[:type],
+              content: parse_tool_input(content_block[:content]).deep_symbolize_keys
+            }
+          end
 
-          content_block.merge(input: LlmGateway::Utils.deep_symbolize_keys(parse_tool_input(content_block[:input])))
+          content_block
         end
       end
 
       def parse_tool_input(input)
-        return {} if input.nil? || input.empty?
+        return {} if input.blank?
 
         JSON.parse(input)
       rescue JSON::ParserError
         {}
       end
 
-      def symbolize_keys(hash)
-        hash.to_h.transform_keys { |key| key.respond_to?(:to_sym) ? key.to_sym : key }
-      end
-
       def string_value(value)
         value.nil? ? "" : value.to_s
       end
+
+      def ensure_timestamp!
+        @timestamp ||= (Time.now.to_f * 1000).to_i
+      end
     end
   end
 end

data/lib/llm_gateway/adapters/openai/chat_completions/option_mapper.rb

@@ -66,7 +66,7 @@ module LlmGateway
           module_function
 
           def map(options)
-            mapped_options = options.reject { |key, _| MANAGED_OPTIONS.include?(key) }
+            mapped_options = options.except(*MANAGED_OPTIONS)
             mapped_options[:max_completion_tokens] = options[:max_completion_tokens] || DEFAULT_MAX_COMPLETION_TOKENS
 
             cache_key = options[:cache_key]