diff --git a/lib/ruby_llm/schema/dsl/schema_builders.rb b/lib/ruby_llm/schema/dsl/schema_builders.rb index 14d1f43..961db5d 100644 --- a/lib/ruby_llm/schema/dsl/schema_builders.rb +++ b/lib/ruby_llm/schema/dsl/schema_builders.rb @@ -41,12 +41,14 @@ def string_schema(description: nil, enum: nil, min_length: nil, max_length: nil, def number_schema(description: nil, minimum: nil, maximum: nil, greater_than: nil, less_than: nil, **options) metadata = extract_metadata!(options) + enum = options.delete(:enum) multiple_of = options.delete(:multiple_of) const = options.delete(:const) { NOT_GIVEN } raise_unknown_options!("number", options) add_const(metadata.merge({ type: "number", + enum: enum, description: description, minimum: minimum, maximum: maximum, @@ -58,12 +60,14 @@ def number_schema(description: nil, minimum: nil, maximum: nil, greater_than: ni def integer_schema(description: nil, minimum: nil, maximum: nil, greater_than: nil, less_than: nil, **options) metadata = extract_metadata!(options) + enum = options.delete(:enum) multiple_of = options.delete(:multiple_of) const = options.delete(:const) { NOT_GIVEN } raise_unknown_options!("integer", options) add_const(metadata.merge({ type: "integer", + enum: enum, description: description, minimum: minimum, maximum: maximum, diff --git a/spec/ruby_llm/schema/properties/numbers_spec.rb b/spec/ruby_llm/schema/properties/numbers_spec.rb index d1decf2..9860edf 100644 --- a/spec/ruby_llm/schema/properties/numbers_spec.rb +++ b/spec/ruby_llm/schema/properties/numbers_spec.rb @@ -36,6 +36,13 @@ expect(properties[:version]).to eq({type: "number", const: 1.5}) end + it "supports number enum values" do + schema_class.number :rating, enum: [1.5, 2.5, 3.5] + + properties = schema_class.properties + expect(properties[:rating]).to eq({type: "number", enum: [1.5, 2.5, 3.5]}) + end + it "supports number type with description" do schema_class.number :price, description: "Price field" @@ -67,4 +74,11 @@ properties = schema_class.properties expect(properties[:version]).to eq({type: "integer", const: 1}) end + + it "supports integer enum values" do + schema_class.integer :level, enum: [0, 1, 2] + + properties = schema_class.properties + expect(properties[:level]).to eq({type: "integer", enum: [0, 1, 2]}) + end end