dspy 0.4.0 → 0.5.1

data/lib/dspy/mixins/instrumentation_helpers.rb

@@ -0,0 +1,119 @@
+# typed: strict
+# frozen_string_literal: true
+
+require 'sorbet-runtime'
+require_relative '../instrumentation'
+
+module DSPy
+  module Mixins
+    # Shared instrumentation helper methods for DSPy modules
+    module InstrumentationHelpers
+      extend T::Sig
+
+      private
+
+      # Prepares base instrumentation payload for prediction-based modules
+      sig { params(signature_class: T.class_of(DSPy::Signature), input_values: T::Hash[Symbol, T.untyped]).returns(T::Hash[Symbol, T.untyped]) }
+      def prepare_base_instrumentation_payload(signature_class, input_values)
+        {
+          signature_class: signature_class.name,
+          model: lm.model,
+          provider: lm.provider,
+          input_fields: input_values.keys.map(&:to_s)
+        }
+      end
+
+      # Instruments a prediction operation with base payload
+      sig { params(event_name: String, signature_class: T.class_of(DSPy::Signature), input_values: T::Hash[Symbol, T.untyped], additional_payload: T::Hash[Symbol, T.untyped]).returns(T.untyped) }
+      def instrument_prediction(event_name, signature_class, input_values, additional_payload = {})
+        base_payload = prepare_base_instrumentation_payload(signature_class, input_values)
+        full_payload = base_payload.merge(additional_payload)
+        
+        # Use smart consolidation: skip nested events when higher-level events are being emitted
+        if should_emit_event?(event_name)
+          Instrumentation.instrument(event_name, full_payload) do
+            yield
+          end
+        else
+          # Skip instrumentation, just execute the block
+          yield
+        end
+      end
+
+      # Emits a validation error event
+      sig { params(signature_class: T.class_of(DSPy::Signature), validation_type: String, error_message: String).void }
+      def emit_validation_error(signature_class, validation_type, error_message)
+        Instrumentation.emit('dspy.prediction.validation_error', {
+          signature_class: signature_class.name,
+          validation_type: validation_type,
+          validation_errors: { validation_type.to_sym => error_message }
+        })
+      end
+
+      # Emits a prediction completion event
+      sig { params(signature_class: T.class_of(DSPy::Signature), success: T::Boolean, additional_data: T::Hash[Symbol, T.untyped]).void }
+      def emit_prediction_complete(signature_class, success, additional_data = {})
+        Instrumentation.emit('dspy.prediction.complete', {
+          signature_class: signature_class.name,
+          success: success
+        }.merge(additional_data))
+      end
+
+      # Determines if an event should be emitted using smart consolidation
+      sig { params(event_name: String).returns(T::Boolean) }
+      def should_emit_event?(event_name)
+        # Smart consolidation: skip nested events when higher-level events are being emitted
+        if is_nested_context?
+          # If we're in a nested context, only emit higher-level events
+          event_name.match?(/^dspy\.(chain_of_thought|react)$/)
+        else
+          # If we're not in a nested context, emit all events normally
+          true
+        end
+      end
+
+      # Determines if this is a top-level event (not nested)
+      sig { params(event_name: String).returns(T::Boolean) }
+      def is_top_level_event?(event_name)
+        # Check if we're in a nested call by looking at the call stack
+        caller_locations = caller_locations(1, 20)
+        return false if caller_locations.nil?
+        
+        # Look for other instrumentation calls in the stack
+        instrumentation_calls = caller_locations.select do |loc|
+          loc.label.include?('instrument_prediction') || 
+          loc.label.include?('instrument') ||
+          loc.path.include?('instrumentation')
+        end
+        
+        # If we have more than one instrumentation call, this is nested
+        instrumentation_calls.size <= 1
+      end
+
+      # Determines if we're in a nested call context
+      sig { returns(T::Boolean) }
+      def is_nested_call?
+        !is_top_level_event?('')
+      end
+
+      # Determines if we're in a nested context where higher-level events are being emitted
+      sig { returns(T::Boolean) }
+      def is_nested_context?
+        caller_locations = caller_locations(1, 30)
+        return false if caller_locations.nil?
+        
+        # Look for higher-level DSPy modules in the call stack
+        # We consider ChainOfThought and ReAct as higher-level modules
+        higher_level_modules = caller_locations.select do |loc|
+          loc.path.include?('chain_of_thought') || 
+          loc.path.include?('re_act') ||
+          loc.path.include?('react')
+        end
+        
+        # If we have higher-level modules in the call stack, we're in a nested context
+        higher_level_modules.any?
+      end
+
+    end
+  end
+end

