nitro_intelligence 0.0.1

data/lib/nitro_intelligence/langfuse_extension.rb

@@ -0,0 +1,126 @@
+# This is an adaptation of https://github.com/simplepractice/langfuse-rb/blob/main/lib/langfuse.rb
+# that lets users set custom trace IDs and create multiple Langfuse clients in the same process.
+#
+# 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.
+#
+module NitroIntelligence
+  class LangfuseExtension
+    attr_reader :config
+
+    def initialize
+      config = Langfuse::Config.new
+      yield(config) if block_given?
+
+      @config = config
+      @client = Langfuse::Client.new(config)
+      @tracer_provider = LangfuseTracerProvider.new(config)
+    end
+
+    def shutdown(timeout: 30)
+      @client.shutdown
+      @tracer_provider.shutdown(timeout:)
+    end
+
+    def create_score(name:, value:, id: nil, trace_id: nil, session_id: nil, observation_id: nil, comment: nil, # rubocop:disable Metrics/ParameterLists
+                     metadata: nil, environment: nil, data_type: :numeric, dataset_run_id: nil, config_id: nil)
+      @client.create_score(
+        name:,
+        value:,
+        id:,
+        trace_id:,
+        session_id:,
+        observation_id:,
+        comment:,
+        metadata:,
+        environment:,
+        data_type:,
+        dataset_run_id:,
+        config_id:
+      )
+    end
+
+    def start_observation(name, attrs = {}, as_type: :span, parent_span_context: nil, start_time: nil,
+                          skip_validation: false)
+      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
+
+      otel_tracer = @tracer_provider.tracer
+      otel_span = Langfuse.send(:create_otel_span,
+                                name:,
+                                start_time:,
+                                parent_span_context:,
+                                otel_tracer:)
+
+      # Serialize attributes
+      # Only set attributes if span is still recording (should always be true here, but guard for safety)
+      if otel_span.recording?
+        otel_attrs = Langfuse::OtelAttributes.create_observation_attributes(type_str, attrs.to_h, mask: nil)
+        otel_attrs.each { |key, value| otel_span.set_attribute(key, value) }
+      end
+
+      # Wrap in appropriate class (attributes already set on span above — pass nil to avoid double-masking)
+      observation = Langfuse.send(:wrap_otel_span, otel_span, type_str, otel_tracer)
+
+      # Events auto-end immediately when created
+      observation.end if type_str == Langfuse::OBSERVATION_TYPES[:event]
+
+      observation
+    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
+
+        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:
+      )
+    end
+  end
+end

data/lib/nitro_intelligence/langfuse_tracer_provider.rb

@@ -0,0 +1,59 @@
+# This is an adaptation of https://github.com/simplepractice/langfuse-rb/blob/main/lib/langfuse/otel_setup.rb
+# that lets users create multiple Langfuse clients in the same process.
+# See also components/nitro_intelligence/lib/nitro_intelligence/langfuse_extension.rb
+#
+# The content of this file should eventually make its way upstream.
+#
+module NitroIntelligence
+  class LangfuseTracerProvider
+    def initialize(config)
+      @tracer_provider = create_tracer_provider(config)
+    end
+
+    def tracer
+      @tracer_provider.tracer("langfuse-rb", Langfuse::VERSION)
+    end
+
+    def shutdown(timeout: 30)
+      @tracer_provider.shutdown(timeout:)
+    end
+
+  private
+
+    def create_tracer_provider(config)
+      raise ArgumentError, "training_async must be true" unless config.tracing_async
+
+      exporter = OpenTelemetry::Exporter::OTLP::Exporter.new(
+        endpoint: "#{config.base_url}/api/public/otel/v1/traces",
+        headers: build_headers(config.public_key, config.secret_key),
+        compression: "gzip"
+      )
+
+      processor = OpenTelemetry::SDK::Trace::Export::BatchSpanProcessor.new(
+        exporter,
+        max_queue_size: config.batch_size * 2,
+        schedule_delay: config.flush_interval * 1000,
+        max_export_batch_size: config.batch_size
+      )
+
+      tracer_provider = OpenTelemetry::SDK::Trace::TracerProvider.new
+      tracer_provider.add_span_processor(processor)
+      span_processor = Langfuse::SpanProcessor.new(config:)
+      tracer_provider.add_span_processor(span_processor)
+
+      OpenTelemetry.propagation = OpenTelemetry::Trace::Propagation::TraceContext::TextMapPropagator.new
+
+      config.logger.info("Langfuse tracing initialized with OpenTelemetry (async mode)")
+
+      tracer_provider
+    end
+
+    def build_headers(public_key, secret_key)
+      credentials = "#{public_key}:#{secret_key}"
+      encoded = Base64.strict_encode64(credentials)
+      {
+        "Authorization" => "Basic #{encoded}",
+      }
+    end
+  end
+end

