nitro_intelligence 3.1.0 → 4.1.0

data/lib/nitro_intelligence/client/base.rb

@@ -1,35 +1,111 @@
 require "nitro_intelligence/client/handlers/audio_transcription_handler"
 require "nitro_intelligence/client/handlers/chat_handler"
+require "nitro_intelligence/client/handlers/decision_handler"
 require "nitro_intelligence/client/handlers/image_handler"
 require "nitro_intelligence/client/handlers/text_to_speech_handler"
 
 module NitroIntelligence
   module Client
+    # A client {Client.new} builds. Every method takes `parameters`, the request's openai-ruby
+    # parameters plus this gem's own (`prompt_name`, `prompt_variables`, `trace_seed`, `metadata`,
+    # `user_id` and the rest the README describes), and defaults the model to the configured
+    # default for its modality.
+    #
+    # Each request works on a copy of the caller's `parameters`, which the handlers fill in with
+    # defaults and gateway headers, so a hash passed to several requests carries nothing from one
+    # request into the next.
     class Base
+      # @api private
       attr_reader :client
 
+      # @api private
       def initialize(client:)
         @client = client
       end
 
+      # Creates a chat completion.
+      #
+      # @param message [String] the user's message, used when `parameters` carries no `messages`
+      # @param parameters [Hash] openai-ruby's chat completion parameters, and this gem's own
+      # @return [OpenAI::Models::Chat::ChatCompletion]
+      # @raise [ArgumentError] when the model is not one `model_config` lists
+      # @raise [Handlers::Observed::ChatHandler::ObservedChatPromptError] on an observed client, when
+      #   the request carries no user message for the model to answer
+      # @example
+      #   client.chat(message: "Why is the sky blue?").choices.first&.message&.content
       def chat(message: "", parameters: {})
-        chat_handler.create(message:, parameters:)
+        chat_handler.create(message:, parameters: parameters.dup)
       end
 
-      # Input images should be byte strings. Returns NitroIntelligence::ImageGeneration
+      # Asks a decision model questions about some content over the gateway's System One route,
+      # `POST /v1/systemone`.
+      #
+      # @param state [String, Hash, Array] the content the questions are about
+      # @param questions [Hash] each question's name mapped to its `type` (`noul`, `choice` or
+      #   `score`), `instructions` and, for `choice` and `score`, `criteria`: the options keyed by
+      #   name, or the levels in order
+      # @param parameters [Hash] `model`, `assets` (images `state` and the questions refer to),
+      #   `metadata` and `request_options`, and on an observed client `trace_name`, `trace_seed`,
+      #   `user_id`, `session_id` and `tags`. It takes no managed prompt: `prompt_name` and
+      #   `prompt_variables` are ignored, so the instructions belong in `questions`.
+      # @return [Decision]
+      # @raise [ArgumentError] when the model is not one `model_config` lists with the type
+      #   `decision`
+      # @example
+      #   client.decide(
+      #     state: "I was charged twice.",
+      #     questions: { billing: { type: "noul", instructions: "Is this message about billing?" } }
+      #   ).answers[:billing].noul
+      def decide(state:, questions:, parameters: {})
+        decision_handler.create(state:, questions:, parameters: parameters.dup)
+      end
+
+      # Generates an image, or edits `target_image`.
+      #
+      # @param message [String] what to generate, or how to edit the target
+      # @param target_image [String, nil] the image to edit, as a byte string
+      # @param reference_images [Array<String>] further images to draw on, as byte strings
+      # @param parameters [Hash] `model`, `aspect_ratio` and `resolution`, among those `model_config`
+      #   lists for the model, and this gem's own
+      # @return [ImageGeneration]
+      # @raise [ArgumentError] when the model, aspect ratio or resolution is not one `model_config`
+      #   lists
+      # @example
+      #   result = client.generate_image(message: "A bear installing a window")
+      #   File.binwrite("bear.#{result.generated_image.file_extension}", result.generated_image.byte_string)
       def generate_image(message: "", target_image: nil, reference_images: [], parameters: {})
-        image_handler.create(message:, target_image:, reference_images:, parameters:)
+        image_handler.create(message:, target_image:, reference_images:, parameters: parameters.dup)
       end
 
