chore: extract reused JSON schema references even in unions (#761)

Author: ms-jpq

lib/openai/helpers/structured_output/json_schema_converter.rb

@@ -12,12 +12,6 @@ module JsonSchemaConverter
           end
         end.freeze
         # @api private
-        COUNTER = Object.new.tap do
-          _1.define_singleton_method(:inspect) do
-            "#<#{OpenAI::Helpers::StructuredOutput::JsonSchemaConverter}::COUNTER>"
-          end
-        end.freeze
-        # @api private
         NO_REF = Object.new.tap do
           _1.define_singleton_method(:inspect) do
             "#<#{OpenAI::Helpers::StructuredOutput::JsonSchemaConverter}::NO_REF>"
@@ -81,14 +75,13 @@ def to_nilable(schema)
           def cache_def!(state, type:, &blk)
             defs, path = state.fetch_values(:defs, :path)
             if (stored = defs[type])
-              stored[OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::COUNTER] += 1
-              stored.fetch(OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::POINTER)
+              pointers = stored.fetch(OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::POINTER)
+              pointers.first.except(OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::NO_REF).tap { pointers << _1 }
             else
               ref_path = String.new
               ref = {"$ref": ref_path}
               stored = {
-                OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::POINTER => ref,
-                OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::COUNTER => 1
+                OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::POINTER => [ref]
               }
               defs.store(type, stored)
               schema = blk.call
@@ -112,17 +105,21 @@ def to_json_schema(type)
             )
             reused_defs = {}
             defs.each_value do |acc|
-              ref = acc.fetch(OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::POINTER)
-              if (no_ref = ref.delete(OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::NO_REF))
-                acc[OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::COUNTER] -= 1
+              sch = acc.except(OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::POINTER)
+              pointers = acc.fetch(OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::POINTER)
+
+              no_refs, refs = pointers.partition do
+                _1.delete(OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::NO_REF)
               end
-              cnt = acc.fetch(OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::COUNTER)
 
-              sch = acc.except(
-                OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::POINTER,
-                OpenAI::Helpers::StructuredOutput::JsonSchemaConverter::COUNTER
-              )
-              cnt > 1 && !no_ref ? reused_defs.store(ref.fetch(:$ref), sch) : ref.replace(sch)
+              case refs
+              in [_, ref, *]
+                reused_defs.store(ref.fetch(:$ref), sch)
+              in [ref]
+                ref.replace(sch)
+              else
+              end
+              no_refs.each { _1.replace(sch) }
             end
 
             xformed = reused_defs.transform_keys { _1.delete_prefix("#/$defs/") }

test/openai/helpers/structured_output_test.rb

@@ -311,19 +311,20 @@ def test_definition_reusing
         ]
       },
       M11 => {
-        type: "object",
-        properties: {
+        :$defs => {".a/?.0/[]" => {type: "array", items: {type: "string"}}},
+        :type => "object",
+        :properties => {
           a: {
             anyOf: [
               {type: "array", items: {type: "string"}},
               {type: "array", items: {type: "string"}}
             ]
           },
-          b: {type: "array", items: {type: "string"}},
-          c: {type: "array", items: {type: "string"}}
+          b: {:$ref => "#/$defs/.a/?.0/[]"},
+          c: {:$ref => "#/$defs/.a/?.0/[]"}
         },
-        required: %w[a b c],
-        additionalProperties: false
+        :required => %w[a b c],
+        :additionalProperties => false
       }
     }