data/lib/nitro_intelligence/media/image.rb

@@ -0,0 +1,36 @@
+require "base64"
+require "mini_magick"
+
+require "nitro_intelligence/media/media"
+
+module NitroIntelligence
+  class Image < Media
+    attr_reader :height, :width
+
+    def self.from_base64(base64_string)
+      # Strip data_uri from string
+      base64_string = base64_string.sub(/^data:[^;]*;base64,/, "")
+      byte_string = Base64.strict_decode64(base64_string)
+
+      new(byte_string)
+    end
+
+    def initialize(file)
+      super
+
+      image = MiniMagick::Image.read(StringIO.new(file))
+
+      @mime_type = image.mime_type
+      @height = image.height
+      @width = image.width
+
+      parse_mime_type
+    end
+
+  private
+
+    def parse_mime_type
+      @file_type, @file_extension = @mime_type.split("/", 2)
+    end
+  end
+end

data/lib/nitro_intelligence/media/image_generation.rb

@@ -0,0 +1,130 @@
+require "base64"
+require "digest"
+require "mini_magick"
+require "stringio"
+require "time"
+
+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
+
+      attr_accessor :aspect_ratio, :model, :resolution
+
+      def initialize
+        @aspect_ratio = DEFAULT_ASPECT_RATIO
+        @model = NitroIntelligence.model_catalog.default_image_model.name
+        @resolution = DEFAULT_RESOLUTION
+      end
+    end
+
+    attr_reader :byte_string, :config, :file_type, :file_extension, :generated_image, :messages, :reference_images,
+                :target_image
+    attr_accessor :trace_id
+
+    def initialize(
+      message: "",
+      target_image: nil,
+      reference_images: []
+    )
+      @config = Config.new
+      @generated_image = nil
+      @model = NitroIntelligence.model_catalog.lookup_by_name(@config.model)
+      @reference_images = reference_images.map { |img| Image.new(img) }
+      @trace_id = nil
+
+      # Overrides
+      yield(@config) if block_given?
+      validate_config!
+
+      if target_image
+        @target_image = Image.new(target_image)
+        configure_aspect_ratio
+      end
+
+      build_messages(message, @target_image, @reference_images)
+    end
+
+    def files
+      [target_image, reference_images, generated_image].flatten.compact
+    end
+
+    def parse_file(chat_completion)
+      base64_string = chat_completion.choices.first&.message.to_h.fetch(:images, {})&.first&.dig(:image_url, :url)
+
+      return unless base64_string
+
+      @generated_image = Image.from_base64(base64_string)
+      @generated_image.direction = "output"
+      @generated_image
+    end
+
+  private
+
+    def build_messages(message, target_image, reference_images)
+      messages = [{ role: "user", content: [] }]
+
+      messages.first[:content].append({ type: "text", text: message }) if message.present?
+
+      if target_image
+        messages.first[:content].append(image_message(mime_type: target_image.mime_type, base64: target_image.base64))
+      end
+
+      reference_images.each do |img|
+        messages.first[:content].append(image_message(mime_type: img.mime_type, base64: img.base64))
+      end
+
+      @messages = messages
+    end
+
+    def calculate_aspect_ratios
+      @model.aspect_ratios.index_by { |x| x.split(":").map(&:to_f).reduce(:/) }
+    end
+
+    def closest_aspect_ratio(width, height)
+      actual_ratio = width.to_f / height
+      calculated_aspect_ratios = calculate_aspect_ratios
+      best_match = calculated_aspect_ratios.keys.min_by { |ratio_val| (ratio_val - actual_ratio).abs }
+      calculated_aspect_ratios[best_match]
+    end
+
+    def configure_aspect_ratio
+      @config.aspect_ratio = closest_aspect_ratio(@target_image.width, @target_image.height)
+    end
+
+    def image_message(url: nil, mime_type: nil, base64: nil)
+      url = "data:#{mime_type};base64,#{base64}" if url.nil?
+
+      {
+        type: "image_url",
+        image_url: {
+          url:,
+        },
+      }
+    end
+
+    def validate_config!
+      # Check model supported
+      @model = NitroIntelligence.model_catalog.lookup_by_name(@config.model)
+      raise ArgumentError, "Unsupported model: '#{@config.model}'" unless @model
+
+      # Check aspect_ratio supported
+      unless @model.aspect_ratios.include?(@config.aspect_ratio)
+        raise ArgumentError,
+              "Unsupported aspect ratio: '#{@config.aspect_ratio}'. " \
+              "Supported ratios for #{@config.model} are: #{@model.aspect_ratios}"
+      end
+
+      # Check resolution supported
+      return if @model.resolutions.include?(@config.resolution)
+
+      raise ArgumentError,
+            "Unsupported resolution: '#{@config.resolution}'. " \
+            "Supported resolutions for #{@config.model} are: #{@model.resolutions}"
+    end
+  end
+end

