desiru 0.1.1 → 0.2.0

data/lib/desiru/core/prediction.rb

@@ -0,0 +1,108 @@
+# frozen_string_literal: true
+
+module Desiru
+  module Core
+    class Prediction
+      attr_reader :completions, :example
+
+      def initialize(example = nil, **kwargs)
+        @example = example || Example.new
+        @completions = {}
+        @metadata = {}
+
+        kwargs.each do |key, value|
+          if key == :completions
+            @completions = value
+          elsif key == :metadata
+            @metadata = value
+          else
+            @completions[key] = value
+          end
+        end
+      end
+
+      def [](key)
+        if @completions.key?(key)
+          @completions[key]
+        elsif @example
+          # First check the raw data in the example
+          @example[key] ||
+            # Then check the inputs and labels
+            @example.inputs[key] ||
+            @example.labels[key]
+        end
+      end
+
+      def []=(key, value)
+        @completions[key] = value
+      end
+
+      def get(key, default = nil)
+        @completions.fetch(key) { @example[key] || default }
+      end
+
+      def keys
+        (@completions.keys + (@example&.keys || [])).uniq
+      end
+
+      def values
+        keys.map { |k| self[k] }
+      end
+
+      def to_h
+        result = @example ? @example.to_h : {}
+        result.merge(@completions)
+      end
+
+      def to_example
+        Example.new(**to_h)
+      end
+
+      def metadata
+        @metadata.dup
+      end
+
+      def set_metadata(key, value)
+        @metadata[key] = value
+      end
+
+      def method_missing(method_name, *args, &)
+        if method_name.to_s.end_with?('=')
+          key = method_name.to_s.chop.to_sym
+          self[key] = args.first
+        elsif @completions.key?(method_name)
+          @completions[method_name]
+        elsif @example.respond_to?(method_name)
+          @example.send(method_name, *args, &)
+        else
+          super
+        end
+      end
+
+      def respond_to_missing?(method_name, include_private = false)
+        method_name.to_s.end_with?('=') ||
+          @completions.key?(method_name) ||
+          @example.respond_to?(method_name) ||
+          super
+      end
+
+      def ==(other)
+        return false unless other.is_a?(self.class)
+
+        @completions == other.completions &&
+          @example == other.example &&
+          @metadata == other.instance_variable_get(:@metadata)
+      end
+
+      def inspect
+        "#<#{self.class.name} completions=#{@completions.inspect} " \
+          "example=#{@example.inspect} metadata=#{@metadata.inspect}>"
+      end
+
+      # Class method to create a Prediction from an Example
+      def self.from_example(example)
+        new(example)
+      end
+    end
+  end
+end

data/lib/desiru/core/trace.rb