data/lib/dspy/mixins/struct_builder.rb

@@ -0,0 +1,133 @@
+# typed: strict
+# frozen_string_literal: true
+
+require 'sorbet-runtime'
+
+module DSPy
+  module Mixins
+    # Shared module for building enhanced structs with input/output properties
+    module StructBuilder
+      extend T::Sig
+
+      private
+
+      # Builds a new struct class with properties from multiple sources
+      sig { params(property_sources: T::Hash[Symbol, T::Hash[Symbol, T.untyped]], additional_fields: T::Hash[Symbol, T.untyped]).returns(T.class_of(T::Struct)) }
+      def build_enhanced_struct(property_sources, additional_fields = {})
+        # Capture self to access methods from within the class block
+        builder = self
+        
+        Class.new(T::Struct) do
+          extend T::Sig
+          
+          # Add properties from each source
+          property_sources.each do |_source_name, props|
+            props.each do |name, prop|
+              type = builder.send(:extract_type_from_prop, prop)
+              options = builder.send(:extract_options_from_prop, prop)
+              
+              if options[:default]
+                const name, type, default: options[:default]
+              elsif options[:factory]
+                const name, type, factory: options[:factory]
+              else
+                const name, type
+              end
+            end
+          end
+          
+          # Add additional fields specific to the enhanced struct
+          additional_fields.each do |name, field_config|
+            type = builder.send(:extract_type_from_prop, field_config)
+            options = builder.send(:extract_options_from_prop, field_config)
+            
+            if options[:default]
+              const name, type, default: options[:default]
+            elsif options[:factory]
+              const name, type, factory: options[:factory]
+            else
+              const name, type
+            end
+          end
+          
+          include StructSerialization
+        end
+      end
+
+      # Builds properties from a props hash (from T::Struct.props)
+      sig { params(props: T::Hash[Symbol, T.untyped]).void }
+      def build_properties_from_hash(props)
+        props.each { |name, prop| build_single_property(name, prop) }
+      end
+
+      # Builds a single property with type and options
+      sig { params(name: Symbol, prop: T.untyped).void }
+      def build_single_property(name, prop)
+        type = extract_type_from_prop(prop)
+        options = extract_options_from_prop(prop)
+        
+        if options[:default]
+          const name, type, default: options[:default]
+        elsif options[:factory]
+          const name, type, factory: options[:factory]
+        else
+          const name, type
+        end
+      end
+
+      # Extracts type from property configuration
+      sig { params(prop: T.untyped).returns(T.untyped) }
+      def extract_type_from_prop(prop)
+        case prop
+        when Hash
+          prop[:type]
+        when Array
+          # Handle [Type, description] format
+          prop.first
+        else
+          prop
+        end
+      end
+
+      # Extracts options from property configuration
+      sig { params(prop: T.untyped).returns(T::Hash[Symbol, T.untyped]) }
+      def extract_options_from_prop(prop)
+        case prop
+        when Hash
+          prop.except(:type, :type_object, :accessor_key, :sensitivity, :redaction)
+        else
+          {}
+        end
+      end
+    end
+
+    # Module for adding serialization capabilities to enhanced structs
+    module StructSerialization
+      extend T::Sig
+
+      sig { returns(T::Hash[Symbol, T.untyped]) }
+      def to_h
+        hash = input_values_hash
+        hash.merge(output_properties_hash)
+      end
+
+      private
+
+      sig { returns(T::Hash[Symbol, T.untyped]) }
+      def input_values_hash
+        if instance_variable_defined?(:@input_values)
+          instance_variable_get(:@input_values) || {}
+        else
+          {}
+        end
+      end
+
+      sig { returns(T::Hash[Symbol, T.untyped]) }
+      def output_properties_hash
+        self.class.props.keys.each_with_object({}) do |key, hash|
+          hash[key] = send(key)
+        end
+      end
+    end
+  end
+end

data/lib/dspy/mixins/type_coercion.rb

