llm_gateway 0.5.0 → 0.6.0

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" }
@@ -50,7 +50,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 },
@@ -64,28 +73,32 @@ module LlmGateway
         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
@@ -96,11 +109,19 @@ module LlmGateway
         event_patch = symbolize_keys(event_patch)
         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 +187,21 @@ module LlmGateway
         end
       end
 
-      def build_event(event_patch)
+      def build_event(event_patch, partial:)
         event_patch = symbolize_keys(event_patch)
         type = event_patch.fetch(:type).to_sym
 
         case type
-        when :message_start, :message_delta, :message_end
+        when :message_start, :message_delta
+          delta = symbolize_keys(event_patch[:delta] || {})
+          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 +209,23 @@ 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:
           )
         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
           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 +233,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] = {
@@ -224,9 +255,6 @@ module LlmGateway
           blocks[event.content_index][:input] += 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,13 +268,42 @@ 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(symbolize_keys(usage).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"
@@ -270,6 +327,10 @@ module LlmGateway
       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/stream_mapper.rb

@@ -92,9 +92,9 @@ module LlmGateway
                 delta: {
                   id: data[:id],
                   model: data[:model],
-                  role: delta[:role] || "assistant"
-                }.compact,
-                usage_increment: {}
+                  role: delta[:role] || "assistant",
+                  timestamp: timestamp_milliseconds(data[:created])
+                }.compact
               }
             ]
           end
@@ -198,34 +198,58 @@ module LlmGateway
               *close_active_block_patches(active_block_type:),
               {
                 type: :message_delta,
-                delta: { stop_reason: normalize_stop_reason(finish_reason) },
-                usage_increment: {}
+                delta: { stop_reason: normalize_stop_reason(finish_reason) }
               }
             ]
           end
 
           def final_usage_patches(data)
+            patch = {
+              type: :message_delta,
+              delta: {}
+            }
+            patch[:usage] = usage(data) if data.key?(:usage)
+
             [
-              {
-                type: accumulator.message_hash.empty? ? :message_start : :message_delta,
-                delta: {},
-                usage_increment: usage_increment(data)
-              }
+              patch,
+              { type: :message_end }
             ]
           end
 
-          def usage_increment(data)
+          def usage(data)
             usage = data[:usage] || {}
+            cache_read = token_count(
+              usage.dig(:prompt_tokens_details, :cached_tokens),
+              usage[:prompt_cache_hit_tokens]
+            )
+            cache_write = token_count(
+              usage.dig(:prompt_tokens_details, :cache_write_tokens),
+              usage[:cache_write_tokens]
+            )
+            prompt_tokens = token_count(usage[:prompt_tokens])
+            input = [ prompt_tokens - cache_read - cache_write, 0 ].max
+            output = token_count(usage[:completion_tokens])
 
             {
-              input_tokens: usage[:prompt_tokens] || 0,
-              cache_creation_input_tokens: 0,
-              cache_read_input_tokens: usage.dig(:prompt_tokens_details, :cached_tokens) || 0,
-              output_tokens: usage[:completion_tokens] || 0,
-              reasoning_tokens: usage.dig(:completion_tokens_details, :reasoning_tokens) || 0
+              input:,
+              cache_write:,
+              cache_read:,
+              output:,
+              total: input + cache_write + cache_read + output,
+              raw: usage
             }
           end
 
+          def token_count(*values)
+            values.compact.first.to_i
+          end
+
+          def timestamp_milliseconds(unix_seconds)
+            return nil if unix_seconds.nil?
+
+            (unix_seconds.to_f * 1000).to_i
+          end
+
           def normalize_stop_reason(finish_reason)
             case finish_reason
             when "tool_calls"

data/lib/llm_gateway/adapters/openai/responses/stream_mapper.rb

@@ -55,9 +55,9 @@ module LlmGateway
                 delta: {
                   id: response[:id],
                   model: response[:model],
-                  role: "assistant"
-                }.compact,
-                usage_increment: {}
+                  role: "assistant",
+                  timestamp: timestamp_milliseconds(response[:created_at])
+                }.compact
               }
             ]
           end
@@ -72,8 +72,7 @@ module LlmGateway
               [
                 {
                   type: :message_start,
-                  delta: { role: item[:role] || "assistant" },
-                  usage_increment: {}
+                  delta: { role: item[:role] || "assistant" }
                 }
               ]
             when "function_call"
