fix quickstart: ensure original_signature is Signature

lib/desiru/module.rb

@@ -15,15 +15,7 @@ class Module
     attr_reader :signature, :model, :config, :demos, :metadata
 
     def initialize(signature, model: nil, config: {}, demos: [], metadata: {})
-      @signature = case signature
-                   when Signature
-                     signature
-                   when String
-                     Signature.new(signature)
-                   else
-                     raise ModuleError, 'Signature must be a String or Signature instance'
-                   end
-
+      @signature = Signature.wrap(signature)
       @model = model || Desiru.configuration.default_model
       @config = default_config.merge(config)
       @demos = demos

lib/desiru/modules/chain_of_thought.rb

@@ -4,11 +4,13 @@ module Desiru
   module Modules
     # Chain of Thought module - adds reasoning steps before producing outputs
     class ChainOfThought < Predict
+      attr_reader :original_signature
+
       def initialize(signature, **)
         # Extend signature to include reasoning field
         extended_sig = extend_signature_with_reasoning(signature)
         super(extended_sig, **)
-        @original_signature = signature
+        @original_signature = Signature.wrap(signature)
       end
 
       protected

lib/desiru/signature.rb

@@ -76,6 +76,17 @@ def output_fields
     alias inputs input_fields
     alias outputs output_fields
 
+    def self.wrap(signature_string_or_instance)
+      case signature_string_or_instance
+      when Signature
+        signature_string_or_instance
+      when String
+        Signature.new(signature_string_or_instance)
+      else
+        raise ModuleError, 'Signature must be a String or Signature instance'
+      end
+    end
+
     def initialize(signature_string, descriptions: {})
       @raw_signature = signature_string
       @descriptions = descriptions

spec/desiru/modules/chain_of_thought_spec.rb

@@ -0,0 +1,40 @@
+# frozen_string_literal: true
+
+require 'spec_helper'
+require 'desiru/modules/chain_of_thought'
+
+RSpec.describe Desiru::Modules::ChainOfThought do
+  let(:mock_model) do
+    double('Model',
+           complete: { content: <<~CNT },
+             reasoning: Some reason for this mock answer
+             answer: 42
+           CNT
+           temperature: 0.7,
+           respond_to?: true)
+  end
+
+  describe '#initialize' do
+    context 'with signature string' do
+      let(:signature_string) { 'question: string -> answer: string' }
+      let(:cot_module) { described_class.new(signature_string, model: mock_model) }
+
+      it 'wraps original_signature with Signature instance' do
+        expect(cot_module.original_signature).to respond_to(:output_fields)
+        expect(cot_module.original_signature).to be_a(Desiru::Signature)
+      end
+    end
+  end
+
+  describe "#forward" do
+    context 'with signature string' do
+      let(:signature_string) { 'question: string -> answer: string' }
+      let(:cot_module) { described_class.new(signature_string, model: mock_model) }
+
+      it 'can build the prompt and complete a response' do
+        result = cot_module.call(question: "Two dice are tossed. What is the probability that the sum equals two?")
+        expect(result.answer).to eq "42"
+      end
+    end
+  end
+end