Fix reference in objects

Author: danielfriis

README.md

@@ -271,9 +271,14 @@ class MySchema < RubyLLM::Schema
     string :longitude
   end
 
+  # Using a reference in an array
   array :coordinates, of: :location
-  
-  object :home_location do
+
+  # Using a reference in an object via the `reference` option
+  object :home_location, reference: :location
+
+  # Using a reference in an object via block
+  object :user do
     reference :location
   end
 end

lib/ruby_llm/schema/dsl.rb

@@ -41,8 +41,8 @@ def null(name = nil, description: nil, required: true)
       end
 
       # Complex type methods
-      def object(name = nil, description: nil, required: true, &block)
-        add_property(name, build_property_schema(:object, description: description, &block), required: required)
+      def object(name = nil, reference: nil, description: nil, required: true, &block)
+        add_property(name, build_property_schema(:object, description: description, reference: reference, &block), required: required)
       end
 
       def array(name, of: nil, description: nil, required: true, min_items: nil, max_items: nil, &block)
@@ -123,16 +123,26 @@ def build_property_schema(type, **options, &)
         when :null
           {type: "null", description: options[:description]}.compact
         when :object
-          sub_schema = Class.new(Schema)
-          sub_schema.class_eval(&)
+          # If the reference option is provided, return the reference
+          return reference(options[:reference]) if options[:reference]
 
-          {
-            type: "object",
-            properties: sub_schema.properties,
-            required: sub_schema.required_properties,
-            additionalProperties: additional_properties,
-            description: options[:description]
-          }.compact
+          sub_schema = Class.new(Schema)
+          
+          # Evaluate the block and capture the result
+          result = sub_schema.class_eval(&)
+          
+          # If the block returned a reference and no properties were added, use the reference
+          if result.is_a?(Hash) && result["$ref"] && sub_schema.properties.empty?
+            result.merge(options[:description] ? {description: options[:description]} : {})
+          else
+            {
+              type: "object",
+              properties: sub_schema.properties,
+              required: sub_schema.required_properties,
+              additionalProperties: additional_properties,
+              description: options[:description]
+            }.compact
+          end
         when :any_of
           schemas = collect_property_schemas_from_block(&)
           {

spec/ruby_llm/schema_spec.rb

@@ -220,6 +220,58 @@
       object_schema = any_of_schemas.find { |s| s[:type] == "object" }
       expect(object_schema[:properties][:nested_field]).to eq({type: "string"})
     end
+
+    it "supports reference to a defined schema by block" do
+      schema_class.define :address do
+        string :street
+        string :city
+      end
+      
+      schema_class.object :user do
+        string :name
+        object :address do
+          reference :address
+        end
+      end
+
+      instance = schema_class.new
+      json_output = instance.to_json_schema
+
+      expect(json_output[:schema][:properties][:user][:properties][:address]).to eq({"$ref" => "#/$defs/address"})
+      expect(json_output[:schema]["$defs"][:address]).to eq({
+        type: "object",
+        properties: {
+          street: {type: "string"},
+          city: {type: "string"}
+        },
+        required: %i[street city]
+      })
+    end
+
+    it "supports reference to a defined schema by reference option" do
+      schema_class.define :address do
+        string :street
+        string :city
+      end
+      
+      schema_class.object :user do
+        string :name
+        object :address, reference: :address
+      end
+
+      instance = schema_class.new
+      json_output = instance.to_json_schema
+
+      expect(json_output[:schema][:properties][:user][:properties][:address]).to eq({"$ref" => "#/$defs/address"})
+      expect(json_output[:schema]["$defs"][:address]).to eq({
+        type: "object",
+        properties: {
+          street: {type: "string"},
+          city: {type: "string"}
+        },
+        required: %i[street city]
+      })
+    end
   end
 
   # ===========================================