llm_cost_tracker 0.8.0 → 0.10.0

data/lib/llm_cost_tracker/parsers/openai_service_charges.rb

@@ -1,6 +1,7 @@
 # frozen_string_literal: true
 
 require_relative "../billing/line_item"
+require_relative "../providers/openai/model_families"
 
 module LlmCostTracker
   module Parsers
@@ -8,15 +9,52 @@ module LlmCostTracker
       RESPONSE_OUTPUT_COMPONENTS = {
         "web_search_call" => :web_search_request,
         "file_search_call" => :file_search_call,
-        "code_interpreter_call" => :container_session
+        "code_interpreter_call" => :container_session,
+        "mcp_call" => :mcp_call
       }.freeze
 
       module_function
 
-      def line_items_from_output(output_items)
+      def line_items_from_output(output_items, request: nil, model: nil)
         deduped = {}
         Array(output_items).each { |item| store_output_item(deduped, item) }
-        deduped.values.filter_map { |item| build_line_item(item) }
+        deduped.values
+               .select { |item| billable?(item) }
+               .filter_map { |item| build_line_item(item, request: request, model: model) }
+      end
+
+      def service_line_items_for(response, request: nil, model: nil)
+        output_items = Array(response["output"])
+        output_items += chat_completions_web_search_items(response, model: model) if output_items.empty?
+        line_items_from_output(output_items, request: request, model: model)
+      end
+
+      CHAT_COMPLETIONS_ANNOTATION_PROVIDER_FIELD = "choices.message.annotations.url_citation"
+      CHAT_COMPLETIONS_SEARCH_MODEL_PROVIDER_FIELD = "request.model"
+
+      def chat_completions_web_search_items(response, model: nil)
+        return [] unless response["choices"]
+
+        provider_field = chat_completions_search_provider_field(response["choices"], model)
+        return [] unless provider_field
+
+        [{ "type" => "web_search_call", "id" => response["id"], "action" => { "type" => "search" },
+           "provider_field" => provider_field }]
+      end
+
+      def chat_completions_search_provider_field(choices, model)
+        return CHAT_COMPLETIONS_ANNOTATION_PROVIDER_FIELD if chat_completions_used_web_search?(choices)
+        return CHAT_COMPLETIONS_SEARCH_MODEL_PROVIDER_FIELD if chat_completions_search_model?(model)
+
+        nil
+      end
+
+      def chat_completions_used_web_search?(choices)
+        Array(choices).any? do |choice|
+          Array(choice.dig("message", "annotations")).any? do |annotation|
+            annotation.is_a?(Hash) && annotation["type"] == "url_citation"
+          end
+        end
       end
 
       def billable?(item)
@@ -31,7 +69,7 @@ module LlmCostTracker
       end
 
       def store_output_item(output_items, item)
-        return unless billable?(item)
+        return unless item.is_a?(Hash) && RESPONSE_OUTPUT_COMPONENTS.key?(item["type"])
 
         component = RESPONSE_OUTPUT_COMPONENTS[item["type"]]
         key = if component == :container_session && item["container_id"]
@@ -42,8 +80,10 @@ module LlmCostTracker
         output_items[key] = item
       end
 
-      def build_line_item(item)
-        component_key = RESPONSE_OUTPUT_COMPONENTS[item["type"]]
+      def build_line_item(item, request: nil, model: nil)
+        return nil unless item.is_a?(Hash)
+
+        component_key = component_key_for(item, request: request, model: model)
         return nil unless component_key
 
         provider_item_id = if component_key == :container_session
@@ -56,12 +96,44 @@ module LlmCostTracker
           quantity: 1,
           cost_status: Billing::CostStatus::UNKNOWN,
           pricing_basis: :provider_usage,
-          provider_field: "response.output.#{item['type']}",
+          provider_field: item["provider_field"] || "response.output.#{item['type']}",
           provider_item_id: provider_item_id,
           details: line_item_details(item)
         )
       end
 
