langfuse-ruby 0.1.7 → 0.2.0

data/lib/langfuse/client.rb

@@ -6,36 +6,73 @@ require 'faraday/multipart'
 require 'json'
 require 'base64'
 require 'concurrent'
+require 'logger'
+require 'digest'
 
 module Langfuse
   class Client
+    # The ingestion API limits batch payloads to 3.5 MB in total
+    MAX_BATCH_SIZE_BYTES = 3_500_000
+
+    # Allowed format for the tracing environment field
+    ENVIRONMENT_PATTERN = /\A(?!langfuse)[a-z0-9\-_]{1,40}\z/
+
+    # Log device that resolves $stdout at write time so output redirection
+    # (e.g. in tests) keeps working after the logger was created.
+    class StdoutLogDevice
+      def write(message)
+        $stdout.write(message)
+      end
+
+      def close; end
+    end
+
     attr_reader :public_key, :secret_key, :host, :debug, :timeout, :retries, :flush_interval, :auto_flush,
-                :ingestion_mode
+                :ingestion_mode, :environment, :sample_rate, :flush_at, :mask, :logger
 
     def initialize(public_key: nil, secret_key: nil, host: nil, debug: false, timeout: 30, retries: 3,
-                   flush_interval: nil, auto_flush: nil, ingestion_mode: nil)
-      @public_key = public_key || ENV['LANGFUSE_PUBLIC_KEY'] || Langfuse.configuration.public_key
-      @secret_key = secret_key || ENV['LANGFUSE_SECRET_KEY'] || Langfuse.configuration.secret_key
-      @host = host || ENV['LANGFUSE_HOST'] || Langfuse.configuration.host
-      @debug = debug || Langfuse.configuration.debug
-      @timeout = timeout || Langfuse.configuration.timeout
-      @retries = retries || Langfuse.configuration.retries
-      @flush_interval = flush_interval || ENV['LANGFUSE_FLUSH_INTERVAL']&.to_i || Langfuse.configuration.flush_interval
-      @auto_flush = if auto_flush.nil?
-                      ENV['LANGFUSE_AUTO_FLUSH'] == 'false' ? false : Langfuse.configuration.auto_flush
-                    else
-                      auto_flush
-                    end
+                   flush_interval: nil, auto_flush: nil, ingestion_mode: nil, environment: nil,
+                   sample_rate: nil, mask: nil, flush_at: nil, logger: nil, shutdown_on_exit: nil)
+      @public_key = config_value(public_key, 'LANGFUSE_PUBLIC_KEY', :public_key)
+      @secret_key = config_value(secret_key, 'LANGFUSE_SECRET_KEY', :secret_key)
+      @host = host || ENV['LANGFUSE_HOST'] || ENV['LANGFUSE_BASE_URL'] || Langfuse.configuration.host
+      @debug = debug || ENV['LANGFUSE_DEBUG'] == 'true' || Langfuse.configuration.debug
+      @timeout = config_value(timeout, nil, :timeout) { 30 }
+      @retries = config_value(retries, nil, :retries) { 3 }
+      @flush_interval = config_value(flush_interval, 'LANGFUSE_FLUSH_INTERVAL', :flush_interval) { 5 }
+      @flush_at = config_value(flush_at, 'LANGFUSE_FLUSH_AT', :flush_at) { 15 }
+      @auto_flush = resolve_auto_flush(auto_flush)
       @ingestion_mode = resolve_ingestion_mode(ingestion_mode)
+      @logger = logger || Langfuse.configuration.logger || build_default_logger
+      @environment = resolve_environment(environment)
+      @sample_rate = resolve_sample_rate(sample_rate)
+      @mask = resolve_mask(mask)
+      @shutdown_on_exit = shutdown_on_exit.nil? ? Langfuse.configuration.shutdown_on_exit : shutdown_on_exit
+      @shutdown = false
 
       raise AuthenticationError, 'Public key is required' unless @public_key
       raise AuthenticationError, 'Secret key is required' unless @secret_key
 
       @connection = build_connection
       @otel_connection = build_otel_connection if @ingestion_mode == :otel
