nitro_intelligence 0.0.1 → 1.0.0

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

@@ -0,0 +1,74 @@
+module NitroIntelligence
+  module Client
+    module Handlers
+      module Observed
+        class ChatHandler
+          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:)
+            trace_name = parameters[:trace_name] || prompt&.name || @observer.project_client.project.slug
+
+            @observer.observe(
+              "chat-completion",
+              type: :generation,
+              parameters:,
+              trace_name:,
+              prompt:
+            ) do |_generation|
+              workflow(parameters:)
+            end
+          end
+
+        private
+
+          def handle_prompt(parameters:)
+            return nil if parameters[:prompt_name].blank?
+
+            prompt = @observer.project_client.project.prompt_store.get_prompt(
+              prompt_name: parameters[:prompt_name],
+              prompt_label: parameters[:prompt_label],
+              prompt_version: parameters[:prompt_version]
+            )
+            prompt_variables = parameters[:prompt_variables] || {}
+
+            if prompt.present?
+              parameters[:messages] = prompt.interpolate(
+                messages: parameters[:messages],
+                variables: prompt_variables
+              )
+
+              parameters.merge!(prompt.config) unless parameters[:prompt_config_disabled]
+            end
+
+            prompt
+          end
+
+          def workflow(parameters:)
+            chat_completion = @base_handler.perform_request(parameters:)
+            input = parameters[:messages]
+            output = chat_completion.choices.first.message.to_h
+
+            trace_attributes = {
+              model: chat_completion.model,
+              input:,
+              output:,
+              usage_details: {
+                prompt_tokens: chat_completion.usage.prompt_tokens,
+                completion_tokens: chat_completion.usage.completion_tokens,
+                total_tokens: chat_completion.usage.total_tokens,
+              },
+            }
+
+            [chat_completion, trace_attributes]
+          end
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,109 @@
+require "nitro_intelligence/media/image_generation"
+
+module NitroIntelligence
+  module Client
+    module Handlers
+      module Observed
+        class ImageHandler
+          def initialize(base_handler:, observer:)
+            @base_handler = base_handler
+            @observer = observer
+          end
+
+          def create(message: "", target_image: nil, reference_images: [], parameters: {})
+            image_generation = build_image_generation(message:, target_image:, reference_images:, parameters:)
+
+            @base_handler.validate_and_resolve!(parameters, image_generation)
+
+            # Modifies parameters in place
+            prompt = handle_prompt(parameters:)
+            trace_name = parameters[:trace_name] || prompt&.name || @observer.project_client.project.slug
+
+            @observer.observe(
+              "image-generation",
+              type: :generation,
+              parameters:,
+              trace_name:,
+              prompt:
+            ) do |generation|
+              workflow(generation:, image_generation:, parameters:)
+            end
+
+            image_generation
+          end
+
+        private
+
+          def build_image_generation(message:, target_image:, reference_images:, parameters:)
+            NitroIntelligence::ImageGeneration.new(message:, target_image:, reference_images:) do |config|
+              config.aspect_ratio = parameters[:aspect_ratio] if parameters.key?(:aspect_ratio)
+              config.model = parameters[:model] if parameters.key?(:model)
+              config.resolution = parameters[:resolution] if parameters.key?(:resolution)
+            end
+          end
+
+          def handle_image_generation_uploads(input, output, image_generation)
+            # 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
+            )
+            upload_handler.upload(
+              image_generation.trace_id,
+              upload_queue: Queue.new(image_generation.files)
+            )
+
+            # Replace base64 strings with media references
+            upload_handler.replace_base64_with_media_references(input)
+            upload_handler.replace_base64_with_media_references(output)
+          end
+
+          def handle_prompt(parameters:)
+            return nil if parameters[:prompt_name].blank?
+
+            prompt = @observer.project_client.project.prompt_store.get_prompt(
+              prompt_name: parameters[:prompt_name],
+              prompt_label: parameters[:prompt_label],
+              prompt_version: parameters[:prompt_version]
+            )
+            prompt_variables = parameters[:prompt_variables] || {}
+
+            if prompt.present?
+              parameters[:messages] = prompt.interpolate(
+                messages: parameters[:messages],
+                variables: prompt_variables
+              )
+
+              parameters.merge!(prompt.config) unless parameters[:prompt_config_disabled]
+            end
+
+            prompt
+          end
+
+          def workflow(generation:, image_generation:, parameters:)
+            chat_completion = @base_handler.perform_request(parameters:)
+
+            image_generation.trace_id = generation.trace_id
+            image_generation.parse_file(chat_completion)
+
+            input = parameters[:messages]
+            output = chat_completion.choices.first.message.to_h
+            handle_image_generation_uploads(input, output, image_generation)
+
+            trace_attributes = {
+              model: chat_completion.model,
+              input:,
+              output:,
+              usage_details: {
+                prompt_tokens: chat_completion.usage.prompt_tokens,
+                completion_tokens: chat_completion.usage.completion_tokens,
+                total_tokens: chat_completion.usage.total_tokens,
+              },
+            }
+
+            [chat_completion, trace_attributes]
+          end
+        end
+      end
+    end
+  end
+end

