Try again

Author: danielfriis

lib/ruby_llm/template/chat_extension.rb

@@ -26,10 +26,15 @@ def with_template(template_name, context = {})
         if loader.available_roles.include?("schema")
           schema_result = loader.render_template("schema", context)
 
-          if schema_result&.respond_to?(:to_json_schema)
-            # It's a RubyLLM::Schema instance
-            schema_data = schema_result.to_json_schema
-            with_schema(schema_data)
+          if schema_result
+            if schema_result.is_a?(Class) && schema_result < RubyLLM::Schema
+              # It's a schema class, pass it directly to with_schema
+              with_schema(schema_result)
+            elsif schema_result.respond_to?(:to_json_schema)
+              # It's a RubyLLM::Schema instance
+              schema_data = schema_result.to_json_schema
+              with_schema(schema_data)
+            end
           end
         end
 

lib/ruby_llm/template/loader.rb

@@ -112,7 +112,7 @@ def load_schema_class(context = {})
         end
 
         if schema_class
-          return schema_class.new
+          return schema_class
         end
 
         # 3. If nothing worked, provide detailed error