nitro_intelligence 4.0.0 → 4.1.1

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,8 @@ module NitroIntelligence
           perform_request(audio_file:, message:, parameters:)
         end
 
-        def perform_request(audio_file:, message: "", parameters: {}, correlation_trace_id: nil)
-          add_request_headers(parameters, MODALITY_HEADER => "audio", REQUESTED_MODEL_HEADER => parameters[:model])
-          add_correlation_headers(parameters, trace_id: correlation_trace_id)
+        def perform_request(audio_file:, message: "", parameters: {}, correlation_trace_id: nil, mask_content: false)
+          add_gateway_headers(parameters, correlation_trace_id:, mask_content:, modality: "audio")
           @client.audio.transcriptions.create(
             prompt: message,
             file: audio_file,
@@ -29,7 +29,7 @@ module NitroIntelligence
           }
 
           parameters.replace(default_parameters.merge(parameters))
-          Client.validate_model(parameters[:model])
+          Client.validate_model(parameters[:model], type: TextModel)
         end
       end
     end

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,21 @@ module NitroIntelligence
           )
         end
 
+        # Every request to the gateway names the model asked for, and the modality where the
+        # route alone does not, carries the correlation headers and, for a project that masks
+        # content, asks the gateway to redact it from its logs.
+        def add_gateway_headers(parameters, correlation_trace_id:, mask_content:, modality: nil)
+          add_request_headers(parameters, MODALITY_HEADER => modality, REQUESTED_MODEL_HEADER => parameters[:model])
+          add_correlation_headers(parameters, trace_id: correlation_trace_id)
+          add_message_redaction_header(parameters, mask_content:)
+        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,8 @@ module NitroIntelligence
           perform_request(parameters:)
         end
 
-        def perform_request(parameters: {}, correlation_trace_id: nil)
-          add_request_headers(parameters, REQUESTED_MODEL_HEADER => parameters[:model])
-          add_correlation_headers(parameters, trace_id: correlation_trace_id)
+        def perform_request(parameters: {}, correlation_trace_id: nil, mask_content: false)
+          add_gateway_headers(parameters, correlation_trace_id:, mask_content:)
           @client.chat.completions.create(**parameters.slice(*ALLOWED_EXTRA_PARAMETERS))
         end
 
@@ -31,7 +31,7 @@ module NitroIntelligence
           }
 
           parameters.replace(default_parameters.merge(parameters))
-          Client.validate_model(parameters[:model])
+          Client.validate_model(parameters[:model], type: TextModel)
         end
       end
     end

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

@@ -0,0 +1,68 @@
+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_gateway_headers(parameters, correlation_trace_id:, 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], type: DecisionModel)
+        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,12 +20,13 @@ module NitroIntelligence
           image_generation
         end
 
-        def perform_request(parameters: {}, correlation_trace_id: nil)
-          add_request_headers(parameters, MODALITY_HEADER => "image", REQUESTED_MODEL_HEADER => parameters[:model])
-          add_correlation_headers(parameters, trace_id: correlation_trace_id)
+        def perform_request(parameters: {}, correlation_trace_id: nil, mask_content: false)
+          add_gateway_headers(parameters, correlation_trace_id:, mask_content:, modality: "image")
           @client.chat.completions.create(**parameters.slice(*ALLOWED_EXTRA_PARAMETERS))
         end
 