data/lib/nitro_intelligence/client/observed.rb

@@ -0,0 +1,38 @@
+require "nitro_intelligence/client/handlers/observed/audio_transcription_handler"
+require "nitro_intelligence/client/handlers/observed/chat_handler"
+require "nitro_intelligence/client/handlers/observed/image_handler"
+require "nitro_intelligence/observability/prompt_store"
+require "nitro_intelligence/observability/upload_handler"
+require "nitro_intelligence/trace"
+
+module NitroIntelligence
+  module Client
+    class Observed < Base
+      def initialize(client:, observer:)
+        super(client:)
+        @observer = observer
+      end
+
+    private
+
+      def chat_handler
+        @chat_handler ||= Handlers::Observed::ChatHandler.new(
+          base_handler: Handlers::ChatHandler.new(client: @client), observer: @observer
+        )
+      end
+
+      def audio_transcription_handler
+        @audio_transcription_handler ||=
+          Handlers::Observed::AudioTranscriptionHandler.new(
+            base_handler: Handlers::AudioTranscriptionHandler.new(client: @client), observer: @observer
+          )
+      end
+
+      def image_handler
+        @image_handler ||= Handlers::Observed::ImageHandler.new(
+          base_handler: Handlers::ImageHandler.new(client: @client), observer: @observer
+        )
+      end
+    end
+  end
+end

data/lib/nitro_intelligence/client/observers/langfuse_observer.rb

@@ -0,0 +1,75 @@
+module NitroIntelligence
+  module Client
+    module Observers
+      class LangfuseObserver
+        attr_reader :project_client
+
+        def initialize(project_client:)
+          @project_client = project_client
+        end
+
+        def observe(operation_name, type:, parameters:, trace_name:, prompt: nil) # rubocop:disable Metrics/AbcSize, Metrics/MethodLength
+          metadata = parameters[:metadata]
+          seed = parameters[:trace_seed]
+          user_id = parameters[:user_id] || NitroIntelligence.configuration.observability_user_id
+          trace_id = NitroIntelligence::Trace.create_id(seed:) if seed.present?
+
+          if prompt
+            metadata[:prompt_name] = prompt.name
+            metadata[:prompt_version] = prompt.version
+          end
+
+          metadata = metadata.transform_values(&:to_s)
+
+          Langfuse.propagate_attributes(
+            user_id:,
+            metadata:
+          ) do
+            @project_client.observability_client.observe(
+              operation_name,
+              as_type: type,
+              trace_id:,
+              environment: NitroIntelligence.environment.to_s,
+              model: parameters[:model],
+              metadata:
+            ) do |generation|
+              generation.update_trace(name: trace_name, release: NitroIntelligence.configuration.current_revision)
+              generation.update({ prompt: { name: prompt.name, version: prompt.version } }) if prompt
+
+              result, trace_attributes = yield(generation)
+
+              if trace_attributes
+                handle_truncation(trace_attributes[:input], trace_attributes[:output], trace_attributes[:model])
+
+                generation.model = trace_attributes[:model] if trace_attributes[:model]
+                generation.usage_details = trace_attributes[:usage_details] if trace_attributes[:usage_details]
+                generation.input = trace_attributes[:input] if trace_attributes[:input]
+                generation.output = trace_attributes[:output] if trace_attributes[:output]
+
+                generation.update_trace(input: trace_attributes[:input], output: trace_attributes[:output])
+              end
+
+              result
+            end
+          end
+        end
+
+      private
+
+        def handle_truncation(_input, output, model_name)
+          model = NitroIntelligence.model_catalog.lookup_by_name(model_name)
+
+          return unless model&.omit_output_fields
+
+          model.omit_output_fields.each do |omit_output_field|
+            last_key = omit_output_field.last
+            parent_keys = omit_output_field[0...-1]
+            parent = parent_keys.empty? ? output : output.dig(*parent_keys)
+
+            parent[last_key] = "[Truncated...]" if parent.is_a?(Hash) && parent.key?(last_key)
+          end
+        end
+      end
+    end
+  end
+end