@@ -0,0 +1,330 @@
+# frozen_string_literal: true
+
+module Desiru
+  module Core
+    class Trace
+      attr_reader :module_name, :signature, :inputs, :outputs, :metadata, :timestamp
+
+      def initialize(module_name:, signature:, inputs: {}, outputs: {}, metadata: {})
+        @module_name = module_name
+        @signature = signature
+        @inputs = inputs
+        @outputs = outputs
+        @metadata = metadata
+        @timestamp = Time.now
+      end
+
+      def to_h
+        {
+          module_name: @module_name,
+          signature: @signature.is_a?(String) ? @signature : @signature.to_h,
+          inputs: @inputs,
+          outputs: @outputs,
+          metadata: @metadata,
+          timestamp: @timestamp
+        }
+      end
+
+      def to_example
+        Example.new(**@inputs, **@outputs)
+      end
+
+      def success?
+        @metadata[:success] != false
+      end
+
+      def error?
+        @metadata.key?(:error)
+      end
+
+      def error
+        @metadata[:error]
+      end
+
+      def duration
+        @metadata[:duration]
+      end
+
+      def duration_ms
+        return 0 unless duration
+
+        (duration * 1000).to_f
+      end
+
+      def ==(other)
+        return false unless other.is_a?(self.class)
+
+        @module_name == other.module_name &&
+          @signature == other.signature &&
+          @inputs == other.inputs &&
+          @outputs == other.outputs &&
+          @metadata == other.metadata
+      end
+    end
+
+    class TraceCollector
+      attr_reader :traces
+
+      def initialize
+        @traces = []
+        @enabled = true
+        @filters = []
+      end
+
+      def collect(trace)
+        return unless @enabled && trace.is_a?(Trace)
+        return if @filters.any? { |filter| !filter.call(trace) }
+
+        @traces << trace
+      end
+
+      def add_filter(&block)
+        @filters << block if block_given?
+      end
+
+      def clear_filters
+        @filters.clear
+      end
+
+      def enable
+        @enabled = true
+      end
+
+      def disable
+        @enabled = false
+      end
+
+      def enabled?
+        @enabled
+      end
+
+      def clear
+        @traces.clear
+      end
+
+      def size
+        @traces.size
+      end
+
+      def empty?
+        @traces.empty?
+      end
+
+      def recent(count = 10)
+        @traces.last(count)
+      end
+
+      def by_module(module_name)
+        @traces.select { |trace| trace.module_name == module_name }
+      end
+
+      def successful
+        @traces.select(&:success?)
+      end
+
+      def failed
+        @traces.reject(&:success?)
+      end
+
+      def to_examples
+        @traces.map(&:to_example)
+      end
+
+      def export
+        @traces.map(&:to_h)
+      end
+
+      def filter_by_module(module_name)
+        @traces.select { |trace| trace.module_name == module_name }
+      end
+
+      def filter_by_success(success: true)
+        if success
+          @traces.select(&:success?)
+        else
+          @traces.reject(&:success?)
+        end
+      end
+
+      def filter_by_time_range(start_time, end_time)
+        @traces.select { |trace| trace.timestamp.between?(start_time, end_time) }
+      end
+
+      def statistics
+        return default_statistics if @traces.empty?
+
+        total = @traces.size
+        successful = @traces.count(&:success?)
+
+        # Group by module
+        by_module = @traces.group_by(&:module_name).transform_values do |module_traces|
+          durations = module_traces.map(&:duration_ms).compact
+          {
+            count: module_traces.size,
+            avg_duration_ms: durations.empty? ? 0 : durations.sum / durations.size.to_f
+          }
+        end
+
+        # Calculate average duration
+        all_durations = @traces.map(&:duration_ms).compact
+        avg_duration = all_durations.empty? ? 0 : all_durations.sum / all_durations.size.to_f
+
+        {
+          total_traces: total,
+          success_rate: total.positive? ? successful.to_f / total : 0,
+          average_duration_ms: avg_duration,
+          by_module: by_module
+        }
+      end
+
+      private
+
+      def default_statistics
+        {
+          total_traces: 0,
+          success_rate: 0,
+          average_duration_ms: 0,
+          by_module: {}
+        }
+      end
+    end
+
+    class TraceContext
+      attr_reader :collector
+
+      def initialize(collector = nil)
+        @collector = collector || TraceCollector.new
+        @stack = []
+      end
+
+      # Class method to temporarily use a specific collector
+      def self.with_collector(collector)
+        # Save current context if it exists
+        previous_context = Thread.current[:desiru_trace_context]
+
+        # Create new context with the provided collector
+        context = new(collector)
+        Thread.current[:desiru_trace_context] = context
+
+        # Store context for thread propagation
+        Thread.current[:desiru_trace_collector_for_threads] = collector
+
+        # Execute the block
+        yield
+      ensure
+        # Restore previous context
+        Thread.current[:desiru_trace_context] = previous_context
+        Thread.current[:desiru_trace_collector_for_threads] = nil
+      end
+
+      # Get the current trace context
+      def self.current
+        Core.trace_context
+      end
+
+      def start_trace(module_name:, signature:, inputs: {})
+        trace_data = {
+          module_name: module_name,
+          signature: signature,
+          inputs: inputs,
+          start_time: Time.now,
+          metadata: {}
+        }
+        @stack.push(trace_data)
+      end
+
+      def add_metadata(metadata)
+        return if @stack.empty?
+
+        current_trace = @stack.last
+        current_trace[:metadata] = (current_trace[:metadata] || {}).merge(metadata)
+      end
+
+      def end_trace(outputs: {}, metadata: {})
+        return if @stack.empty?
+
+        trace_data = @stack.pop
+        duration = Time.now - trace_data[:start_time]
+
+        # Merge accumulated metadata with end trace metadata
+        combined_metadata = (trace_data[:metadata] || {}).merge(metadata).merge(
+          duration: duration,
+          success: true
+        )
+
+        trace = Trace.new(
+          module_name: trace_data[:module_name],
+          signature: trace_data[:signature],
+          inputs: trace_data[:inputs],
+          outputs: outputs,
+          metadata: combined_metadata
+        )
+
+        @collector.collect(trace)
+        trace
+      end
+
+      def record_error(error, outputs: {}, metadata: {})
+        return if @stack.empty?
+
+        trace_data = @stack.pop
+        duration = Time.now - trace_data[:start_time]
+
+        # Merge accumulated metadata with error metadata
+        combined_metadata = (trace_data[:metadata] || {}).merge(metadata).merge(
+          duration: duration,
+          success: false,
+          error: error.message,
+          error_class: error.class.name
+        )
+
+        trace = Trace.new(
+          module_name: trace_data[:module_name],
+          signature: trace_data[:signature],
+          inputs: trace_data[:inputs],
+          outputs: outputs,
+          metadata: combined_metadata
+        )
+
+        @collector.collect(trace)
+        trace
+      end
+
+      def with_trace(module_name:, signature:, inputs: {})
+        start_trace(module_name: module_name, signature: signature, inputs: inputs)
+
+        begin
+          outputs = yield
+          end_trace(outputs: outputs)
+          outputs
+        rescue StandardError => e
+          record_error(e)
+          raise
+        end
+      end
+    end
+
+    class << self
+      def trace_collector
+        @trace_collector ||= TraceCollector.new
+      end
+
+      def trace_context
+        # Check for thread-local context first (set by with_collector)
+        if Thread.current[:desiru_trace_context]
+          Thread.current[:desiru_trace_context]
+        elsif Thread.current[:desiru_trace_collector_for_threads]
+          # Create a new context for this thread using the parent's collector
+          Thread.current[:desiru_trace_context] = TraceContext.new(Thread.current[:desiru_trace_collector_for_threads])
+        else
+          @trace_context ||= TraceContext.new(trace_collector)
+        end
+      end
+
+      def reset_traces!
+        @trace_collector = TraceCollector.new
+        @trace_context = TraceContext.new(@trace_collector)
+      end
+    end
+  end
+end