@@ -106,33 +105,55 @@ module LlmGateway
 
           def response_completed_patches(response)
             response ||= {}
+            patch = {
+              type: :message_delta,
+              delta: {
+                id: response[:id],
+                model: response[:model],
+                role: "assistant",
+                timestamp: timestamp_milliseconds(response[:created_at]),
+                stop_reason: stop_reason_for(response)
+              }.compact
+            }
+            patch[:usage] = usage(response) if response.key?(:usage)
 
             [
-              {
-                type: accumulator.message_hash.empty? ? :message_start : :message_delta,
-                delta: {
-                  id: response[:id],
-                  model: response[:model],
-                  role: "assistant",
-                  stop_reason: stop_reason_for(response)
-                }.compact,
-                usage_increment: usage_increment(response)
-              }
+              patch,
+              { type: :message_end }
             ]
           end
 
-          def usage_increment(response)
+          def usage(response)
             usage = response[:usage] || {}
+            cache_read = token_count(usage.dig(:input_tokens_details, :cached_tokens))
+            cache_write = token_count(
+              usage.dig(:input_tokens_details, :cache_write_tokens),
+              usage[:cache_write_tokens]
+            )
+            input_tokens = token_count(usage[:input_tokens])
+            input = [ input_tokens - cache_read - cache_write, 0 ].max
+            output = token_count(usage[:output_tokens])
 
             {
-              input_tokens: usage[:input_tokens] || 0,
-              cache_creation_input_tokens: 0,
-              cache_read_input_tokens: usage.dig(:input_tokens_details, :cached_tokens) || 0,
-              output_tokens: usage[:output_tokens] || 0,
-              reasoning_tokens: usage.dig(:output_tokens_details, :reasoning_tokens) || 0
+              input:,
+              cache_write:,
+              cache_read:,
+              output:,
+              total: input + cache_write + cache_read + output,
+              raw: usage
             }
           end
 
+          def token_count(*values)
+            values.compact.first.to_i
+          end
+
+          def timestamp_milliseconds(unix_seconds)
+            return nil if unix_seconds.nil?
+
+            (unix_seconds.to_f * 1000).to_i
+          end
+
           def stop_reason_for(response)
             output = response[:output] || []
             last_item = output.last || {}

data/lib/llm_gateway/adapters/stream_mapper.rb

@@ -5,14 +5,21 @@ require_relative "normalized_stream_accumulator"
 module LlmGateway
   module Adapters
     class StreamMapper
+      def initialize(provider:, api:)
+        @provider = provider
+        @api = api
+      end
+
       def result
-        accumulator.result
+        accumulator.final_message
       end
 
       private
 
+      attr_reader :provider, :api
+
       def accumulator
-        @accumulator ||= LlmGateway::Adapters::NormalizedStreamAccumulator.new
+        @accumulator ||= LlmGateway::Adapters::NormalizedStreamAccumulator.new(provider:, api:)
       end
 
       def push_patches(patches, &block)

data/lib/llm_gateway/adapters/structs.rb

@@ -9,35 +9,6 @@ class BaseStruct < Dry::Struct
   transform_keys(&:to_sym)
 end
 
-class AssistantStreamEvent < BaseStruct
-  EventType = Types::Coercible::Symbol.enum(:text_start, :text_delta, :text_end, :tool_start, :tool_delta, :tool_end, :reasoning_start, :reasoning_delta, :reasoning_end)
-
-  attribute :type, EventType
-  attribute :delta, Types::Coercible::String.default { "" }
-  attribute :content_index, Types::Integer
-end
-
-
-class AssistantToolStartEvent < AssistantStreamEvent
-  attribute :id, Types::String
-  attribute :name, Types::String
-  attribute :content_index, Types::Integer
-end
-
-
-class AssistantStreamReasoningEvent < AssistantStreamEvent
-  attribute :signature, Types::Coercible::String.default { "" }
-  attribute :content_index, Types::Integer
-end
-
-class AssistantStreamMessageEvent < BaseStruct
-  EventType = Types::Coercible::Symbol.enum(:message_start, :message_delta, :message_end)
-
-  attribute :type, EventType
-  attribute :delta, Types::Coercible::Hash.default { {} }
-  attribute :usage_increment, Types::Coercible::Hash.default { {} }
-end
-
 class TextContent < BaseStruct
   attribute :type, Types::String.enum("text")
   attribute :text, Types::String