-      # Audio file should be file object with extension in the filename
-      # Use file_extension for now to prevent a library dependency
+      # Transcribes audio into text.
+      #
+      # @param message [String] a prompt to guide the transcription, such as its vocabulary or style
+      # @param audio_file [File] the audio, opened in binary mode, with a filename whose extension
+      #   gives its format
+      # @param parameters [Hash] openai-ruby's transcription parameters, and this gem's own
+      # @return [OpenAI::Models::Audio::Transcription, OpenAI::Models::Audio::TranscriptionVerbose,
+      #   OpenAI::Models::Audio::TranscriptionDiarized, StringIO] whichever openai-ruby returns for
+      #   the `response_format`: a `Transcription` for `json`, the default, and the response body as a
+      #   StringIO for `text`, `srt` and `vtt`. An observed client records the transcription's `text`,
+      #   so it supports only the formats that return an object.
+      # @raise [ArgumentError] when the model is not one `model_config` lists
+      # @example
+      #   client.transcribe_audio(audio_file: File.open("sample.m4a", "rb")).text
       def transcribe_audio(message: +"", audio_file: nil, parameters: {})
-        audio_transcription_handler.create(message:, audio_file:, parameters:)
+        audio_transcription_handler.create(message:, audio_file:, parameters: parameters.dup)
       end
 
-      # Returns StringIO
+      # Speaks text aloud.
+      #
+      # @param message [String] the text to speak
+      # @param parameters [Hash] openai-ruby's speech parameters, with `voice` and `response_format`
+      #   among those `model_config` lists for the model, and this gem's own
+      # @return [StringIO] the audio
+      # @raise [ArgumentError] when the model, voice or format is not one `model_config` lists
+      # @example
+      #   File.binwrite("hello.mp3", client.text_to_speech(message: "Hello").string)
       def text_to_speech(message: "", parameters: {})
-        text_to_speech_handler.create(message:, parameters:)
+        text_to_speech_handler.create(message:, parameters: parameters.dup)
       end
 
     private
@@ -42,6 +118,10 @@ module NitroIntelligence
         @chat_handler ||= Handlers::ChatHandler.new(client: @client)
       end
 
+      def decision_handler
+        @decision_handler ||= Handlers::DecisionHandler.new(client: @client)
+      end
+
       def image_handler
         @image_handler ||= Handlers::ImageHandler.new(client: @client)
       end

data/lib/nitro_intelligence/client/client.rb

@@ -1,11 +1,31 @@
 require "nitro_intelligence/client/factory"
 
 module NitroIntelligence
+  # Inference through the inference gateway: chat, image generation and editing, audio
+  # transcription and text-to-speech. See the README's "Basic Usage".
   module Client
-    def self.new(observability_project_slug: nil)
-      Factory.new(observability_project_slug:).build
+    # Builds a client, which records every request in Cerebro when given a project to report to.
+    #
+    # A project slug that `observability_projects` does not configure logs a warning and gives a
+    # client without observability rather than raising. Building a client makes no request to
+    # Cerebro, so a configured project with wrong credentials still gives an observed client, and the
+    # error surfaces on the first request that reaches Cerebro, such as a prompt lookup.
+    #
+    # @param observability_project_slug [String, nil] the Cerebro project to record requests in, and
+    #   to fetch prompts from
+    # @param inference_max_retries [Integer, nil] overrides the configured `inference_max_retries`
+    #   for this client
+    # @param inference_timeout [Numeric, nil] overrides the configured `inference_timeout` for this
+    #   client, in seconds
+    # @return [Client::Base]
+    # @example
+    #   client = NitroIntelligence::Client.new(observability_project_slug: "my-feature-project")
+    #   client.chat(message: "Why is the sky blue?")
+    def self.new(observability_project_slug: nil, inference_max_retries: nil, inference_timeout: nil)
+      Factory.new(observability_project_slug:, inference_max_retries:, inference_timeout:).build
     end
 
+    # @api private
     def self.validate_model(model)
       raise ArgumentError, "Unsupported model: '#{model}'" unless NitroIntelligence.model_catalog.exists?(model)
     end

data/lib/nitro_intelligence/client/factory.rb

@@ -5,9 +5,12 @@ require "nitro_intelligence/observability/project"
 
 module NitroIntelligence
   module Client
+    # @api private
     class Factory
-      def initialize(observability_project_slug:)
+      def initialize(observability_project_slug:, inference_max_retries: nil, inference_timeout: nil)
         @observability_project_slug = observability_project_slug