data/lib/desiru/core/traceable.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module Desiru
+  module Core
+    module Traceable
+      def call(inputs = {})
+        return super unless trace_enabled?
+
+        # Handle anonymous classes
+        class_name = self.class.name || "AnonymousModule"
+        module_name = class_name.split('::').last
+
+        Core.trace_context.start_trace(
+          module_name: module_name,
+          signature: signature,
+          inputs: inputs
+        )
+
+        begin
+          result = super
+
+          outputs = if result.is_a?(ModuleResult)
+                      result.outputs
+                    else
+                      result
+                    end
+
+          metadata = if result.is_a?(ModuleResult)
+                       result.metadata
+                     else
+                       {}
+                     end
+
+          Core.trace_context.end_trace(
+            outputs: outputs,
+            metadata: metadata
+          )
+
+          result
+        rescue StandardError => e
+          Core.trace_context.record_error(e)
+          raise
+        end
+      end
+
+      def trace_enabled?
+        return true unless defined?(@trace_enabled)
+
+        @trace_enabled
+      end
+
+      def enable_trace!
+        @trace_enabled = true
+      end
+
+      def disable_trace!
+        @trace_enabled = false
+      end
+    end
+  end
+end

data/lib/desiru/core.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+require_relative 'core/example'
+require_relative 'core/prediction'
+require_relative 'core/trace'
+require_relative 'core/traceable'
+require_relative 'core/compiler'
+
+module Desiru
+  module Core
+  end
+end

data/lib/desiru/module.rb

@@ -2,6 +2,7 @@
 
 require_relative 'async_capable'
 require_relative 'assertions'
+require_relative 'core/traceable'
 
 module Desiru
   # Base class for all Desiru modules
@@ -9,6 +10,7 @@ module Desiru
   class Module
     extend Forwardable
     include AsyncCapable
+    prepend Core::Traceable
 
     attr_reader :signature, :model, :config, :demos, :metadata
 
@@ -216,6 +218,12 @@ module Desiru
       end
     end
 
+    def key?(key)
+      @data.key?(key.to_sym) || @data.key?(key.to_s)
+    end
+
+    alias has_key? key?
+
     def method_missing(method_name, *args, &)
       method_str = method_name.to_s
       if method_str.end_with?('?')