known-types-openai 0.0.4

require 'json'
require 'pathname'

require 'codify/rust' # https://rubygems.org/gems/codify.rb

FILE_HEADER = "This is free and unencumbered software released into the public domain."
BASE_URL = 'https://platform.openai.com'
OPENAPI_FILE = 'openapi.yaml'
DERIVES = %i(Clone Debug)
CFG_DERIVES = %i(serde)

Rust = Codify::Rust
OneOf = Rust::Types::Tuple0.new(:OneOf)
AnyOf = Rust::Types::Tuple0.new(:AnyOf)
AllOf = Rust::Types::Tuple0.new(:AllOf)

task default: %w(codegen)
task codegen: %w(codegen:components codegen:groups)

require_relative '.rake/openapi_helpers'
require_relative '.rake/openapi_flattener'
require_relative '.rake/openapi_hoister'

include OpenAPIHelpers
include OpenAPIFlattener
include OpenAPIHoister

namespace :codegen do
  task flatten: [OPENAPI_FILE] do |t|
    spec = OpenAI::Spec.parse(t.prerequisites.first)
    output = {}
    spec.schemas_raw.each do |schema_ref, schema|
      #next unless schema_ref == 'AssistantObject'
      flatten_definition(schema_ref, schema, output)
      # puts; puts JSON.pretty_generate(output)
      #break if schema_ref == 'AssistantObject'
    end
    hoist_nullables!(output)
    puts JSON.pretty_generate(output)
    exit
    # # exit
    output.each do |schema_ref, schema|
      puts
      puts "// #{schema_ref}:"
      puts JSON.pretty_generate(schema)
    end
    p [spec.schemas_raw.keys.size, output.keys.size]
  end

  task debug: [OPENAPI_FILE] do |t|
    spec = OpenAI::Spec.parse(t.prerequisites.first)
    spec.schemas.each do |schema|
      #puts JSON.pretty_generate(schema.to_h)
      definition = schema.to_rust
      p [definition.name, definition]
      definition.each_subtype do |definition2|
        next if definition2.primitive?
        p [definition2.name, definition2]
      end
    end
  end

  task components: [OPENAPI_FILE] do |t|
    spec = OpenAI::Spec.parse(t.prerequisites.first)

    File.open('src/components.rs', 'w') do |out|
      out.puts "// #{FILE_HEADER}"
      out.puts
      out.puts "//! OpenAI API components"
      out.puts
      out.puts "#![allow(non_camel_case_types)]"
      out.puts
      out.puts "use crate::prelude::{String, Vec};"

      definitions = spec.schemas.map(&:to_rust)
      definitions.each do |definition|
        module_name = camel_to_snake(definition.name)
        module_path = Pathname("components/#{module_name}.rs")
        if (Pathname('src') / module_path).exist?
          out.puts
          out.puts "include!(\"#{module_path}\");"
        else
          out.puts
          definition.write(out)
        end

        definition.each_subtype do |definition|
          next if definition.primitive?
          out.puts
          definition.write(out)
        end
      end
    end
  end # :components

  task groups: [OPENAPI_FILE] do |t|
    spec = OpenAI::Spec.parse(t.prerequisites.first)

    File.open('src/groups.rs', 'w') do |out|
      out.puts "// #{FILE_HEADER}"
      out.puts
      out.puts "//! OpenAI API types organized by group"
      out.puts
      spec.groups.each do |group|
        module_name = group.snake_case_id
        out.puts "pub mod #{module_name};"
      end
    end

    spec.groups.each do |group|
      module_name = group.snake_case_id
      module_path = Pathname("src/groups/#{module_name}.rs")
      File.open(module_path, 'w') do |out|
        out.puts "// #{FILE_HEADER}"
        out.puts
        out.puts "//! **OpenAI API: #{group.title}**"

        next if group.description.empty?
        out.puts "//!"
        wrap_text(inline_to_reference_links(group.description), 80-4).each do |line|
          out.puts "//! #{line}"
        end

        link_refs = link_refs(group.description)
        next if link_refs.empty?
        out.puts "//!"
        link_refs.each do |link_ref|
          out.puts "//! #{link_ref}"
        end
      end
    end
  end # :groups
end # :codegen