+      def component_key_for(item, request:, model:)
+        component = RESPONSE_OUTPUT_COMPONENTS[item["type"]]
+        return component unless component == :web_search_request
+        return component unless web_search_preview_used?(request) || chat_completions_search_model?(model)
+
+        reasoning_model?(model) ? :web_search_preview_request_reasoning : :web_search_preview_request_non_reasoning
+      end
+
+      def web_search_preview_used?(request)
+        tools = request && (request[:tools] || request["tools"])
+        return false unless tools.respond_to?(:each)
+
+        tools.any? do |tool|
+          type = tool.is_a?(Hash) ? (tool[:type] || tool["type"]) : tool
+          type.to_s.include?("web_search_preview")
+        end
+      end
+
+      def chat_completions_search_model?(model)
+        return false unless model
+
+        name = model.to_s.split("/", 2).last
+        LlmCostTracker::Providers::Openai::ModelFamilies.chat_completions_search?(name)
+      end
+
+      def reasoning_model?(model)
+        return false unless model
+
+        name = model.to_s.split("/", 2).last
+        LlmCostTracker::Providers::Openai::ModelFamilies.reasoning?(name)
+      end
+
       def line_item_details(item)
         {
           "status" => item["status"],
@@ -70,17 +142,13 @@ module LlmCostTracker
         }.compact
       end
 
-      def openai_service_line_items(response)
-        line_items_from_output(response["output"])
-      end
-
-      def openai_stream_service_line_items(events)
+      def openai_stream_service_line_items(events, request: nil, model: nil)
         output_items = []
         each_event_data(events) do |data|
           output_items.concat(Array(data.dig("response", "output")))
           output_items << data["item"] if data["item"]
         end
-        line_items_from_output(output_items)
+        line_items_from_output(output_items, request: request, model: model)
       end
     end
   end

data/lib/llm_cost_tracker/parsers/openai_usage.rb

@@ -1,17 +1,29 @@
 # frozen_string_literal: true
 
 require_relative "openai_service_charges"
+require_relative "../providers/openai/hosts"
+require_relative "../providers/openai/model_families"
 
 module LlmCostTracker
   module Parsers
     module OpenaiUsage
       include OpenaiServiceCharges
 
-      OPENAI_DATA_RESIDENCY_HOST_PATTERN = /\A[a-z]{2,3}\.api\.openai\.com\z/
+      class << self
+        def combined_pricing_mode(host:, model:, service_tier:)
+          modes = [Pricing.normalize_mode(service_tier)]
+          modes << "data_residency" if regional_processing?(host: host, model: model)
+          modes = modes.compact.uniq
+          modes.empty? ? nil : modes.join("_")
+        end
 
-      private
+        def regional_processing?(host:, model:)
+          LlmCostTracker::Providers::Openai::Hosts.data_residency?(host) &&
+            LlmCostTracker::Providers::Openai::ModelFamilies.data_residency?(model)
+        end
+      end
 
-      def parse_openai_usage(request_url:, request_body:, response_status:, response_body:)
+      def parse(request_url:, request_body:, response_status:, response_body:, **)
         return nil unless response_status == 200
 
         response = safe_json_parse(response_body)
@@ -23,7 +35,7 @@ module LlmCostTracker
 
         model = response["model"] || request["model"]
 