@@ -0,0 +1,67 @@
+# typed: strict
+# frozen_string_literal: true
+
+require 'sorbet-runtime'
+
+module DSPy
+  module Mixins
+    # Shared module for type coercion logic across DSPy modules
+    module TypeCoercion
+      extend T::Sig
+
+      private
+
+      # Coerces output attributes to match their expected types
+      sig { params(output_attributes: T::Hash[Symbol, T.untyped], output_props: T::Hash[Symbol, T.untyped]).returns(T::Hash[Symbol, T.untyped]) }
+      def coerce_output_attributes(output_attributes, output_props)
+        output_attributes.map do |key, value|
+          prop_type = output_props[key]&.dig(:type)
+          coerced_value = coerce_value_to_type(value, prop_type)
+          [key, coerced_value]
+        end.to_h
+      end
+
+      # Coerces a single value to match its expected type
+      sig { params(value: T.untyped, prop_type: T.untyped).returns(T.untyped) }
+      def coerce_value_to_type(value, prop_type)
+        return value unless prop_type
+
+        case prop_type
+        when ->(type) { enum_type?(type) }
+          extract_enum_class(prop_type).deserialize(value)
+        when Float, ->(type) { simple_type_match?(type, Float) }
+          value.to_f
+        when Integer, ->(type) { simple_type_match?(type, Integer) }
+          value.to_i
+        else
+          value
+        end
+      end
+
+      # Checks if a type is an enum type
+      sig { params(type: T.untyped).returns(T::Boolean) }
+      def enum_type?(type)
+        (type.is_a?(Class) && type < T::Enum) ||
+          (type.is_a?(T::Types::Simple) && type.raw_type < T::Enum)
+      end
+
+      # Extracts the enum class from a type
+      sig { params(prop_type: T.untyped).returns(T.class_of(T::Enum)) }
+      def extract_enum_class(prop_type)
+        if prop_type.is_a?(Class) && prop_type < T::Enum
+          prop_type
+        elsif prop_type.is_a?(T::Types::Simple) && prop_type.raw_type < T::Enum
+          prop_type.raw_type
+        else
+          T.cast(prop_type, T.class_of(T::Enum))
+        end
+      end
+
+      # Checks if a type matches a simple type (like Float, Integer)
+      sig { params(type: T.untyped, target_type: T.untyped).returns(T::Boolean) }
+      def simple_type_match?(type, target_type)
+        type.is_a?(T::Types::Simple) && type.raw_type == target_type
+      end
+    end
+  end
+end

data/lib/dspy/predict.rb

@@ -4,6 +4,9 @@ require 'sorbet-runtime'
 require_relative 'module'
 require_relative 'instrumentation'
 require_relative 'prompt'
+require_relative 'mixins/struct_builder'
+require_relative 'mixins/type_coercion'
+require_relative 'mixins/instrumentation_helpers'
 
 module DSPy
   # Exception raised when prediction fails validation
@@ -22,6 +25,9 @@ module DSPy
 
   class Predict < DSPy::Module
     extend T::Sig
+    include Mixins::StructBuilder
+    include Mixins::TypeCoercion
+    include Mixins::InstrumentationHelpers
 
     sig { returns(T.class_of(Signature)) }
     attr_reader :signature_class
@@ -79,112 +85,63 @@ module DSPy
 
     sig { params(input_values: T.untyped).returns(T.untyped) }
     def forward_untyped(**input_values)
-      # Prepare instrumentation payload
-      input_fields = input_values.keys.map(&:to_s)
-      
-      Instrumentation.instrument('dspy.predict', {
-        signature_class: @signature_class.name,
-        model: lm.model,
-        provider: lm.provider,
-        input_fields: input_fields
-      }) do
+      instrument_prediction('dspy.predict', @signature_class, input_values) do
         # Validate input
-        begin
-          _input_struct = @signature_class.input_struct_class.new(**input_values)
-        rescue ArgumentError => e
-          # Emit validation error event
-          Instrumentation.emit('dspy.predict.validation_error', {
-            signature_class: @signature_class.name,
-            validation_type: 'input',
-            validation_errors: { input: e.message }
-          })
-          raise PredictionInvalidError.new({ input: e.message })
-        end
-
-        # Call LM
+        validate_input_struct(input_values)
+        
+        # Call LM and process response
         output_attributes = lm.chat(self, input_values)