+        @inference_max_retries = inference_max_retries
+        @inference_timeout = inference_timeout
       end
 
       def build
@@ -42,10 +45,16 @@ module NitroIntelligence
         project_client
       end
 
+      # The gateway owns failover between providers, and an upstream 429 it passes
+      # on has already been decided there, so the SDK's own retries default to off
+      # rather than to openai-ruby's two. A nil timeout keeps openai-ruby's default.
       def inference_client
+        config = NitroIntelligence.config
         @inference_client ||= OpenAI::Client.new(
-          api_key: NitroIntelligence.config.inference_api_key,
-          base_url: NitroIntelligence.config.inference_base_url
+          api_key: config.inference_api_key,
+          base_url: config.inference_base_url,
+          max_retries: @inference_max_retries.nil? ? config.inference_max_retries : @inference_max_retries,
+          **{ timeout: @inference_timeout || config.inference_timeout }.compact
         )
       end
     end

data/lib/nitro_intelligence/client/handlers/audio_transcription_handler.rb

@@ -4,6 +4,7 @@ require "nitro_intelligence/client/handlers/base_handler"
 module NitroIntelligence
   module Client
     module Handlers
+      # @api private
       class AudioTranscriptionHandler < BaseHandler
         ALLOWED_EXTRA_PARAMETERS = OpenAI::Models::Audio::TranscriptionCreateParams.fields.keys.uniq.freeze
 
@@ -12,9 +13,10 @@ module NitroIntelligence
           perform_request(audio_file:, message:, parameters:)
         end
 
-        def perform_request(audio_file:, message: "", parameters: {}, correlation_trace_id: nil)
+        def perform_request(audio_file:, message: "", parameters: {}, correlation_trace_id: nil, mask_content: false)
           add_request_headers(parameters, MODALITY_HEADER => "audio", REQUESTED_MODEL_HEADER => parameters[:model])
           add_correlation_headers(parameters, trace_id: correlation_trace_id)
+          add_message_redaction_header(parameters, mask_content:)
           @client.audio.transcriptions.create(
             prompt: message,
             file: audio_file,

data/lib/nitro_intelligence/client/handlers/base_handler.rb

@@ -2,7 +2,9 @@ require "json"
 
 module NitroIntelligence
   module Client
+    # @api private
     module Handlers
+      # @api private
       class BaseHandler
         MODALITY_HEADER = "nip-modality".freeze
         REQUESTED_MODEL_HEADER = "nip-requested-model".freeze
@@ -12,6 +14,10 @@ module NitroIntelligence
         TRACE_ID_HEADER = "x-litellm-trace-id".freeze
         SPEND_LOGS_METADATA_HEADER = "x-litellm-spend-logs-metadata".freeze
 
+        # Asks the gateway to replace the messages and the response with
+        # "redacted-by-litellm" in its own logs. The gateway reads any value as on.
+        MESSAGE_REDACTION_HEADER = "x-litellm-enable-message-redaction".freeze
+
         # Headers over ~8KB are rejected by most proxies. Metadata is caller
         # supplied, so cap it rather than turning a large hash into a failed request.
         MAX_SPEND_LOGS_METADATA_BYTES = 4096
@@ -82,9 +88,12 @@ module NitroIntelligence
           nil
         end
 
+        # Copies the request options rather than changing them, because a caller may pass one
+        # hash to several requests, and the headers added here belong to this request alone.
         def add_request_headers(parameters, headers)
-          request_options = (parameters[:request_options] ||= {})
-          (request_options[:extra_headers] ||= {}).merge!(headers.compact)
+          request_options = parameters[:request_options].to_h
+          parameters[:request_options] =
+            request_options.merge(extra_headers: request_options[:extra_headers].to_h.merge(headers.compact))
           parameters
         end
 
@@ -110,6 +119,12 @@ module NitroIntelligence
           )
         end
 
+        # A project that masks content in Cerebro keeps it out of the gateway's logs
+        # too, so one setting decides where a request's content is recorded.
+        def add_message_redaction_header(parameters, mask_content:)
+          add_request_headers(parameters, MESSAGE_REDACTION_HEADER => ("true" if mask_content))
+        end
+
         def spend_logs_metadata(metadata)
           return nil if metadata.blank?
 

data/lib/nitro_intelligence/client/handlers/chat_handler.rb

@@ -4,6 +4,7 @@ require "nitro_intelligence/client/handlers/base_handler"
 module NitroIntelligence
   module Client
     module Handlers
