dspy 0.3.1 → 0.5.0

data/lib/dspy/chain_of_thought.rb

@@ -5,25 +5,110 @@ require 'sorbet-runtime'
 require_relative 'predict'
 require_relative 'signature'
 require_relative 'instrumentation'
+require_relative 'mixins/struct_builder'
 
 module DSPy
   # Enhances prediction by encouraging step-by-step reasoning
   # before providing a final answer using Sorbet signatures.
   class ChainOfThought < Predict
     extend T::Sig
+    include Mixins::StructBuilder
 
     FieldDescriptor = DSPy::Signature::FieldDescriptor
 
     sig { params(signature_class: T.class_of(DSPy::Signature)).void }
     def initialize(signature_class)
       @original_signature = signature_class
+      enhanced_signature = build_enhanced_signature(signature_class)
+      
+      # Call parent constructor with enhanced signature
+      super(enhanced_signature)
+      @signature_class = enhanced_signature
+    end
+
+    # Override prompt-based methods to maintain ChainOfThought behavior
+    sig { override.params(new_prompt: Prompt).returns(ChainOfThought) }
+    def with_prompt(new_prompt)
+      # Create a new ChainOfThought with the same original signature
+      instance = self.class.new(@original_signature)
+      
+      # Ensure the instruction includes "Think step by step" if not already present
+      enhanced_instruction = if new_prompt.instruction.include?("Think step by step")
+                               new_prompt.instruction
+                             else
+                               "#{new_prompt.instruction} Think step by step."
+                             end
+      
+      # Create enhanced prompt with ChainOfThought-specific schemas
+      enhanced_prompt = Prompt.new(
+        instruction: enhanced_instruction,
+        input_schema: @signature_class.input_json_schema,
+        output_schema: @signature_class.output_json_schema,
+        few_shot_examples: new_prompt.few_shot_examples,
+        signature_class_name: @signature_class.name
+      )
+      
+      instance.instance_variable_set(:@prompt, enhanced_prompt)
+      instance
+    end
+
+    sig { override.params(instruction: String).returns(ChainOfThought) }
+    def with_instruction(instruction)
+      enhanced_instruction = ensure_chain_of_thought_instruction(instruction)
+      super(enhanced_instruction)
+    end
+
+    sig { override.params(examples: T::Array[FewShotExample]).returns(ChainOfThought) }
+    def with_examples(examples)
+      # Convert examples to include reasoning if they don't have it
+      enhanced_examples = examples.map do |example|
+        if example.reasoning.nil? || example.reasoning.empty?
+          # Try to extract reasoning from the output if it contains a reasoning field
+          reasoning = example.output[:reasoning] || "Step by step reasoning for this example."
+          DSPy::FewShotExample.new(
+            input: example.input,
+            output: example.output,
+            reasoning: reasoning
+          )
+        else
+          example
+        end
+      end
+      
+      super(enhanced_examples)
+    end
+
+    # Access to the original signature for optimization
+    sig { returns(T.class_of(DSPy::Signature)) }
+    attr_reader :original_signature
 
-      # Create enhanced output struct with reasoning
+    # Override forward_untyped to add ChainOfThought-specific instrumentation
+    sig { override.params(input_values: T.untyped).returns(T.untyped) }
+    def forward_untyped(**input_values)
+      instrument_prediction('dspy.chain_of_thought', @original_signature, input_values) do
+        # Call parent prediction logic
+        prediction_result = super(**input_values)
+        
+        # Analyze reasoning if present
+        analyze_reasoning(prediction_result)
+        
+        prediction_result
+      end
+    end
+
+    private
+
+    # Builds enhanced signature with reasoning capabilities
+    sig { params(signature_class: T.class_of(DSPy::Signature)).returns(T.class_of(DSPy::Signature)) }
+    def build_enhanced_signature(signature_class)
       enhanced_output_struct = create_enhanced_output_struct(signature_class)
+      create_signature_class(signature_class, enhanced_output_struct)
+    end
 
-      # Create enhanced signature class
-      enhanced_signature = Class.new(DSPy::Signature) do
-        # Set the description
+    # Creates signature class with enhanced description and reasoning field
+    sig { params(signature_class: T.class_of(DSPy::Signature), enhanced_output_struct: T.class_of(T::Struct)).returns(T.class_of(DSPy::Signature)) }
+    def create_signature_class(signature_class, enhanced_output_struct)
+      Class.new(DSPy::Signature) do
         description "#{signature_class.description} Think step by step."
 
         # Use the same input struct and copy field descriptors
@@ -47,52 +132,49 @@ module DSPy
           attr_reader :input_struct_class, :output_struct_class
         end
       end
+    end
 