@@ -87,12 +58,101 @@ class ToolResult < BaseStruct
   attribute :content, Types::String
 end
 
-class AssistantMessage < BaseStruct
+class PartialAssistantMessage < BaseStruct
   ContentBlock =
     Types.Instance(TextContent) |
     Types.Instance(ReasoningContent) |
     Types.Instance(ToolCall)
 
+  attribute? :id, Types::String.optional
+  attribute? :model, Types::String.optional
+  attribute? :role, Types::String.enum("assistant").optional
+  attribute :timestamp, Types::Integer
+  attribute? :stop_reason, Types::String.enum("stop", "length", "tool_use", "toolUse", "error", "aborted").optional
+  attribute? :content, Types::Array.of(ContentBlock).optional
+
+  def self.new(attributes = {})
+    attrs = attributes.to_h.transform_keys(&:to_sym)
+    attrs[:content] = Array(attrs[:content]).map { |block| build_content_block(block) } if attrs.key?(:content)
+    super(attrs)
+  end
+
+  def self.build_content_block(block)
+    return block if block.is_a?(TextContent) || block.is_a?(ReasoningContent) || block.is_a?(ToolCall)
+
+    case block[:type] || block["type"]
+    when "text"
+      TextContent.new(block)
+    when "reasoning"
+      ReasoningContent.new(block)
+    when "thinking"
+      ReasoningContent.new(
+        type: "reasoning",
+        reasoning: block[:thinking] || block["thinking"] || block[:reasoning] || block["reasoning"],
+        signature: block[:signature] || block["signature"]
+      )
+    when "tool_use"
+      ToolCall.new(block)
+    else
+      raise ArgumentError, "Unsupported content block type: #{block[:type] || block['type']}"
+    end
+  end
+
+  private_class_method :build_content_block
+end
+
+class AssistantStreamEvent < BaseStruct
+  EventType = Types::Coercible::Symbol.enum(:text_start, :text_delta, :text_end, :tool_start, :tool_delta, :tool_end, :reasoning_start, :reasoning_delta, :reasoning_end)
+
+  attribute :type, EventType
+  attribute :delta, Types::Coercible::String.default { "" }
+  attribute :content_index, Types::Integer
+  attribute :partial, Types.Instance(PartialAssistantMessage)
+
+  def content
+    case type
+    when :text_end
+      finalized_content_block&.text
+    when :reasoning_end
+      finalized_content_block&.reasoning
+    when :tool_end
+      finalized_content_block
+    end
+  end
+
+  def text
+    content if type == :text_end
+  end
+
+  def reasoning
+    content if type == :reasoning_end
+  end
+
+  def tool_call
+    finalized_content_block if type == :tool_end
+  end
+
+  alias tool tool_call
+
+  private
+
+  def finalized_content_block
+    partial.content&.[](content_index)
+  end
+end
+
+class AssistantToolStartEvent < AssistantStreamEvent
+  attribute :id, Types::String
+  attribute :name, Types::String
+  attribute :content_index, Types::Integer
+end
+
+class AssistantStreamReasoningEvent < AssistantStreamEvent
+  attribute :signature, Types::Coercible::String.default { "" }
+  attribute :content_index, Types::Integer
+end
+
+class AssistantMessage < PartialAssistantMessage
   attribute :id, Types::String
   attribute :model, Types::String
   attribute :usage, Types::Hash
@@ -103,12 +163,6 @@ class AssistantMessage < BaseStruct
   attribute? :error_message, Types::String.optional
   attribute :content, Types::Array.of(ContentBlock)
 
-  def self.new(attributes)
-    attrs = attributes.to_h.transform_keys(&:to_sym)
-    attrs[:content] = Array(attrs[:content]).map { |block| build_content_block(block) }
-    super(attrs)
-  end
-
   def to_h
     result = {
       id: id,
@@ -120,26 +174,22 @@ class AssistantMessage < BaseStruct
       api: api,
       content: content.map(&:to_h)
     }
+    result[:timestamp] = timestamp unless timestamp.nil?
     result[:error_message] = error_message unless error_message.nil?
     result
   end