+      # @api private
       class ChatHandler < BaseHandler
         ALLOWED_EXTRA_PARAMETERS = OpenAI::Models::Chat::CompletionCreateParams.fields.keys.uniq.freeze
 
@@ -12,9 +13,10 @@ module NitroIntelligence
           perform_request(parameters:)
         end
 
-        def perform_request(parameters: {}, correlation_trace_id: nil)
+        def perform_request(parameters: {}, correlation_trace_id: nil, mask_content: false)
           add_request_headers(parameters, REQUESTED_MODEL_HEADER => parameters[:model])
           add_correlation_headers(parameters, trace_id: correlation_trace_id)
+          add_message_redaction_header(parameters, mask_content:)
           @client.chat.completions.create(**parameters.slice(*ALLOWED_EXTRA_PARAMETERS))
         end
 

data/lib/nitro_intelligence/client/handlers/decision_handler.rb

@@ -0,0 +1,75 @@
+require "openai"
+require "nitro_intelligence/client/handlers/base_handler"
+require "nitro_intelligence/decision"
+
+module NitroIntelligence
+  module Client
+    module Handlers
+      # Asks a decision model questions about some content over the gateway's System One route,
+      # `POST /v1/systemone`. openai-ruby has no resource for it, so the request goes through
+      # the client's generic `request`, which keeps its credentials, retries, timeouts and errors.
+      #
+      # @api private
+      class DecisionHandler < BaseHandler
+        # `/v1/systemone` is the route the gateway's docs name and key route limits list.
+        PATH = "v1/systemone".freeze
+        # What is left of that route to add to an `inference_base_url` that already ends in `/v1`.
+        PATH_AFTER_V1 = "systemone".freeze
+
+        # System One request fields other than the model, state and questions, taken from
+        # `parameters`. Everything else there is for the client, the gateway or tracing.
+        BODY_PARAMETERS = %i[assets].freeze
+
+        # Declaring a model, rather than taking the decoded hash, is what gets the response its
+        # `last_response`, and with it the gateway's cost headers. Only `model` is declared:
+        # `Decision` reads `answers` and `usage` raw through `BaseModel#[]`, which openai-ruby
+        # documents as public API for keys outside a model's spec, from 0.86 through 0.99 at least.
+        class Response < OpenAI::BaseModel
+          required :model, String
+        end
+
+        def create(state:, questions:, parameters: {})
+          validate_and_resolve!(parameters)
+          perform_request(state:, questions:, parameters:)
+        end
+
+        def perform_request(state:, questions:, parameters: {}, correlation_trace_id: nil, mask_content: false)
+          add_request_headers(parameters, REQUESTED_MODEL_HEADER => parameters[:model])
+          add_correlation_headers(parameters, trace_id: correlation_trace_id)
+          add_message_redaction_header(parameters, mask_content:)
+          response = @client.request(
+            method: :post,
+            path:,
+            body: { model: parameters[:model], state:, questions:, **parameters.slice(*BODY_PARAMETERS).compact },
+            model: Response,
+            options: parameters[:request_options]
+          )
+          Decision.new(response)
+        end
+
+        def validate_and_resolve!(parameters)
+          default_parameters = {
+            metadata: {},
+            model: NitroIntelligence.model_catalog.default_decision_model&.name,
+          }
+
+          parameters.replace(default_parameters.merge(parameters))
+          Client.validate_model(parameters[:model])
+          return if NitroIntelligence.model_catalog.lookup_by_name(parameters[:model]).is_a?(DecisionModel)
+
+          # A decision request has nothing in common with a chat one, so a model of another type
+          # would fail at the gateway with an error that does not say why.
+          raise ArgumentError, "Model '#{parameters[:model]}' is not a decision model"
+        end
+
+      private
+
+        # Both reach `/v1/systemone`: the whole route on a base URL without a version, or only the
+        # rest of it on one that already ends in `/v1`.
+        def path
+          @client.base_url.path.chomp("/").end_with?("/v1") ? PATH_AFTER_V1 : PATH
+        end
+      end
+    end
+  end
+end

data/lib/nitro_intelligence/client/handlers/image_handler.rb

@@ -5,6 +5,7 @@ require "nitro_intelligence/media/image_generation"
 module NitroIntelligence
   module Client
     module Handlers