-      # Call parent constructor with enhanced signature
-      super(enhanced_signature)
-      @signature_class = enhanced_signature
+    # Creates enhanced output struct with reasoning field
+    sig { params(signature_class: T.class_of(DSPy::Signature)).returns(T.class_of(T::Struct)) }
+    def create_enhanced_output_struct(signature_class)
+      output_props = signature_class.output_struct_class.props
+      
+      build_enhanced_struct(
+        { output: output_props },
+        { reasoning: [String, "Step by step reasoning process"] }
+      )
     end
 
-    # Override forward_untyped to add ChainOfThought-specific instrumentation
-    sig { override.params(input_values: T.untyped).returns(T.untyped) }
-    def forward_untyped(**input_values)
-      # Prepare instrumentation payload
-      input_fields = input_values.keys.map(&:to_s)
+    # Ensures instruction includes chain of thought prompt
+    sig { params(instruction: String).returns(String) }
+    def ensure_chain_of_thought_instruction(instruction)
+      instruction.include?("Think step by step") ? instruction : "#{instruction} Think step by step."
+    end
+
+    # Analyzes reasoning in prediction result and emits instrumentation events
+    sig { params(prediction_result: T.untyped).void }
+    def analyze_reasoning(prediction_result)
+      return unless prediction_result.respond_to?(:reasoning) && prediction_result.reasoning
       
-      # Instrument ChainOfThought lifecycle
-      result = Instrumentation.instrument('dspy.chain_of_thought', {
-        signature_class: @original_signature.name,
-        model: lm.model,
-        provider: lm.provider,
-        input_fields: input_fields
-      }) do
-        # Call parent prediction logic
-        prediction_result = super(**input_values)
-        
-        # Analyze reasoning if present
-        if prediction_result.respond_to?(:reasoning) && prediction_result.reasoning
-          reasoning_content = prediction_result.reasoning.to_s
-          reasoning_length = reasoning_content.length
-          reasoning_steps = count_reasoning_steps(reasoning_content)
-          
-          # Emit reasoning analysis event
-          Instrumentation.emit('dspy.chain_of_thought.reasoning_complete', {
-            signature_class: @original_signature.name,
-            reasoning_steps: reasoning_steps,
-            reasoning_length: reasoning_length,
-            has_reasoning: !reasoning_content.empty?
-          })
-        end
-        
-        prediction_result
-      end
+      reasoning_content = prediction_result.reasoning.to_s
+      return if reasoning_content.empty?
       
-      result
+      emit_reasoning_analysis(reasoning_content)
     end
 
-    private
+    # Emits reasoning analysis instrumentation event
+    sig { params(reasoning_content: String).void }
+    def emit_reasoning_analysis(reasoning_content)
+      Instrumentation.emit('dspy.chain_of_thought.reasoning_complete', {
+        signature_class: @original_signature.name,
+        reasoning_steps: count_reasoning_steps(reasoning_content),
+        reasoning_length: reasoning_content.length,
+        has_reasoning: true
+      })
+    end
 
     # Count reasoning steps by looking for step indicators
+    sig { params(reasoning_text: String).returns(Integer) }
     def count_reasoning_steps(reasoning_text)
       return 0 if reasoning_text.nil? || reasoning_text.empty?
       
@@ -113,50 +195,5 @@ module DSPy
       # Fallback: count sentences if no clear steps
       max_count > 0 ? max_count : reasoning_text.split(/[.!?]+/).reject(&:empty?).length
     end
-
-    sig { params(signature_class: T.class_of(DSPy::Signature)).returns(T.class_of(T::Struct)) }
-    def create_enhanced_output_struct(signature_class)
-      # Get original output props
-      original_props = signature_class.output_struct_class.props
-
-      # Create new struct class with reasoning added
-      Class.new(T::Struct) do
-        # Add all original fields
-        original_props.each do |name, prop|
-          # Extract the type and other options
-          type = prop[:type]
-          options = prop.except(:type, :type_object, :accessor_key, :sensitivity, :redaction)
-
-          # Handle default values
-          if options[:default]
-            const name, type, default: options[:default]
-          elsif options[:factory]
-            const name, type, factory: options[:factory]
-          else
-            const name, type
-          end
-        end
-
-        # Add reasoning field (ChainOfThought always provides this)
-        const :reasoning, String
-
-        # Add to_h method to serialize the struct to a hash
-        define_method :to_h do
-          hash = {}
-
-          # Start with input values if available
-          if self.instance_variable_defined?(:@input_values)
-            hash.merge!(self.instance_variable_get(:@input_values))
-          end
-
-          # Then add output properties
-          self.class.props.keys.each do |key|
-            hash[key] = self.send(key)
-          end
-
-          hash
-        end
-      end
-    end
   end
 end