data/lib/nitro_intelligence/media/media.rb

@@ -0,0 +1,19 @@
+require "base64"
+
+module NitroIntelligence
+  class Media
+    attr_accessor :direction, :reference_id
+    attr_reader :base64, :byte_string, :file_extension, :file_type, :mime_type
+
+    # Input should be byte string. e.g. File.binread('file.ext')
+    def initialize(file)
+      @base64 = Base64.strict_encode64(file)
+      @byte_string = file
+      @direction = "input"
+      @file_extension = nil
+      @file_type = nil
+      @mime_type = nil
+      @reference_id = nil
+    end
+  end
+end

data/lib/nitro_intelligence/media/upload_handler.rb

@@ -0,0 +1,135 @@
+require "base64"
+require "digest"
+require "httparty"
+require "time"
+
+module NitroIntelligence
+  class UploadHandler
+    def initialize(observability_host, observability_auth_token)
+      @observability_host = observability_host
+      @observability_auth_token = observability_auth_token
+      @uploaded_media = []
+    end
+
+    def replace_base64_with_media_references(payload)
+      # Make it easier to lookup message style image_urls
+      media_lookup = @uploaded_media.index_by { |image| "data:#{image.mime_type};base64,#{image.base64}" }
+
+      replace_base64_image_url = ->(node) do
+        case node
+        when Hash
+          url_key = if node.key?(:url)
+                      :url
+                    else
+                      (node.key?("url") ? "url" : nil)
+                    end
+
+          if url_key
+            url_value = node[url_key]
+
+            # Replace base64 strings if they match with our uploaded media
+            if (media = media_lookup[url_value])
+              # Overwrite base64 string with Langfuse media ref
+              # This *should* be rendering inline in the gui
+              # https://github.com/langfuse/langfuse/issues/4555
+              # https://github.com/langfuse/langfuse/issues/5030
+              node[url_key] = "@@@langfuseMedia:type=#{media.mime_type}|id=#{media.reference_id}|source=bytes@@@"
+            # Sometimes models can generate unwanted images that will not have a reference
+            # these untracked base64 strings can easily push 4.5mb Langfuse limit
+            elsif url_value.is_a?(String) && url_value.start_with?("data:")
+              node[url_key] = "[Discarded media]"
+            end
+          end
+
+          node.each_value { |val| replace_base64_image_url.call(val) }
+
+        when Array
+          node.each { |val| replace_base64_image_url.call(val) }
+        end
+      end
+
+      replace_base64_image_url.call(payload)
+
+      payload
+    end
+
+    def upload(trace_id, upload_queue: Queue.new)
+      until upload_queue.empty?
+        media = upload_queue.pop
+
+        content_length = media.byte_string.bytesize
+        content_sha256 = Base64.strict_encode64(Digest::SHA256.digest(media.byte_string))
+
+        # returns {"mediaId" -> "", "uploadUrl" => ""}
+        upload_url_response = get_upload_url({
+                                               traceId: trace_id,
+                                               contentType: media.mime_type,
+                                               contentLength: content_length,
+                                               sha256Hash: content_sha256,
+                                               field: media.direction,
+                                             })
+
+        # NOTE: uploadUrl is None if the file is stored in Langfuse already as then there is no need to upload it again.
+        upload_response = upload_media(
+          upload_url_response["mediaId"],
+          upload_url_response["uploadUrl"],
+          media.mime_type,
+          content_sha256,
+          media.byte_string
+        )
+
+        associate_media(upload_url_response["mediaId"], upload_response) if upload_response.present?
+
+        media.reference_id = upload_url_response["mediaId"]
+        @uploaded_media.append(media)
+      end
+
+      @uploaded_media
+    end
+
+  private
+
+    def get_upload_url(request_body)
+      HTTParty.post(
+        "#{@observability_host}/api/public/media",
+        body: request_body.to_json,
+        headers: {
+          "Content-Type" => "application/json",
+          "Authorization" => "Basic #{@observability_auth_token}",
+        }
+      )
+    end
+
+    def associate_media(media_id, upload_response)
+      request_body = {
+        uploadedAt: Time.now.utc.iso8601(6),
+        uploadHttpStatus: upload_response.code,
+        uploadHttpError: upload_response.code == 200 ? nil : upload_response.body,
+      }
+
+      HTTParty.patch(
+        "#{@observability_host}/api/public/media/#{media_id}",
+        body: request_body.to_json,
+        headers: {
+          "Content-Type" => "application/json",
+          "Authorization" => "Basic #{@observability_auth_token}",
+        }
+      )
+    end
+
+    def upload_media(media_id, upload_url, content_type, content_sha256, content_bytes)
+      if media_id.present? && upload_url.present?
+        return HTTParty.put(
+          upload_url,
+          headers: {
+            "Content-Type" => content_type,
+            "x-amz-checksum-sha256" => content_sha256,
+          },
+          body: content_bytes
+        )
+      end
+
+      nil
+    end
+  end
+end