+      # @api private
       class ImageHandler < BaseHandler
         ALLOWED_EXTRA_PARAMETERS = OpenAI::Models::Chat::CompletionCreateParams.fields.keys.uniq.freeze
 
@@ -19,9 +20,10 @@ module NitroIntelligence
           image_generation
         end
 
-        def perform_request(parameters: {}, correlation_trace_id: nil)
+        def perform_request(parameters: {}, correlation_trace_id: nil, mask_content: false)
           add_request_headers(parameters, MODALITY_HEADER => "image", REQUESTED_MODEL_HEADER => parameters[:model])
           add_correlation_headers(parameters, trace_id: correlation_trace_id)
+          add_message_redaction_header(parameters, mask_content:)
           @client.chat.completions.create(**parameters.slice(*ALLOWED_EXTRA_PARAMETERS))
         end
 
@@ -40,7 +42,9 @@ module NitroIntelligence
               },
             },
           }
-          parameters.replace(default_parameters.merge(parameters))
+          # Deep, so that request options a caller passes (headers, max_retries,
+          # timeout) sit alongside the image config rather than replacing it.
+          parameters.replace(default_parameters.deep_merge(parameters))
           Client.validate_model(parameters[:model])
         end
 

data/lib/nitro_intelligence/client/handlers/observed/audio_transcription_handler.rb

@@ -6,7 +6,12 @@ module NitroIntelligence
   module Client
     module Handlers
       module Observed
+        # @api private
         class AudioTranscriptionHandler
+          # Raised by {Client::Base#transcribe_audio} on an observed client when the prompt it names
+          # is not a text prompt.
+          #
+          # @api public
           class ObservedAudioTranscriptionPromptError < StandardError; end
 
           def initialize(base_handler:, observer:)
@@ -17,9 +22,8 @@ module NitroIntelligence
           def create(audio_file:, message: +"", parameters: {})
             @base_handler.validate_and_resolve!(parameters)
 