-        UsageCapture.build(
+        Event.build(
           provider: provider_for(request_url),
           provider_response_id: response["id"],
           pricing_mode: pricing_mode(
@@ -32,101 +44,141 @@ module LlmCostTracker
             service_tier: response["service_tier"] || request["service_tier"]
           ),
           model: model,
-          token_usage: token_usage(usage: usage, cache_read: cache_read),
+          token_usage: token_usage(usage: usage, cache_read: cache_read, model: model),
           usage_source: :response,
-          service_line_items: openai_service_line_items(response)
+          service_line_items: service_line_items_for(response, request: request, model: response["model"])
         )
       end
 
-      def parse_openai_stream_usage(response_status:, request_url: nil, request_body: nil, events: [])
+      def parse_stream(response_status:, request_url: nil, request_body: nil, events: [], **)
         return nil unless response_status == 200
 
         request = safe_json_parse(request_body)
-        model = find_event_value(events) { |data| data["model"] || data.dig("response", "model") } || request["model"]
         usage = detect_stream_usage(events)
-        response_id = find_event_value(events) { |data| data["id"] || data.dig("response", "id") }
-        pricing_mode = pricing_mode(
-          request_url: request_url,
+        context = stream_capture_context(events: events, request: request, request_url: request_url)
+
+        return build_known_stream_usage(usage: usage, **context) if usage
+
+        warn_missing_stream_usage(request_url: request_url, request: request)
+        build_unknown_stream_usage(**context)
+      end
+
+      def auto_enable_stream_usage?(request_url)
+        openai_chat_completions_url?(request_url)
+      end
+
+      private
+
+      def stream_capture_context(events:, request:, request_url:)
+        model = find_event_value(events) do |data|
+          data["model"] || data.dig("response", "model") || data.dig("chunk", "model")
+        end || request["model"]
+        {
+          provider: provider_for(request_url),
           model: model,
-          service_tier: stream_pricing_mode(events) || request["service_tier"]
-        )
-        service_line_items = openai_stream_service_line_items(events)
-
-        if usage
-          cache_read = cache_read_input_tokens(usage)
-          UsageCapture.build(
-            provider: provider_for(request_url),
-            provider_response_id: response_id,
-            pricing_mode: pricing_mode,
-            model: model,
-            token_usage: token_usage(usage: usage, cache_read: cache_read),
-            stream: true,
-            usage_source: :stream_final,
-            service_line_items: service_line_items
-          )
-        else
-          build_unknown_stream_usage(
-            provider: provider_for(request_url),
+          provider_response_id: find_event_value(events) do |data|
+            data["id"] || data.dig("response", "id") || data.dig("chunk", "id")
+          end,
+          pricing_mode: pricing_mode(
+            request_url: request_url,
             model: model,
-            provider_response_id: response_id,
-            pricing_mode: pricing_mode,
-            service_line_items: service_line_items
-          )
-        end
+            service_tier: stream_pricing_mode(events) || request["service_tier"]
+          ),
+          service_line_items: openai_stream_service_line_items(events, request: request, model: model)
+        }
+      end
+
+      def build_known_stream_usage(usage:, provider:, model:, provider_response_id:, pricing_mode:, service_line_items:)
+        cache_read = cache_read_input_tokens(usage)
+        Event.build(
+          provider: provider,
+          provider_response_id: provider_response_id,
+          pricing_mode: pricing_mode,
+          model: model,
+          token_usage: token_usage(usage: usage, cache_read: cache_read, model: model),
+          stream: true,
+          usage_source: :stream_final,
+          service_line_items: service_line_items
+        )
+      end
+
+      def warn_missing_stream_usage(request_url:, request:)
+        return unless request.is_a?(Hash) && request["stream"]
+        return unless openai_chat_completions_url?(request_url)
+        return if request.dig("stream_options", "include_usage")
+
+        Logging.warn(
+          "OpenAI-compatible chat-completions stream finished without a final usage chunk. " \
+          "Set `stream_options: { include_usage: true }` in your request body so the gem can " \
+          "record token counts. This call was stored with usage_source=unknown."
+        )
+      end
+
+      def openai_chat_completions_url?(request_url)
+        uri = parsed_uri(request_url)
+        uri && uri.path.to_s.end_with?("/chat/completions")
       end
 
       def detect_stream_usage(events)
         find_event_value(events, reverse: true) do |data|
-          usage = data["usage"] || data.dig("response", "usage")
+          usage = data["usage"] || data.dig("response", "usage") || data.dig("chunk", "usage")
           usage if usage.is_a?(Hash)
         end
       end
 
       def stream_pricing_mode(events)
         find_event_value(events, reverse: true) do |data|
-          data["service_tier"] || data.dig("response", "service_tier")
+          data["service_tier"] || data.dig("response", "service_tier") || data.dig("chunk", "service_tier")
         end
       end
 
       def pricing_mode(request_url:, model:, service_tier:)
-        modes = [Pricing.normalize_mode(service_tier)]
-        modes << "data_residency" if openai_regional_processing?(request_url: request_url, model: model)
-        modes = modes.compact.uniq
-        modes.empty? ? nil : modes.join("_")
+        OpenaiUsage.combined_pricing_mode(host: parsed_uri(request_url)&.host, model: model, service_tier: service_tier)
       end
 
-      def openai_regional_processing?(request_url:, model:)
-        uri = parsed_uri(request_url)
-        return false unless uri&.host.to_s.downcase.match?(OPENAI_DATA_RESIDENCY_HOST_PATTERN)
-
-        openai_data_residency_model?(model)
-      end
-
-      def openai_data_residency_model?(model)
-        model.to_s.match?(/\Agpt-5\.(?:4|5)(?:-(?:mini|nano|pro))?(?:-\d{4}-\d{2}-\d{2})?\z/)
-      end
-
-      def token_usage(usage:, cache_read:)
+      def token_usage(usage:, cache_read:, model: nil)
         audio_input = audio_input_tokens(usage)
         audio_output = audio_output_tokens(usage)
+        image_input = image_input_tokens(usage)
+        image_output_details = image_output_tokens(usage)
+        text_output_details = text_output_tokens(usage)
+        raw_output = (usage["completion_tokens"] || usage["output_tokens"]).to_i
+        image_output, regular_output_remainder = split_stream_image_output(
+          raw_output: raw_output, image_output_details: image_output_details,
+          text_output_details: text_output_details, audio_output: audio_output,
+          default_to_image: LlmCostTracker::Providers::Openai::ModelFamilies.image_output?(model)
+        )
 
         TokenUsage.build(
-          input_tokens: regular_input_tokens(usage: usage, cache_read: cache_read, audio_input: audio_input),
-          output_tokens: regular_output_tokens(usage: usage, audio_output: audio_output),
+          input_tokens: regular_input_tokens(
+            usage: usage, cache_read: cache_read, audio_input: audio_input, image_input: image_input
+          ),
+          output_tokens: regular_output_remainder,
           total_tokens: usage["total_tokens"],
           cache_read_input_tokens: cache_read,
           audio_input_tokens: audio_input,
           audio_output_tokens: audio_output,
+          image_input_tokens: image_input,
+          image_output_tokens: image_output,
           hidden_output_tokens: hidden_output_tokens(usage)
         )
       end
 
-      def regular_input_tokens(usage:, cache_read:, audio_input:)
-        [(usage["prompt_tokens"] || usage["input_tokens"]).to_i - cache_read - audio_input, 0].max
+      def split_stream_image_output(raw_output:, image_output_details:, text_output_details:, audio_output:,
+                                    default_to_image: false)
+        if image_output_details.zero? && text_output_details.zero?
+          remainder = [raw_output - audio_output, 0].max
+          return default_to_image ? [remainder, 0] : [0, remainder]
+        end
+
+        text_output = text_output_details
+        text_output = [raw_output - image_output_details - audio_output, 0].max if text_output.zero?
+        [image_output_details, text_output]
       end
 
-      def regular_output_tokens(usage:, audio_output:)
-        [(usage["completion_tokens"] || usage["output_tokens"]).to_i - audio_output, 0].max
+      def regular_input_tokens(usage:, cache_read:, audio_input:, image_input:)
+        raw = (usage["prompt_tokens"] || usage["input_tokens"]).to_i
+        [raw - cache_read - audio_input - image_input, 0].max
       end
 
       def cache_read_input_tokens(usage)
@@ -149,6 +201,21 @@ module LlmCostTracker
         details["audio_tokens"].to_i
       end
 
+      def image_input_tokens(usage)
+        details = input_token_details(usage)
+        details["image_tokens"].to_i
+      end
+
+      def image_output_tokens(usage)
+        details = output_token_details(usage)
+        details["image_tokens"].to_i
+      end
+
+      def text_output_tokens(usage)
+        details = output_token_details(usage)
+        details["text_tokens"].to_i
+      end
+
       def input_token_details(usage)
         usage["prompt_tokens_details"] || usage["input_tokens_details"] || usage["input_token_details"] || {}
       end

data/lib/llm_cost_tracker/parsers.rb

@@ -2,19 +2,46 @@
 
 module LlmCostTracker
   module Parsers
-    BUILT_INS = [Openai.new, OpenaiCompatible.new, Anthropic.new, Gemini.new].freeze
+    autoload :Base,                 "llm_cost_tracker/parsers/base"
+    autoload :OpenaiUsage,          "llm_cost_tracker/parsers/openai_usage"
+    autoload :OpenaiServiceCharges, "llm_cost_tracker/parsers/openai_service_charges"
+    autoload :SSE,                  "llm_cost_tracker/parsers/sse"
+    autoload :Openai,               "llm_cost_tracker/parsers/openai"
+    autoload :Azure,                "llm_cost_tracker/parsers/azure"
+    autoload :OpenaiCompatible,     "llm_cost_tracker/parsers/openai_compatible"
+    autoload :Anthropic,            "llm_cost_tracker/parsers/anthropic"
+    autoload :Gemini,               "llm_cost_tracker/parsers/gemini"
+
+    MUTEX = Mutex.new
+    PARSER_CONSTANTS = %i[Openai Azure OpenaiCompatible Anthropic Gemini].freeze
 
     module_function
 
     def find_for(url)
-      BUILT_INS.find { |parser| parser.match?(url) }
+      PARSER_CONSTANTS.each do |name|
+        klass = const_get(name)
+        return instance_for(klass) if klass.match?(url)
+      end
+      nil
     end
 
     def find_for_provider(provider)
       provider_name = provider.to_s.downcase
-      BUILT_INS.find do |parser|
-        parser.provider_names.include?(provider_name)
+      PARSER_CONSTANTS.each do |name|
+        klass = const_get(name)
+        return instance_for(klass) if klass.provider_names.include?(provider_name)
+      end
+      nil
+    end
+
+    def instance_for(klass)
+      cached = (@instances ||= {})[klass]
+      return cached if cached
+
+      MUTEX.synchronize do
+        @instances[klass] ||= klass.new
       end
     end
+    private_class_method :instance_for
   end
 end