data/lib/nitro_intelligence/models/model.rb

@@ -0,0 +1,22 @@
+module NitroIntelligence
+  class Model
+    attr_reader :name, :omit_output_fields
+
+    def initialize(name:, omit_output_fields: [], **)
+      @name = name
+      @omit_output_fields = omit_output_fields.map { |field| field.split(".").map(&:to_sym) }
+    end
+  end
+
+  class TextModel < Model; end
+
+  class ImageModel < Model
+    attr_reader :aspect_ratios, :resolutions
+
+    def initialize(name:, omit_output_fields: [], aspect_ratios: [], resolutions: [], **)
+      super
+      @aspect_ratios = aspect_ratios
+      @resolutions = resolutions
+    end
+  end
+end

data/lib/nitro_intelligence/models/model_catalog.rb

@@ -0,0 +1,17 @@
+require "nitro_intelligence/models/model_factory"
+
+module NitroIntelligence
+  class ModelCatalog
+    attr_reader :models, :default_text_model, :default_image_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_image_model = lookup_by_name(model_config[:default_image_model])
+    end
+
+    def lookup_by_name(name)
+      @models.find { |model| model.name == name }
+    end
+  end
+end

data/lib/nitro_intelligence/models/model_factory.rb

@@ -0,0 +1,19 @@
+require "nitro_intelligence/models/model"
+
+module NitroIntelligence
+  class ModelFactory
+    def self.build(model_metadata)
+      model_metadata = model_metadata.symbolize_keys
+
+      if image_model?(model_metadata)
+        ImageModel.new(**model_metadata)
+      else
+        TextModel.new(**model_metadata)
+      end
+    end
+
+    def self.image_model?(model_metadata)
+      model_metadata.key?(:aspect_ratios) || model_metadata.key?(:resolutions)
+    end
+  end
+end

data/lib/nitro_intelligence/null_cache.rb

@@ -0,0 +1,20 @@
+module NitroIntelligence
+  # Default no-operation cache implementation used
+  # as a sensible default for initializing NitroIntelligence
+  class NullCache
+    def read(_key) = nil
+    def write(_key, _value, **_options) = true
+    def delete(_key) = true
+
+    def fetch(key, **)
+      value = read(key)
+      return value unless value.nil?
+
+      return unless block_given?
+
+      computed = yield
+      write(key, computed, **)
+      computed
+    end
+  end
+end

data/lib/nitro_intelligence/prompt/prompt.rb

@@ -0,0 +1,56 @@
+module NitroIntelligence
+  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