module OpenAPI
  class Schema
    attr_reader :id, :ref, :type, :nullable, :default
    attr_reader :title, :description

    def initialize(id, spec)
      spec = spec.dup || {}
      @id = id ? id.to_sym : nil
      @ref = spec.delete(:'$ref')
      @type = spec[:type] ? spec.delete(:type).to_sym : nil
      @nullable = spec.delete(:nullable)
      @default = spec.delete(:default)
      @title = spec.delete(:title)
      @description = spec.delete(:description)
      @spec = spec
    end

    def to_h
      @spec.dup.merge!({
        id: @id,
        :'$ref' => @ref,
        type: @type,
        nullable: @nullable,
        title: @title,
        description: @description,
      })
    end

    def nullable?() !!@nullable end
    def summary?() !self.summary.to_s.empty? end
    def summary() first_sentence(self.description) end

    def ref?() !!@ref end
    def recursive_ref?() !!@spec[:'$recursiveRef'] end
    def one_of?() !!@spec[:oneOf] end
    def any_of?() !!@spec[:anyOf] end
    def all_of?() !!@spec[:allOf] end
    def items?() !!@spec[:items] end
    def items() Schema.new("#{self.id}_Item", @spec[:items]) end
    def properties?() !!@spec[:properties] end

    def one_of
      make_enums(@spec[:oneOf])
    end

    def any_of
      make_enums(@spec[:anyOf])
    end

    def all_of
      make_enums(@spec[:allOf])
    end

    def make_enums(inputs)
      if inputs && inputs.size == 1
        [Schema.new(self.id, inputs.first)]
      else
        (inputs || []).map.with_index { |x, i| Schema.new("#{self.id}_#{i + 1}", x) }
      end
    end

    def properties
      (@spec[:properties] || {}).inject({}) do |result, (k, v)|
        result[Identifier.new(k.to_sym)] = Schema.new("#{self.id}_#{snake_to_camel(k)}", v) # FIXME
        result
      end
    end

    def to_rust_variant
      case
        when self.ref? then self.ref.split('/').last
        when self.properties? || @type == :object
          'Object'
        when self.items? || @type == :array
          case self.items.type
            when :integer then 'ArrayOfIntegers'
            when :string then 'ArrayOfStrings'
            else 'Array'
          end
        when self.one_of? then 'OneOf'
        when self.any_of? then 'AnyOf'
        when self.all_of? then 'AllOf'
        else case @type
          when :null then 'Null'
          when :boolean then 'Boolean'
          when :integer then 'Integer'
          when :number then 'Number'
          when :string then 'String'
          else raise NotImplementedError, self.inspect
        end
      end
    end

    def to_rust(level = 0)
      type = case
        when self.ref?
          type = Rust::Types::Named.new(self.ref.split('/').last)
          type = Rust::Types::Option.new(type) if self.nullable?
          type
        when self.properties? || @type == :object
          # Some 414 top-level schemas are `type: "object"`.
          derive_default = self.properties.values.all? do |property|
            #p self if property.type.nil?
            %i(null boolean integer number string).include?(property.type)
          end
          derives = DERIVES + (derive_default ? %i(Default) : [])
          Rust::Struct.new(self.id, derives: derives, cfg_derives: CFG_DERIVES) do |definition|
            definition.comment = self.summary
            self.properties.each do |(property_id, property_type)|
              field_type = property_type.to_rust(level + 1)
              field_type = Rust::Types::Option.new(field_type).flatten if property_type.nullable?
              field = Rust::StructField.new(property_id.to_rust, field_type)
              field.comment = property_type.summary
              field.rename = property_id.id if property_id.needs_rename?
              definition.fields << field
            end
          end
        when self.items? || self.type == :array
          # Only 4 top-level schemas are `type: "array"`.
          item_type = self.items.to_rust(level + 1) rescue Rust::Types::String
          Rust::Types::Vec.new(item_type)
        when self.one_of? || self.any_of?
          # Some 23 top-level schemas are `oneOf`. Only 10 top-level schemas are `anyOf`.
          schemas = self.one_of? ? self.one_of : self.any_of
          if schemas.size == 1
            schemas.first.to_rust
          elsif schemas.all? { |t| t.type == :string }
            # e.g. RealtimeResponseCreateParams_Conversation
            Rust::TypeAlias.new(self.id, Rust::Types::String) do |definition|
              definition.comment = self.summary
            end
            # TODO: Rust::Types::String
          elsif schemas.all? { |t| t.type == :string || t.type == :null }
            # e.g. ApproximateLocation_City
            Rust::TypeAlias.new(self.id, Rust::Types::Option.new(Rust::Types::String)) do |definition|
              definition.comment = self.summary
            end
            # TODO: Rust::Types::Option.new(Rust::Types::String)
          elsif schemas.all? { |t| t.type == :object && t.properties.size == 2 && t.properties.any? { |k, _| k.to_sym == :type } }
            # e.g. CreateEvalCompletionsRunDataSource, CreateEvalResponsesRunDataSource
            Rust::Enum.new(self.id, derives: DERIVES, cfg_derives: CFG_DERIVES) do |definition|
              definition.comment = self.summary
              schemas.each do |t|
                property1, property2 = t.properties.values
                variant_name = snake_to_camel(property1.to_h[:enum].first)
                variant_type = property2.to_rust(level + 1)
                variant = Rust::EnumVariant.new(variant_name, variant_type)
                variant.comment = t.summary
                definition.variants << variant
              end
            end
          else
            Rust::Enum.new(self.id, derives: DERIVES, cfg_derives: CFG_DERIVES) do |definition|
              definition.comment = self.summary
              definition.variants << Rust::EnumVariant.new(:Null) if self.nullable?
              #out.puts "#[default]" if self.nullable? # TODO: && self.default == 'null'
              schemas.each do |t|
                next if t.recursive_ref? # skip $recursiveRef
                variant = Rust::EnumVariant.new(t.to_rust_variant, t.to_rust(level + 1))
                variant.comment = t.summary
                definition.variants << variant
              end
            end
          end
        when self.all_of?
          # Only 9 top-level schemas are `allOf`.
          schemas = self.all_of
          if schemas.size == 1
            schemas.first.to_rust
          else
            Rust::Newtype.new(self.id, AllOf, derives: DERIVES + %i(Default), cfg_derives: CFG_DERIVES) do |definition|
              definition.comment = self.summary
            end
          end
        else case @type
          # Only `ParallelToolCalls` has `type: "boolean"`.
          # TODO: spec[:default] if self.type == :boolean
          when :null then Rust::Types::Unit
          when :boolean then Rust::Types::Bool
          when :integer then Rust::Types::I64
          when :number then Rust::Types::F64
          # Some 12 schemas are `type: "string"`.
          when :string then Rust::Types::String
          else Rust::Types::String #raise NotImplementedError, self.inspect
        end
      end

      if level.zero? && !type.definition?
        # TODO: DERIVES + %i(Default)
        type = Rust::Newtype.new(self.id, type, derives: DERIVES, cfg_derives: CFG_DERIVES) do |definition|
          definition.comment = self.summary
        end
      end
      #type = @nullable ? Rust::Types::Option.new(type) : type # FIXME
      type
    end
  end # Schema

  class Identifier
    include Comparable

    attr_reader :id

    def initialize(id) @id = id end
    def <=>(other) self.id <=> other.id end
    def ==(other) self.id == other.id end
    def hash() @id.hash end
    def needs_rename?() @id.to_s.match?(/[\[\]\.]/) end
    def to_s() @id.to_s end
    def to_sym() @id.to_sym end
    def to_rust() @id.to_s.gsub('[]', '').gsub(/[-\/\.]/, '_') end
  end # Identifier