data/lib/nitro_intelligence/configuration.rb

@@ -8,6 +8,7 @@ module NitroIntelligence
 
     config_accessor :logger, default: Logger.new($stdout)
     config_accessor :cache_provider, default: NitroIntelligence::NullCache.new
+    config_accessor :current_revision, default: ""
     config_accessor :environment, default: "test"
     config_accessor :agent_server_config, default: {}
     config_accessor :inference_api_key, default: ""
@@ -15,6 +16,7 @@ module NitroIntelligence
     config_accessor :model_config, default: {}
     config_accessor :observability_base_url, default: ""
     config_accessor :observability_projects, default: []
+    config_accessor :observability_user_id, default: ""
 
     class << self
       def configure

data/lib/nitro_intelligence/langfuse_extension.rb

@@ -4,6 +4,9 @@
 # The content of this file should eventually make its way upstream. Setting a custom trace ID is
 # already awaiting approval in https://github.com/simplepractice/langfuse-rb/pull/69.
 #
+
+require "nitro_intelligence/langfuse_tracer_provider"
+
 module NitroIntelligence
   class LangfuseExtension
     attr_reader :config
@@ -14,7 +17,7 @@ module NitroIntelligence
 
       @config = config
       @client = Langfuse::Client.new(config)
-      @tracer_provider = LangfuseTracerProvider.new(config)
+      @tracer_provider = NitroIntelligence::LangfuseTracerProvider.new(config)
     end
 
     def shutdown(timeout: 30)
@@ -40,14 +43,11 @@ module NitroIntelligence
       )
     end
 