-
-        output_attributes = output_attributes.transform_keys(&:to_sym)
-
-        output_props = @signature_class.output_struct_class.props
-        output_attributes = output_attributes.map do |key, value|
-          prop_type = output_props[key][:type] if output_props[key]
-          if prop_type
-            # Check if it's an enum (can be raw Class or T::Types::Simple)
-            enum_class = if prop_type.is_a?(Class) && prop_type < T::Enum
-                           prop_type
-                         elsif prop_type.is_a?(T::Types::Simple) && prop_type.raw_type < T::Enum
-                           prop_type.raw_type
-                         end
-
-            if enum_class
-              [key, enum_class.deserialize(value)]
-            elsif prop_type == Float || (prop_type.is_a?(T::Types::Simple) && prop_type.raw_type == Float)
-              [key, value.to_f]
-            elsif prop_type == Integer || (prop_type.is_a?(T::Types::Simple) && prop_type.raw_type == Integer)
-              [key, value.to_i]
-            else
-              [key, value]
-            end
-          else
-            [key, value]
-          end
-        end.to_h
-
-        # Create combined struct with both input and output values
-        begin
-          combined_struct = create_combined_struct_class
-          all_attributes = input_values.merge(output_attributes)
-          combined_struct.new(**all_attributes)
-        rescue ArgumentError => e
-          raise PredictionInvalidError.new({ output: e.message })
-        rescue TypeError => e
-          raise PredictionInvalidError.new({ output: e.message })
-        end
+        processed_output = process_lm_output(output_attributes)
+        
+        # Create combined result struct
+        create_prediction_result(input_values, processed_output)
       end
     end
 
     private
 
+    # Validates input using signature struct
+    sig { params(input_values: T::Hash[Symbol, T.untyped]).void }
+    def validate_input_struct(input_values)
+      @signature_class.input_struct_class.new(**input_values)
+    rescue ArgumentError => e
+      emit_validation_error(@signature_class, 'input', e.message)
+      raise PredictionInvalidError.new({ input: e.message })
+    end
+
+    # Processes LM output with type coercion
+    sig { params(output_attributes: T::Hash[T.untyped, T.untyped]).returns(T::Hash[Symbol, T.untyped]) }
+    def process_lm_output(output_attributes)
+      output_attributes = output_attributes.transform_keys(&:to_sym)
+      output_props = @signature_class.output_struct_class.props
+      
+      coerce_output_attributes(output_attributes, output_props)
+    end
+
+    # Creates the final prediction result struct
+    sig { params(input_values: T::Hash[Symbol, T.untyped], output_attributes: T::Hash[Symbol, T.untyped]).returns(T.untyped) }
+    def create_prediction_result(input_values, output_attributes)
+      begin
+        combined_struct = create_combined_struct_class
+        all_attributes = input_values.merge(output_attributes)
+        combined_struct.new(**all_attributes)
+      rescue ArgumentError => e
+        raise PredictionInvalidError.new({ output: e.message })
+      rescue TypeError => e
+        raise PredictionInvalidError.new({ output: e.message })
+      end
+    end
+
+    # Creates a combined struct class with input and output properties
     sig { returns(T.class_of(T::Struct)) }
     def create_combined_struct_class
       input_props = @signature_class.input_struct_class.props
       output_props = @signature_class.output_struct_class.props
       
-      # Create a new struct class that combines input and output fields
-      Class.new(T::Struct) do
-        extend T::Sig
-        
-        # Add input fields
-        input_props.each do |name, prop_info|
-          if prop_info[:rules]&.any? { |rule| rule.is_a?(T::Props::NilableRules) }
-            prop name, prop_info[:type], default: prop_info[:default]
-          else
-            const name, prop_info[:type], default: prop_info[:default]
-          end
-        end
-        
-        # Add output fields  
-        output_props.each do |name, prop_info|
-          if prop_info[:rules]&.any? { |rule| rule.is_a?(T::Props::NilableRules) }
-            prop name, prop_info[:type], default: prop_info[:default]
-          else
-            const name, prop_info[:type], default: prop_info[:default]
-          end
-        end
-        
-        # Add to_h method to serialize the struct to a hash
-        define_method :to_h do
-          hash = {}
-          
-          # Add all properties
-          self.class.props.keys.each do |key|
-            hash[key] = self.send(key)
-          end
-          
-          hash
-        end
-      end
+      build_enhanced_struct({
+        input: input_props,
+        output: output_props
+      })
     end
   end
 end