-            # Modifies message in place
-            prompt = handle_prompt(message:, parameters:)
-            trace_name = parameters[:trace_name] || prompt&.name || @observer.project_client.project.slug
+            prompt, message = handle_prompt(message:, parameters:)
+            trace_name = parameters[:trace_name] || prompt&.name || project.slug
 
             @observer.observe(
               "audio-transcription",
@@ -35,12 +39,16 @@ module NitroIntelligence
 
         private
 
+          def project
+            @observer.project_client.project
+          end
+
           def handle_prompt(message:, parameters:)
             prompt = NitroIntelligence::Observability::PromptResolver.for(
-              store: @observer.project_client.project.prompt_store,
+              store: project.prompt_store,
               parameters:
             )
-            return nil if prompt.blank?
+            return [nil, message] if prompt.blank?
 
             # Prompts for audio transcriptions should only be text
             if prompt.type != "text"
@@ -49,25 +57,35 @@ module NitroIntelligence
             end
             interpolated_prompt = prompt.compile(**(parameters[:prompt_variables] || {}))
 
-            message.prepend("#{interpolated_prompt} ").strip!
             parameters.merge!(prompt.config) unless parameters[:prompt_config_disabled]
 
-            prompt
+            # A new string rather than the caller's, which they may send again.
+            [prompt, "#{interpolated_prompt} #{message}".strip]
           end
 
-          def workflow(generation:, message:, audio_file:, parameters:)
-            audio_transcription = @base_handler.perform_request(
-              audio_file:, message:, parameters:, correlation_trace_id: generation.trace_id
-            )
+          # Uploads go to Cerebro's media API directly, past any masking, so a masked
+          # project does not send the audio at all.
+          def upload_input_audio(audio_file, generation)
+            return if project.mask_content?
 
             audio_file.rewind
             upload_handler = NitroIntelligence::Observability::UploadHandler.new(
-              auth_token: @observer.project_client.project.auth_token
+              auth_token: project.auth_token
             )
             upload_handler.upload(
               generation.trace_id,
+              observation_id: generation.id,
               upload_queue: Queue.new([NitroIntelligence::Audio.new(audio_file)])
             )
+          end
+
+          def workflow(generation:, message:, audio_file:, parameters:)
+            audio_transcription = @base_handler.perform_request(
+              audio_file:, message:, parameters:,
+              correlation_trace_id: generation.trace_id, mask_content: project.mask_content?
+            )
+
+            upload_input_audio(audio_file, generation)
 
             # The model and input are already recorded on the observation before the
             # request runs, and the response carries no model of its own

data/lib/nitro_intelligence/client/handlers/observed/chat_handler.rb

@@ -3,8 +3,15 @@ require "nitro_intelligence/observability/prompt_resolver"
 module NitroIntelligence
   module Client
     module Handlers
+      # @api private
       module Observed
+        # @api private
         class ChatHandler
+          # Raised by {Client::Base#chat} on an observed client when the request carries no user
+          # message for the model to answer: the prompt supplies only a system message, and the
+          # caller passed no `message`.
+          #
+          # @api public
           class ObservedChatPromptError < StandardError; end
 
           def initialize(base_handler:, observer:)
@@ -18,7 +25,7 @@ module NitroIntelligence
             prompt = handle_prompt(parameters:)
             validate_message_shape!(parameters:, prompt:)
 
-            trace_name = parameters[:trace_name] || prompt&.name || @observer.project_client.project.slug
+            trace_name = parameters[:trace_name] || prompt&.name || project.slug
 
             @observer.observe(
               "chat-completion",
@@ -34,6 +41,10 @@ module NitroIntelligence
 
         private
 
+          def project
+            @observer.project_client.project
+          end
+
           # A chat completion needs a turn for the model to answer, and it is the model's own
           # chat template that insists on one: Qwen's raises "No user query found in
           # messages." Left to the gateway that costs a round trip and comes back as a 400
@@ -76,7 +87,7 @@ module NitroIntelligence
 
           def handle_prompt(parameters:)
             prompt = NitroIntelligence::Observability::PromptResolver.for(
-              store: @observer.project_client.project.prompt_store,
+              store: project.prompt_store,
               parameters:
             )
             return nil if prompt.blank?
@@ -92,7 +103,9 @@ module NitroIntelligence
           end
 
           def workflow(generation:, parameters:)
-            chat_completion = @base_handler.perform_request(parameters:, correlation_trace_id: generation.trace_id)
+            chat_completion = @base_handler.perform_request(
+              parameters:, correlation_trace_id: generation.trace_id, mask_content: project.mask_content?
+            )
             output = chat_completion.choices.first.message.to_h
 
             trace_attributes = {

data/lib/nitro_intelligence/client/handlers/observed/decision_handler.rb

@@ -0,0 +1,72 @@
+module NitroIntelligence
+  module Client
+    module Handlers
+      module Observed
+        # Recorded as a generation, the closest observation type Langfuse has to a decision: it
+        # carries the model, usage and cost. How a decision should be recorded is still open in
+        # powerhome/cerebro#128.
+        #
+        # @api private
+        class DecisionHandler
+          def initialize(base_handler:, observer:)
+            @base_handler = base_handler
+            @observer = observer
+          end
+
+          def create(state:, questions:, parameters: {})
+            @base_handler.validate_and_resolve!(parameters)
+
+            trace_name = parameters[:trace_name] || project.slug
+
+            # The questions alone say nothing about why the model answered as it did, so the
+            # state they were asked about is recorded with them.
+            @observer.observe(
+              "decision",
+              type: :generation,
+              parameters:,
+              trace_name:,
+              input: { state:, questions: }
+            ) do |generation|
+              workflow(generation:, state:, questions:, parameters:)
+            end
+          end
+
+        private
+
+          def project
+            @observer.project_client.project
+          end
+
+          def workflow(generation:, state:, questions:, parameters:)
+            decision = @base_handler.perform_request(
+              state:, questions:, parameters:, correlation_trace_id: generation.trace_id,
+              mask_content: project.mask_content?
+            )
+
+            trace_attributes = {
+              model: decision.model,
+              # A copy, because the observer truncates the output it records in place.
+              output: decision.raw_answers.deep_dup,
+              usage_details: usage_details(decision.usage),
+              cost_details: @base_handler.cost_details(decision),
+            }
+
+            [decision, trace_attributes]
+          end
+
+          def usage_details(usage)
+            input_tokens = usage[:input_tokens]
+            output_tokens = usage[:output_tokens]
+            return nil if input_tokens.nil? && output_tokens.nil?
+
+            {
+              prompt_tokens: input_tokens,
+              completion_tokens: output_tokens,
+              total_tokens: (input_tokens + output_tokens if input_tokens && output_tokens),
+            }.compact
+          end
+        end
+      end
+    end
+  end
+end