end

module OpenAI
  class Spec
    def self.parse(path)
      self.new(OpenAPIHelpers.load_file(path))
    end

    def initialize(spec) @spec = spec end

    def groups
      @spec[:'x-oaiMeta'][:groups].map { |g| Group.new(g) }.sort_by(&:snake_case_id)
    end

    def schemas
      schemas = flatten_openapi_schemas(self.schemas_raw)
      hoist_nullables!(schemas)
      schemas.map { |k, v| OpenAPI::Schema.new(k, v) }.sort_by(&:id)
    end

    def schemas_raw
      @spec[:components][:schemas]
    end
  end # Spec

  class Group
    attr_reader :spec

    def initialize(spec) @spec = spec end
    def snake_case_id() @spec[:id].gsub('-', '_') end
    def title() @spec[:title] end
    def description() @spec[:description] end
  end
end # OpenAI

def wrap_text(text, max_width = 80)
  text.to_s.gsub(/(.{1,#{max_width}})(\s+|$)/, "\\1\n").split("\n")
end

def first_sentence(text)
  text.to_s.gsub(/\s*\n+/, ' ').match(/^.*?[.!?](?:\s|$)/)&.[](0)&.strip || text.to_s.strip
end

def link_refs(markdown, base_url = BASE_URL)
  base_url = base_url.to_s.chomp('/')
  markdown.to_s.scan(/\[(.*?)\]\((\/[^)]*)\)/)
    .map { |_, path| "[#{path}]: #{base_url}#{path}" }
    .uniq
end

def inline_to_reference_links(markdown)
  markdown.to_s.gsub(/\[(.*?)\]\(((?!https:).*?)\)/, '[\1][\2]')
end

def camel_to_snake(name)
  name.to_s.gsub(/([A-Z])/, '_\1').downcase.gsub(/^_/, '')
end

def snake_to_camel(name)
  name.to_s.gsub(/[_.]/, ' ').split.map(&:capitalize).join
end