-    def start_observation(name, attrs = {}, as_type: :span, parent_span_context: nil, start_time: nil,
+    def start_observation(name, attrs = {}, as_type: :span, trace_id: nil, parent_span_context: nil, start_time: nil, # rubocop:disable Metrics/ParameterLists
                           skip_validation: false)
+      parent_span_context = Langfuse.send(:resolve_trace_context, trace_id, parent_span_context)
       type_str = as_type.to_s
-
-      unless skip_validation || Langfuse.send(:valid_observation_type?, as_type)
-        valid_types = Langfuse::OBSERVATION_TYPES.values.sort.join(", ")
-        raise ArgumentError, "Invalid observation type: #{type_str}. Valid types: #{valid_types}"
-      end
+      Langfuse.send(:validate_observation_type!, as_type, type_str) unless skip_validation
 
       otel_tracer = @tracer_provider.tracer
       otel_span = Langfuse.send(:create_otel_span,
@@ -73,54 +73,11 @@ module NitroIntelligence
     end
 
     def observe(name, attrs = {}, as_type: :span, trace_id: nil, **kwargs, &block)
-      # Merge positional attrs and keyword kwargs
       merged_attrs = attrs.to_h.merge(kwargs)
-      unless trace_id.nil?
-        unless valid_trace_id?(trace_id)
-          raise ArgumentError, "#{trace_id} is not a valid 32 lowercase hex char Langfuse trace ID"
-        end
+      observation = start_observation(name, merged_attrs, as_type:, trace_id:)
+      return observation unless block
 
-        parent_span_context = create_span_context_with_trace_id(trace_id)
-      end
-      observation = start_observation(name, merged_attrs, as_type:, parent_span_context:)
-
-      if block
-        # Block-based API: auto-ends when block completes
-        # Set context and execute block
-        current_context = OpenTelemetry::Context.current
-        begin
-          result = OpenTelemetry::Context.with_current(
-            OpenTelemetry::Trace.context_with_span(observation.otel_span, parent_context: current_context)
-          ) do
-            yield(observation)
-          end
-        ensure
-          # Only end if not already ended (events auto-end in start_observation)
-          observation.end unless as_type.to_s == Langfuse::OBSERVATION_TYPES[:event]
-        end
-        result
-      else
-        # Stateful API - return observation
-        # Events already auto-ended in start_observation
-        observation
-      end
-    end
-
-  private
-
-    def valid_trace_id?(trace_id)
-      !!(trace_id =~ /^[0-9a-f]{32}$/)
-    end
-
-    def create_span_context_with_trace_id(trace_id_as_hex_str)
-      trace_id_as_byte_str = [trace_id_as_hex_str].pack("H*")
-      # NOTE: trace_flags must be SAMPLED or the trace will not appear in Langfuse.
-      # The Python SDK does the same: https://github.com/langfuse/langfuse-python/blob/v4.0.0/langfuse/_client/client.py#L1568
-      trace_flags = OpenTelemetry::Trace::TraceFlags::SAMPLED
-      OpenTelemetry::Trace::SpanContext.new(
-        trace_id: trace_id_as_byte_str,
-        trace_flags:
-      )
+      observation.send(:run_in_context, &block)
     end
   end
 end

data/lib/nitro_intelligence/media/audio.rb

@@ -0,0 +1,50 @@
+require "nitro_intelligence/media/media"
+
+module NitroIntelligence
+  class Audio < Media
+    class AudioFileFormatError < StandardError; end
+
+    def initialize(file)
+      # TODO: Consider a library for dealing with audio files. Dirty implementation
+      raise AudioFileFormatError unless file.respond_to?(:to_path)
+
+      file_extension = File.basename(file).split(".").last
+      file = file.read
+      super
+
+      @file_extension = file_extension
+      @file_type = "audio"
+      @mime_type = determine_mime_type
+    end
+
+  private
+
+    # These are supported by Langfuse
+    def determine_mime_type
+      ext = @file_extension.downcase
+
+      case ext
+      when "m4a", "m4b", "m4r", "m4p", "mp4"
+        "audio/mp4"
+      when "mp3", "mpga", "mp2", "mp1"
+        "audio/mp3"
+      when "wav", "wave"
+        "audio/wav"
+      when "webm", "weba"
+        "audio/webm"
+      when "ogg", "spx"
+        "audio/ogg"
+      when "oga"
+        "audio/oga"
+      when "aac"
+        "audio/aac"
+      when "flac"
+        "audio/flac"
+      when "opus"
+        "audio/opus"
+      else
+        "audio/#{ext}"
+      end
+    end
+  end
+end

data/lib/nitro_intelligence/media/image_generation.rb

@@ -9,7 +9,6 @@ require "nitro_intelligence/media/image"
 module NitroIntelligence
   class ImageGeneration
     class Config
-      CUSTOM_PARAMS = %i[aspect_ratio image_generation resolution].freeze
       DEFAULT_ASPECT_RATIO = "1:1".freeze
       DEFAULT_RESOLUTION = "1K".freeze
 
@@ -17,7 +16,7 @@ module NitroIntelligence
 
       def initialize
         @aspect_ratio = DEFAULT_ASPECT_RATIO
-        @model = NitroIntelligence.model_catalog.default_image_model.name
+        @model = NitroIntelligence.model_catalog.default_image_model&.name
         @resolution = DEFAULT_RESOLUTION
       end
     end
@@ -61,6 +60,9 @@ module NitroIntelligence
       @generated_image = Image.from_base64(base64_string)
       @generated_image.direction = "output"
       @generated_image
+    rescue ArgumentError
+      NitroIntelligence.logger.info("Skipping image parse due to invalid base64; likely already parsed.")
+      nil
     end
 
   private

data/lib/nitro_intelligence/models/model_catalog.rb

@@ -2,16 +2,21 @@ require "nitro_intelligence/models/model_factory"
 
 module NitroIntelligence
   class ModelCatalog
-    attr_reader :models, :default_text_model, :default_image_model
+    attr_reader :models, :default_audio_transcription_model, :default_image_model, :default_text_model
 
     def initialize(model_config)
       @models = (model_config[:models] || []).map { |model_metadata| ModelFactory.build(model_metadata) }
-      @default_text_model = lookup_by_name(model_config[:default_text_model])
+      @default_audio_transcription_model = lookup_by_name(model_config[:default_audio_transcription_model])
       @default_image_model = lookup_by_name(model_config[:default_image_model])
+      @default_text_model = lookup_by_name(model_config[:default_text_model])
     end
 
     def lookup_by_name(name)
       @models.find { |model| model.name == name }
     end
+
+    def exists?(name)
+      lookup_by_name(name).present?
+    end
   end
 end

data/lib/nitro_intelligence/observability/project.rb

@@ -0,0 +1,33 @@
+require "base64"
+require "nitro_intelligence/observability/prompt_store"
+
+module NitroIntelligence
+  module Observability
+    class Project
+      class NotFoundError < StandardError; end
+
+      attr_reader :slug, :id, :public_key, :secret_key, :auth_token, :prompt_store
+
+      def initialize(slug:, id:, public_key:, secret_key:, **_kwargs)
+        @slug = slug
+        @id = id
+        @public_key = public_key
+        @secret_key = secret_key
+        @auth_token = Base64.strict_encode64("#{public_key}:#{secret_key}")
+        @prompt_store = PromptStore.new(
+          observability_project_slug: slug,
+          observability_public_key: public_key,
+          observability_secret_key: secret_key
+        )
+      end
+
+      def self.find_by_slug(slug:)
+        project_config = NitroIntelligence.config.observability_projects.find { |project| project["slug"] == slug }
+
+        return new(**project_config.to_h.transform_keys(&:to_sym)) if project_config
+
+        nil
+      end
+    end
+  end
+end

data/lib/nitro_intelligence/observability/project_client.rb

@@ -0,0 +1,18 @@
+module NitroIntelligence
+  module Observability
+    class ProjectClient
+      class NotFoundError < StandardError; end
+
+      attr_reader :project, :observability_client
+
+      def initialize(project:, observability_client:)
+        @project = project
+        @observability_client = observability_client
+      end
+
+      def shutdown
+        @observability_client.shutdown
+      end
+    end
+  end
+end

data/lib/nitro_intelligence/observability/project_client_registry.rb

@@ -0,0 +1,44 @@
+require "nitro_intelligence/langfuse_extension"
+require "nitro_intelligence/observability/project"
+require "nitro_intelligence/observability/project_client"
+
+module NitroIntelligence
+  module Observability
+    class ProjectClientRegistry
+      def initialize(base_url:)
+        @base_url = base_url
+        @project_clients = {}
+      end
+
+      def fetch(slug)
+        @project_clients[slug] ||= build_project_client(slug)
+      end
+
+      def shutdown_all
+        @project_clients.values.each(&:shutdown)
+      end
+
+    private
+
+      def build_project_client(slug)
+        project = Observability::Project.find_by_slug(slug:)
+
+        unless project
+          raise NitroIntelligence::Observability::Project::NotFoundError,
+                "No observability project config found for slug: #{slug}"
+        end
+
+        observability_client = NitroIntelligence::LangfuseExtension.new do |config|
+          config.public_key = project.public_key
+          config.secret_key = project.secret_key
+          config.base_url = @base_url
+          # Default flush of 60 seconds can be too quick when
+          # dealing with longer responses like image gen
+          config.flush_interval = 120
+        end
+
+        Observability::ProjectClient.new(project:, observability_client:)
+      end
+    end
+  end
+end

data/lib/nitro_intelligence/observability/prompt.rb

@@ -0,0 +1,58 @@
+module NitroIntelligence
+  module Observability
+    class Prompt
+      attr_reader :name, :type, :prompt, :version, :config, :labels, :tags
+
+      VARIABLE_REGEX = /\{\{([a-zA-Z0-9_]+)\}\}/
+
+      def initialize(name:, type:, prompt:, version:, **extra_args)
+        @name = name
+        @type = type
+        @prompt = prompt
+        @version = version
+        @config = extra_args[:config] || {}
+        @labels = extra_args[:labels] || []
+        @tags = extra_args[:tags] || []
+      end
+
+      # Returns prompt "content" from API with prompt variables replaced
+      # Prompt "content" will either be a string or an array of hashes
+      # based on prompt "type" ("text" or "chat")
+      def compile(**replacements)
+        return replace_variables(@prompt, **replacements) if @type == "text"
+
+        @prompt.map do |message|
+          message[:content] = replace_variables(message[:content], **replacements)
+          message
+        end
+      end
+
+      # Takes provided chat messages and inserts the compiled prompt
+      # into the correct position based on prompt "type" ("text" or "chat")
+      def interpolate(messages:, variables:)
+        if @type == "text"
+          messages.prepend({ role: "system", content: compile(**variables) })
+        elsif @type == "chat"
+          compile(**variables) + messages
+        end
+      end
+
+      def variables
+        messages = @type == "text" ? [@prompt] : @prompt.pluck(:content)
+
+        messages.map do |message|
+          message.scan(VARIABLE_REGEX).flatten.map(&:to_sym)
+        end.flatten
+      end
+
+    private
+
+      def replace_variables(input, **replacements)
+        input.gsub(VARIABLE_REGEX) do |match|
+          key = ::Regexp.last_match(1).to_sym
+          replacements.fetch(key, match)
+        end
+      end
+    end
+  end
+end