-      @otel_exporter = OtelExporter.new(connection: @otel_connection, debug: @debug) if @ingestion_mode == :otel
+      @otel_exporter = OtelExporter.new(connection: @otel_connection, debug: @debug, logger: @logger) if @ingestion_mode == :otel
       @event_queue = Concurrent::Array.new
+      @flush_mutex = Mutex.new
+      @flush_condition = ConditionVariable.new
       @flush_thread = start_flush_thread if @auto_flush
+      register_shutdown_hook if @shutdown_on_exit
+    end
+
+    # Generate a trace ID matching the active ingestion mode
+    # (W3C 32-char hex for :otel, UUID for :legacy)
+    def generate_trace_id
+      @ingestion_mode == :otel ? Utils.generate_hex_trace_id : Utils.generate_id
+    end
+
+    # Generate an observation ID matching the active ingestion mode
+    # (W3C 16-char hex for :otel, UUID for :legacy)
+    def generate_observation_id
+      @ingestion_mode == :otel ? Utils.generate_hex_span_id : Utils.generate_id
     end
 
     # Trace operations
@@ -43,7 +80,7 @@ module Langfuse
               input: nil, output: nil, metadata: nil, tags: nil, timestamp: nil, **kwargs)
       Trace.new(
         client: self,
-        id: id || Utils.generate_id,
+        id: id || generate_trace_id,
         name: name,
         user_id: user_id,
         session_id: session_id,
@@ -59,12 +96,13 @@ module Langfuse
     end
 
     # Span operations
-    def span(trace_id:, name: nil, start_time: nil, end_time: nil, input: nil, output: nil,
+    def span(trace_id:, id: nil, name: nil, start_time: nil, end_time: nil, input: nil, output: nil,
              metadata: nil, level: nil, status_message: nil, parent_observation_id: nil,
              version: nil, as_type: nil, **kwargs)
       Span.new(
         client: self,
         trace_id: trace_id,
+        id: id || generate_observation_id,
         name: name,
         start_time: start_time || Utils.current_timestamp,
         end_time: end_time,
@@ -233,13 +271,15 @@ module Langfuse
     end
 
     # Generation operations
-    def generation(trace_id:, name: nil, start_time: nil, end_time: nil, completion_start_time: nil,
+    def generation(trace_id:, id: nil, name: nil, start_time: nil, end_time: nil, completion_start_time: nil,
                    model: nil, model_parameters: nil, input: nil, output: nil, usage: nil,
+                   usage_details: nil, cost_details: nil, prompt: nil,
                    metadata: nil, level: nil, status_message: nil, parent_observation_id: nil,
                    version: nil, **kwargs)
       Generation.new(
         client: self,
         trace_id: trace_id,
+        id: id || generate_observation_id,
         name: name,
         start_time: start_time || Utils.current_timestamp,
         end_time: end_time,
@@ -249,6 +289,9 @@ module Langfuse
         input: input,
         output: output,
         usage: usage,
+        usage_details: usage_details,
+        cost_details: cost_details,
+        prompt: prompt,
         metadata: metadata,
         level: level,
         status_message: status_message,
@@ -259,11 +302,12 @@ module Langfuse
     end
 
     # Event operations
-    def event(trace_id:, name:, start_time: nil, input: nil, output: nil, metadata: nil,
+    def event(trace_id:, name:, id: nil, start_time: nil, input: nil, output: nil, metadata: nil,
               level: nil, status_message: nil, parent_observation_id: nil, version: nil, **kwargs)
       Event.new(
         client: self,
         trace_id: trace_id,
+        id: id || generate_observation_id,
         name: name,
         start_time: start_time,
         input: input,
@@ -291,18 +335,18 @@ module Langfuse
       params[:version] = version if version
       params[:label] = label if label
 
-      puts "Making request to: #{@host}#{path} with params: #{params}" if @debug
+      @logger.debug("Making request to: #{@host}#{path} with params: #{params}")
 
       response = get(path, params)
 
-      puts "Response status: #{response.status}" if @debug
-      puts "Response headers: #{response.headers}" if @debug
-      puts "Response body type: #{response.body.class}" if @debug
+      @logger.debug("Response status: #{response.status}")
+      @logger.debug("Response headers: #{response.headers}")
+      @logger.debug("Response body type: #{response.body.class}")
 
       # Check if response body is a string (HTML) instead of parsed JSON
       if response.body.is_a?(String) && response.body.include?('<!DOCTYPE html>')
-        puts 'Received HTML response instead of JSON:' if @debug
-        puts response.body[0..200] if @debug
+        @logger.debug('Received HTML response instead of JSON:')
+        @logger.debug(response.body[0..200])
         raise APIError,
               'Received HTML response instead of JSON. This usually indicates a 404 error or incorrect API endpoint.'
       end
@@ -330,20 +374,35 @@ module Langfuse
     end
 
     # Score/Evaluation operations
-    def score(name:, value:, trace_id: nil, observation_id: nil, data_type: nil, comment: nil, **kwargs)
+    # Scores can target a trace, an observation (trace_id + observation_id),
+    # a session (session_id) or a dataset run (dataset_run_id).
+    def score(name:, value:, trace_id: nil, observation_id: nil, session_id: nil, dataset_run_id: nil,
+              id: nil, data_type: nil, comment: nil, metadata: nil, config_id: nil, queue_id: nil,
+              environment: nil, **kwargs)
       data = {
+        id: id,
+        trace_id: trace_id,
+        observation_id: observation_id,
+        session_id: session_id,
+        dataset_run_id: dataset_run_id,
         name: name,
         value: value,
         data_type: data_type,
         comment: comment,
+        metadata: metadata,
+        config_id: config_id,
+        queue_id: queue_id,
+        environment: environment,
         **kwargs
-      }
+      }.compact
 
-      data[:trace_id] = trace_id if trace_id
-      data[:observation_id] = observation_id if observation_id
+      if trace_id.nil? && observation_id.nil? && session_id.nil? && dataset_run_id.nil?
+        @logger.warn('Langfuse score should reference a trace_id, observation_id, session_id or dataset_run_id')
+      end
 
       enqueue_event('score-create', data)
     end
+    alias create_score score
 
     # Event queue management
     def enqueue_event(type, body)
@@ -357,15 +416,21 @@ module Langfuse
       ]
 
       unless valid_types.include?(type)
-        puts "Warning: Invalid event type '#{type}'. Skipping event." if @debug
+        @logger.debug("Warning: Invalid event type '#{type}'. Skipping event.")
         return
       end
 
+      prepared_body = Utils.prepare_event_body(body)
+      inject_default_environment(prepared_body)
+      apply_mask(prepared_body)
+
+      return unless sampled_event?(type, prepared_body)
+
       event = {
         id: Utils.generate_id,
         type: type,
         timestamp: Utils.current_timestamp,
-        body: Utils.deep_stringify_keys(body)
+        body: prepared_body
       }
 
       if type == 'trace-update'
@@ -381,20 +446,22 @@ module Langfuse
             # 更新现有的 trace-create 事件
             @event_queue[existing_event_index][:body].merge!(event[:body])
             @event_queue[existing_event_index][:timestamp] = event[:timestamp]
-            puts "Updated existing trace-create event for trace_id: #{trace_id}" if @debug
+            @logger.debug("Updated existing trace-create event for trace_id: #{trace_id}")
           else
             # 如果没找到对应的 trace-create 事件，将 trace-update 转换为 trace-create
             event[:type] = 'trace-create'
             @event_queue << event
-            puts "Converted trace-update to trace-create for trace_id: #{trace_id}" if @debug
+            @logger.debug("Converted trace-update to trace-create for trace_id: #{trace_id}")
           end
-        elsif @debug
-          puts 'Warning: trace-update event missing trace_id, skipping'
+        else
+          @logger.debug('Warning: trace-update event missing trace_id, skipping')
         end
       else
         @event_queue << event
       end
-      puts "Enqueued event: #{type}" if @debug
+      @logger.debug("Enqueued event: #{type}")
+
+      request_flush if @auto_flush && @event_queue.length >= @flush_at
     end
 
     def flush
@@ -407,31 +474,155 @@ module Langfuse
     end
 
     def shutdown
+      return if @shutdown
+
+      @shutdown = true
       @flush_thread&.kill if @auto_flush
       flush unless @event_queue.empty?
     end
 
     private
 
+    def build_default_logger
+      logger = Logger.new(StdoutLogDevice.new)
+      logger.level = @debug ? Logger::DEBUG : Logger::WARN
+      logger.progname = 'langfuse'
+      logger.formatter = proc do |severity, _time, progname, msg|
+        "#{severity} -- #{progname}: #{msg}\n"
+      end
+      logger
+    end
+
+    # Resolve a config value with precedence: explicit arg > env var > config attr > block default
+    def config_value(explicit, env_key, config_attr)
+      return explicit if explicit
+
+      if env_key
+        env_val = ENV.fetch(env_key, nil)
+        return env_val.to_i if env_val && %i[flush_interval flush_at timeout retries].include?(config_attr)
+        return env_val if env_val
+      end
+
+      Langfuse.configuration.send(config_attr) || (yield if block_given?)
+    end
+
+    def resolve_auto_flush(auto_flush)
+      if auto_flush.nil?
+        ENV['LANGFUSE_AUTO_FLUSH'] == 'false' ? false : Langfuse.configuration.auto_flush
+      else
+        auto_flush
+      end
+    end
+
+    def resolve_environment(explicit_environment)
+      environment = explicit_environment || ENV['LANGFUSE_TRACING_ENVIRONMENT'] || Langfuse.configuration.environment
+      return nil if environment.nil? || environment.to_s.empty?
+
+      environment = environment.to_s
+      unless environment.match?(ENVIRONMENT_PATTERN)
+        @logger.warn("Invalid Langfuse environment '#{environment}'. It must match #{ENVIRONMENT_PATTERN.inspect}. " \
+                     'Events may be rejected by the server.')
+      end
+      environment
+    end
+
+    def resolve_sample_rate(explicit_sample_rate)
+      rate = explicit_sample_rate || ENV['LANGFUSE_SAMPLE_RATE']&.to_f || Langfuse.configuration.sample_rate
+      return nil if rate.nil?
+
+      rate = rate.to_f
+      unless rate.between?(0.0, 1.0)
+        @logger.warn("Invalid Langfuse sample_rate #{rate}, expected 0.0..1.0. Disabling sampling.")
+        return nil
+      end
+      rate
+    end
+
+    def resolve_mask(explicit_mask)
+      mask = explicit_mask || Langfuse.configuration.mask
+      return nil if mask.nil?
+
+      unless mask.respond_to?(:call)
+        @logger.warn('Langfuse mask must respond to #call. Ignoring mask.')
+        return nil
+      end
+      mask
+    end
+
+    def register_shutdown_hook
+      at_exit do
+        shutdown
+      rescue StandardError => e
+        @logger.debug("Langfuse shutdown on exit failed: #{e.message}")
+      end
+    end
+
+    def inject_default_environment(body)
+      return unless @environment
+      return if body.key?('environment')
+
+      body['environment'] = @environment
+    end
+
+    def apply_mask(body)
+      return unless @mask
+
+      %w[input output metadata].each do |field|
+        next unless body.key?(field) && !body[field].nil?
+
+        body[field] = begin
+          @mask.call(body[field])
+        rescue StandardError => e
+          @logger.error("Langfuse mask function failed: #{e.message}")
+          '<masked due to failed mask function>'
+        end
+      end
+    end
+
+    # Deterministic trace-based sampling: all events of a trace share the same decision.
+    def sampled_event?(type, body)
+      return true unless @sample_rate
+
+      trace_id = %w[trace-create trace-update].include?(type) ? body['id'] : body['traceId']
+      return true if trace_id.nil?
+
+      return true if trace_sampled?(trace_id)
+
+      @logger.debug("Dropping event for trace #{trace_id} due to sampling (rate: #{@sample_rate})")
+      false
+    end
+
+    def trace_sampled?(trace_id)
+      return true if @sample_rate >= 1.0
+      return false if @sample_rate <= 0.0
+
+      normalized = Digest::SHA256.hexdigest(trace_id.to_s)[0, 8].to_i(16).to_f / 0xffffffff
+      normalized < @sample_rate
+    end
+
+    def request_flush
+      @flush_mutex.synchronize { @flush_condition.signal }
+    end
+
     def debug_event_data(events)
       return unless @debug
 
-      puts "\n=== Event Data Debug Information ==="
+      @logger.debug('=== Event Data Debug Information ===')
       events.each_with_index do |event, index|
-        puts "Event #{index + 1}:"
-        puts "  ID: #{event[:id]}"
-        puts "  Type: #{event[:type]}"
-        puts "  Timestamp: #{event[:timestamp]}"
-        puts "  Body keys: #{event[:body]&.keys || 'nil'}"
+        @logger.debug("Event #{index + 1}:")
+        @logger.debug("  ID: #{event[:id]}")
+        @logger.debug("  Type: #{event[:type]}")
+        @logger.debug("  Timestamp: #{event[:timestamp]}")
+        @logger.debug("  Body keys: #{event[:body]&.keys || 'nil'}")
 
         # 检查常见的问题
-        puts '  ⚠️  WARNING: Empty or nil type!' if event[:type].nil? || event[:type].to_s.empty?
+        @logger.debug('  ⚠️  WARNING: Empty or nil type!') if event[:type].nil? || event[:type].to_s.empty?
 
-        puts '  ⚠️  WARNING: Empty body!' if event[:body].nil?
+        @logger.debug('  ⚠️  WARNING: Empty body!') if event[:body].nil?
 
-        puts '  ---'
+        @logger.debug('  ---')
       end
-      puts "=== End Debug Information ===\n"
+      @logger.debug('=== End Debug Information ===')
     end
 
     def send_batch(events)
@@ -441,10 +632,10 @@ module Langfuse
       # 验证事件数据
       valid_events = events.select do |event|
         if event[:type].nil? || event[:type].to_s.empty?
-          puts "Warning: Event with empty type detected, skipping: #{event[:id]}" if @debug
+          @logger.debug("Warning: Event with empty type detected, skipping: #{event[:id]}")
           false
         elsif event[:body].nil?
-          puts "Warning: Event with empty body detected, skipping: #{event[:id]}" if @debug
+          @logger.debug("Warning: Event with empty body detected, skipping: #{event[:id]}")
           false
         else
           true
@@ -452,7 +643,7 @@ module Langfuse
       end
 
       if valid_events.empty?
-        puts 'No valid events to send' if @debug
+        @logger.debug('No valid events to send')
         return
       end
 
@@ -464,32 +655,111 @@ module Langfuse
     end
 
     def send_batch_legacy(valid_events)
-      batch_data = build_batch_data(valid_events)
-      puts "Sending batch data: #{batch_data}" if @debug
-
-      begin
-        response = post('/api/public/ingestion', batch_data)
-        puts "Flushed #{valid_events.length} events (legacy)" if @debug
-        response
-      rescue StandardError => e
-        puts "Failed to flush events: #{e.message}" if @debug
-        valid_events.each { |event| @event_queue << event }
-        raise
+      chunks = chunk_events(valid_events)
+      response = nil
+
+      chunks.each_with_index do |chunk, index|
+        batch_data = build_batch_data(chunk)
+        @logger.debug("Sending batch data: #{batch_data}")
+
+        begin
+          response = post('/api/public/ingestion', batch_data)
+          log_ingestion_errors(response)
+          @logger.debug("Flushed #{chunk.length} events (legacy)")
+        rescue StandardError => e
+          @logger.debug("Failed to flush events: #{e.message}")
+          chunks[index..].each { |failed_chunk| failed_chunk.each { |event| @event_queue << event } }
+          raise
+        end
       end
+
+      response
     end
 
     def send_batch_otel(valid_events)
-      puts "Sending #{valid_events.length} events via OTEL" if @debug
+      score_events, otel_events = valid_events.partition { |event| event[:type] == 'score-create' }
+
+      response = nil
+
+      unless otel_events.empty?
+        @logger.debug("Sending #{otel_events.length} events via OTEL")
+
+        begin
+          response = @otel_exporter.export(otel_events)
+          handle_response(response)
+          @logger.debug("Flushed #{otel_events.length} events (otel)")
+        rescue StandardError => e
+          @logger.debug("Failed to flush OTEL events: #{e.message}")
+          # Re-queue both OTEL and score events — scores were already drained
+          # from the queue by flush and would otherwise be permanently lost.
+          otel_events.each { |event| @event_queue << event }
+          score_events.each { |event| @event_queue << event }
+          raise
+        end
+      end
 
-      begin
-        response = @otel_exporter.export(valid_events)
-        handle_response(response)
-        puts "Flushed #{valid_events.length} events (otel)" if @debug
-        response
-      rescue StandardError => e
-        puts "Failed to flush OTEL events: #{e.message}" if @debug
-        valid_events.each { |event| @event_queue << event }
-        raise
+      # Scores are not part of the OTLP trace mapping; they always go through
+      # the ingestion API. IDs are normalized to match the OTel-derived IDs.
+      unless score_events.empty?
+        score_events.each { |event| normalize_otel_score_event(event) }
+        response = send_batch_legacy(score_events)
+      end
+
+      response
+    end
+
+    # Align score references with the OTel-derived trace/span IDs so scores
+    # attach to the correct entities when ingesting via the OTel endpoint.
+    def normalize_otel_score_event(event)
+      body = event[:body]
+      return unless body.is_a?(Hash)
+
+      body['traceId'] = OtelExporter.to_otel_trace_id(body['traceId']) if body['traceId']
+      body['observationId'] = OtelExporter.to_otel_span_id(body['observationId']) if body['observationId']
+    end
+
+    # Split events into chunks that respect the ingestion API batch size limit.
+    def chunk_events(events)
+      chunks = [[]]
+      current_size = 0
+
+      events.each do |event|
+        event_size = estimated_event_size(event)
+
+        if event_size > MAX_BATCH_SIZE_BYTES
+          @logger.warn("Langfuse event #{event[:id]} exceeds the maximum batch size of #{MAX_BATCH_SIZE_BYTES} bytes and was dropped")
+          next
+        end
+
+        if current_size + event_size > MAX_BATCH_SIZE_BYTES && !chunks.last.empty?
+          chunks << []
+          current_size = 0
+        end
+
+        chunks.last << event
+        current_size += event_size
+      end
+
+      chunks.reject(&:empty?)
+    end
+
+    def estimated_event_size(event)
+      JSON.generate(event).bytesize
+    rescue StandardError
+      1024
+    end
+
+    # The ingestion API responds with 207 and per-event successes/errors.
+    def log_ingestion_errors(response)
+      body = response.respond_to?(:body) ? response.body : nil
+      return unless body.is_a?(Hash)
+
+      errors = body['errors']
+      return unless errors.is_a?(Array) && errors.any?
+
+      errors.each do |error|
+        @logger.warn("Langfuse ingestion partial failure (status #{error['status']}): " \
+                     "event #{error['id']} - #{error['message']}")
       end
     end
 
@@ -509,11 +779,12 @@ module Langfuse
 
       Thread.new do
         loop do
-          sleep(@flush_interval) # Configurable flush interval
+          # Wait for the flush interval or an early wake-up (flush_at threshold)
+          @flush_mutex.synchronize { @flush_condition.wait(@flush_mutex, @flush_interval) }
           begin
             flush unless @event_queue.empty?
           rescue StandardError => e
-            puts "Error in flush thread: #{e.message}" if @debug
+            @logger.debug("Error in flush thread: #{e.message}")
           end
         end
       end
@@ -522,7 +793,7 @@ module Langfuse
     def resolve_ingestion_mode(explicit_mode)
       return explicit_mode.to_sym if explicit_mode
 
-      env_mode = ENV['LANGFUSE_INGESTION_MODE']
+      env_mode = ENV.fetch('LANGFUSE_INGESTION_MODE', nil)
       return env_mode.to_sym if env_mode && !env_mode.empty?
 
       Langfuse.configuration.ingestion_mode || :legacy
@@ -614,7 +885,7 @@ module Langfuse
     end
 
     def handle_response(response)
-      puts "Handling response with status: #{response.status}" if @debug
+      @logger.debug("Handling response with status: #{response.status}")
 
       case response.status
       when 200..299