+        # The model is not checked here: building the ImageGeneration already checked it, along
+        # with the aspect ratio and resolution it supports.
         def validate_and_resolve!(parameters, image_generation)
           default_parameters = {
             image_generation:,
@@ -43,7 +45,6 @@ module NitroIntelligence
           # 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
 
       private

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

@@ -1,3 +1,4 @@
+require "nitro_intelligence/client/handlers/observed/base_handler"
 require "base64"
 require "nitro_intelligence/media/audio"
 require "nitro_intelligence/observability/prompt_resolver"
@@ -6,20 +7,19 @@ module NitroIntelligence
   module Client
     module Handlers
       module Observed
-        class AudioTranscriptionHandler
+        # @api private
+        class AudioTranscriptionHandler < BaseHandler
+          # 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:)
-            @base_handler = base_handler
-            @observer = observer
-          end
-
           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",
@@ -37,10 +37,10 @@ module NitroIntelligence
 
           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 +49,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/base_handler.rb

@@ -0,0 +1,23 @@
+module NitroIntelligence
+  module Client
+    module Handlers
+      module Observed
+        # Wraps a handler's request in an observation recorded to the observer's Cerebro project.
+        #
+        # @api private
+        class BaseHandler
+          def initialize(base_handler:, observer:)
+            @base_handler = base_handler
+            @observer = observer
+          end
+
+        private
+
+          def project
+            @observer.project_client.project
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -1,24 +1,27 @@
+require "nitro_intelligence/client/handlers/observed/base_handler"
 require "nitro_intelligence/observability/prompt_resolver"
 
 module NitroIntelligence
   module Client
     module Handlers
+      # @api private
       module Observed
-        class ChatHandler
+        # @api private
+        class ChatHandler < BaseHandler
+          # 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:)
-            @base_handler = base_handler
-            @observer = observer
-          end
-
           def create(message: "", parameters: {})
             @base_handler.validate_and_resolve!(parameters, message)
 
             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",
@@ -76,7 +79,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 +95,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,65 @@
+require "nitro_intelligence/client/handlers/observed/base_handler"
+
+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 < BaseHandler
+          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 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

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

@@ -1,3 +1,4 @@
+require "nitro_intelligence/client/handlers/observed/base_handler"
 require "nitro_intelligence/media/image_generation"
 require "nitro_intelligence/observability/prompt_resolver"
 
@@ -5,12 +6,8 @@ module NitroIntelligence
   module Client
     module Handlers
       module Observed
-        class ImageHandler
-          def initialize(base_handler:, observer:)
-            @base_handler = base_handler
-            @observer = observer
-          end
-
+        # @api private
+        class ImageHandler < BaseHandler
           def create(message: "", target_image: nil, reference_images: [], parameters: {})
             image_generation = build_image_generation(message:, target_image:, reference_images:, parameters:)
 
@@ -18,7 +15,7 @@ module NitroIntelligence
 
             # Modifies parameters in place
             prompt = handle_prompt(parameters:)
-            trace_name = parameters[:trace_name] || prompt&.name || @observer.project_client.project.slug
+            trace_name = parameters[:trace_name] || prompt&.name || project.slug
 
             @observer.observe(
               "image-generation",
@@ -44,12 +41,17 @@ module NitroIntelligence
           end
 
           def handle_image_generation_uploads(input, output, image_generation)
+            # Uploads go to Cerebro's media API directly, past any masking, and a masked
+            # project records no input or output for a media reference to sit in.
+            return if project.mask_content?
+
             # If we are doing image generation we should upload the media to observability manually
             upload_handler = NitroIntelligence::Observability::UploadHandler.new(
-              auth_token: @observer.project_client.project.auth_token
+              auth_token: project.auth_token
             )
             upload_handler.upload(
               image_generation.trace_id,
+              observation_id: image_generation.observation_id,
               upload_queue: Queue.new(image_generation.files)
             )
 
@@ -60,7 +62,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?
@@ -76,12 +78,17 @@ module NitroIntelligence
           end
 
           def workflow(generation:, image_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?
+            )
 
             image_generation.trace_id = generation.trace_id
+            image_generation.observation_id = generation.id
             image_generation.parse_file(chat_completion)
 
-            input = parameters[:messages]
+            # A copy, because recording replaces its image URLs and the caller may send the same
+            # messages again.
+            input = parameters[:messages].deep_dup
             output = chat_completion.choices.first.message.to_h
             handle_image_generation_uploads(input, output, image_generation)
 

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

@@ -1,3 +1,4 @@
+require "nitro_intelligence/client/handlers/observed/base_handler"
 require "nitro_intelligence/media/audio"
 require "nitro_intelligence/observability/prompt_resolver"
 
@@ -5,20 +6,20 @@ module NitroIntelligence
   module Client
     module Handlers
       module Observed
-        class TextToSpeechHandler
+        # @api private
+        class TextToSpeechHandler < BaseHandler
+          # Raised by {Client::Base#text_to_speech} on an observed client when the prompt it names
+          # is not a text prompt.
+          #
+          # @api public
           class ObservedTextToSpeechPromptError < StandardError; end
 
-          def initialize(base_handler:, observer:)
-            @base_handler = base_handler
-            @observer = observer
-          end
-
           def create(message: "", parameters: {})
             prompt = handle_prompt(parameters:)
 
             @base_handler.validate_and_resolve!(parameters)
 
-            trace_name = parameters[:trace_name] || prompt&.name || @observer.project_client.project.slug
+            trace_name = parameters[:trace_name] || prompt&.name || project.slug
 
             @observer.observe(
               "text-to-speech",
@@ -28,7 +29,7 @@ module NitroIntelligence
               prompt:,
               input: message
             ) do |generation|
-              workflow(message:, parameters:, trace_id: generation.trace_id)
+              workflow(generation:, message:, parameters:)
             end
           end
 
@@ -36,7 +37,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?
@@ -54,21 +55,30 @@ module NitroIntelligence
             prompt
           end
 
-          def handle_text_to_speech_upload(tts_file, trace_id)
+          def handle_text_to_speech_upload(tts_file, trace_id, observation_id:)
+            # Uploads go to Cerebro's media API directly, past any masking. A masked
+            # project has no media reference to record; the observer masks the output.
+            return "" if project.mask_content?
+
             upload_handler = NitroIntelligence::Observability::UploadHandler.new(
-              auth_token: @observer.project_client.project.auth_token
+              auth_token: project.auth_token
             )
+            audio = Audio.new(tts_file)
+            audio.direction = "output"
+
             upload_handler.upload(
               trace_id,
-              upload_queue: Queue.new([Audio.new(tts_file)])
+              observation_id:,
+              upload_queue: Queue.new([audio])
             )
 
             uploaded_media = upload_handler.uploaded_media.first
             "@@@langfuseMedia:type=#{uploaded_media.mime_type}|id=#{uploaded_media.reference_id}|source=bytes@@@"
           end
 
-          def workflow(message:, parameters:, trace_id:)
-            tts = @base_handler.perform_request(message:, parameters:, correlation_trace_id: trace_id)
+          def workflow(generation:, message:, parameters:)
+            tts = @base_handler.perform_request(message:, parameters:, correlation_trace_id: generation.trace_id,
+                                                mask_content: project.mask_content?)
             output = ""
 
             Tempfile.create(["tts", ".#{parameters[:response_format]}"]) do |tempfile|
@@ -76,7 +86,7 @@ module NitroIntelligence
               tempfile.write(tts.string)
               tempfile.rewind
 
-              output = handle_text_to_speech_upload(tempfile, trace_id)
+              output = handle_text_to_speech_upload(tempfile, generation.trace_id, observation_id: generation.id)
             end
 
             # Usage details and the resolved model are read off a response body, and

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

@@ -4,6 +4,7 @@ require "nitro_intelligence/client/handlers/base_handler"
 module NitroIntelligence
   module Client
     module Handlers
+      # @api private
       class TextToSpeechHandler < BaseHandler
         ALLOWED_EXTRA_PARAMETERS = OpenAI::Models::Audio::SpeechCreateParams.fields.keys.uniq.freeze
 
@@ -12,9 +13,8 @@ module NitroIntelligence
           perform_request(message:, parameters:)
         end
 
-        def perform_request(message: "", parameters: {}, correlation_trace_id: nil)
-          add_request_headers(parameters, MODALITY_HEADER => "audio", REQUESTED_MODEL_HEADER => parameters[:model])
-          add_correlation_headers(parameters, trace_id: correlation_trace_id)
+        def perform_request(message: "", parameters: {}, correlation_trace_id: nil, mask_content: false)
+          add_gateway_headers(parameters, correlation_trace_id:, mask_content:, modality: "audio")
           @client.audio.speech.create(
             input: message,
             **parameters.slice(*ALLOWED_EXTRA_PARAMETERS)
@@ -23,8 +23,7 @@ module NitroIntelligence
 
         def validate_and_resolve!(parameters)
           model_name = parameters[:model] || NitroIntelligence.model_catalog.default_text_to_speech_model&.name
-          model = NitroIntelligence.model_catalog.lookup_by_name(model_name)
-          raise ArgumentError, "Unsupported model: '#{model_name}'" unless model
+          model = Client.validate_model(model_name, type: TextToSpeechModel)
 
           default_parameters = {
             metadata: {},