+end
 
-  def self.build_content_block(block)
-    return block if block.is_a?(TextContent) || block.is_a?(ReasoningContent) || block.is_a?(ToolCall)
+class AssistantStreamMessageEvent < BaseStruct
+  EventType = Types::Coercible::Symbol.enum(:message_start, :message_delta)
 
-    case block[:type] || block["type"]
-    when "text"
-      TextContent.new(block)
-    when "reasoning"
-      ReasoningContent.new(block)
-    when "thinking"
-      ReasoningContent.new(type: "reasoning", reasoning: block[:thinking] || block["thinking"] || block[:reasoning] || block["reasoning"], signature: block[:signature] || block["signature"])
-    when "tool_use"
-      ToolCall.new(block)
-    else
-      raise ArgumentError, "Unsupported content block type: #{block[:type] || block['type']}"
-    end
-  end
+  attribute :type, EventType
+  attribute :delta, Types::Coercible::Hash.default { {} }
+  attribute :usage, Types::Coercible::Hash.default { {} }
+  attribute :partial, Types.Instance(PartialAssistantMessage)
+end
 
-  private_class_method :build_content_block
+class AssistantStreamMessageEndEvent < BaseStruct
+  attribute :type, Types::Coercible::Symbol.enum(:message_end)
+  attribute :message, Types.Instance(AssistantMessage)
 end

data/lib/llm_gateway/base_client.rb

@@ -6,11 +6,9 @@ require "json"
 
 module LlmGateway
   class BaseClient
-    attr_accessor
-    attr_reader :api_key, :model_key, :base_endpoint
+    attr_reader :api_key, :base_endpoint
 
-    def initialize(model_key:, api_key:)
-      @model_key = model_key
+    def initialize(api_key:)
       @api_key = api_key
     end
 

data/lib/llm_gateway/clients/anthropic.rb

@@ -9,10 +9,11 @@ module LlmGateway
   module Clients
     class Anthropic < BaseClient
       CLAUDE_CODE_VERSION = "2.1.2"
+      DEFAULT_MODEL = "claude-3-7-sonnet-20250219"
 
-      def initialize(model_key: "claude-3-7-sonnet-20250219", api_key: ENV["ANTHROPIC_API_KEY"])
+      def initialize(api_key: ENV["ANTHROPIC_API_KEY"])
         @base_endpoint = "https://api.anthropic.com/v1"
-        super(model_key: model_key, api_key: api_key)
+        super(api_key: api_key)
       end
 
       def chat(messages, **kwargs)
@@ -44,11 +45,11 @@ module LlmGateway
 
       private
 
-      def build_body(messages, tools: nil, system: [], cache_retention: nil, **options)
+      def build_body(messages, tools: nil, system: [], cache_retention: nil, model: DEFAULT_MODEL, **options)
         cache_control = anthropic_cache_control_for(cache_retention)
 
         body = {
-          model: model_key,
+          model: model,
           messages: messages
         }
 

data/lib/llm_gateway/clients/groq.rb

@@ -5,14 +5,16 @@ require_relative "../base_client"
 module LlmGateway
   module Clients
     class Groq < BaseClient
-      def initialize(model_key: "openai/gpt-oss-120b", api_key: ENV["GROQ_API_KEY"])
+      DEFAULT_MODEL = "openai/gpt-oss-120b"
+
+      def initialize(api_key: ENV["GROQ_API_KEY"])
         @base_endpoint = "https://api.groq.com/openai/v1"
-        super(model_key: model_key, api_key: api_key)
+        super(api_key: api_key)
       end
 
-      def chat(messages, tools: nil, system: [], **options)
+      def chat(messages, tools: nil, system: [], model: DEFAULT_MODEL, **options)
         body = {
-          model: model_key,
+          model: model,
           messages: system + messages,
           tools: tools
         }
@@ -21,9 +23,9 @@ module LlmGateway
         post("chat/completions", body)
       end
 
-      def stream(messages, tools: nil, system: [], **options, &block)
+      def stream(messages, tools: nil, system: [], model: DEFAULT_MODEL, **options, &block)
         body = {
-          model: model_key,
+          model: model,
           messages: system + messages,
           tools: tools,
           stream_options: { include_usage: true }