Implement working with_response_format for Open AI

Author: jayelkaake

.gitignore

@@ -47,8 +47,8 @@ build-iPhoneSimulator/
 # for a library or gem, you might want to ignore these files since the code is
 # intended to run in multiple environments; otherwise, check them in:
 Gemfile.lock
-# .ruby-version
-# .ruby-gemset
+.ruby-version
+.ruby-gemset
 
 # unless supporting rvm < 1.11.0 or doing something fancy, ignore this:
 .rvmrc
@@ -57,3 +57,4 @@ Gemfile.lock
 # .rubocop-https?--*
 
 repomix-output.*
+/.idea/

Gemfile

@@ -18,6 +18,7 @@ group :development do
   gem 'nokogiri'
   gem 'overcommit', '>= 0.66'
   gem 'pry', '>= 0.14'
+  gem 'pry-byebug', '>= 3.11'
   gem 'rake', '>= 13.0'
   gem 'rdoc'
   gem 'reline'

README.md

@@ -60,6 +60,9 @@ chat.ask "Tell me a story about a Ruby programmer" do |chunk|
   print chunk.content
 end
 
+# Get structured responses easily (OpenAI only for now)
+chat.with_response_format(:integer).ask("What is 2 + 2?").to_i # => 4
+
 # Generate images
 RubyLLM.paint "a sunset over mountains in watercolor style"
 

docs/guides/chat.md

@@ -261,6 +261,54 @@ end
 chat.ask "What is metaprogramming in Ruby?"
 ```
 
+## Receiving Structured Responses
+You can ensure the responses follow a schema you define like this:
+```ruby
+chat = RubyLLM.chat
+
+chat.with_response_format(:integer).ask("What is 2 + 2?").to_i
+# => 4
+
+chat.with_response_format(:string).ask("Say 'Hello World' and nothing else.").content
+# => "Hello World"
+
+chat.with_response_format(:array, items: { type: :string })
+chat.ask('What are the 2 largest countries? Only respond with country names.').content
+# => ["Russia", "Canada"]
+
+chat.with_response_format(:object, properties: { age: { type: :integer } })
+chat.ask('Provide sample customer age between 10 and 100.').content
+# => { "age" => 42 }
+
+chat.with_response_format(
+  :object,
+  properties: { hobbies: { type: :array, items: { type: :string, enum: %w[Soccer Golf Hockey] } } }
+)
+chat.ask('Provide at least 1 hobby.').content
+# => { "hobbies" => ["Soccer"] }
+```
+
+You can also provide the JSON schema you want directly to the method like this:
+```ruby
+chat.with_response_format(type: :object, properties: { age: { type: :integer } })
+# => { "age" => 31 }
+```
+
+In this example the code is automatically switching to OpenAI's json_mode since no object properties are requested:
+```ruby
+chat.with_response_format(:json) # Don't care about structure, just give me JSON
+
+chat.ask('Provide a sample customer data object with name and email keys.').content
+# => { "name" => "Tobias", "email" => "[email protected]" }
+
+chat.ask('Provide a sample customer data object with name and email keys.').content
+# => { "first_name" => "Michael", "email_address" => "[email protected]" }
+```
+
+{: .note }
+**Only OpenAI supported for now:** Only OpenAI models support this feature for now. We will add support for other models shortly.
+
+
 ## Next Steps
 
 This guide covered the core `Chat` interface. Now you might want to explore:
@@ -269,4 +317,4 @@ This guide covered the core `Chat` interface. Now you might want to explore:
 *   [Using Tools]({% link guides/tools.md %}): Enable the AI to call your Ruby code.
 *   [Streaming Responses]({% link guides/streaming.md %}): Get real-time feedback from the AI.
 *   [Rails Integration]({% link guides/rails.md %}): Persist your chat conversations easily.
-*   [Error Handling]({% link guides/error-handling.md %}): Build robust applications that handle API issues.
+*   [Error Handling]({% link guides/error-handling.md %}): Build robust applications that handle API issues.

lib/ruby_llm/chat.rb

@@ -31,6 +31,56 @@ def initialize(model: nil, provider: nil, assume_model_exists: false, context: n
       }
     end
 
+    ##
+    # This method lets you ensure the responses follow a schema you define like this:
+    #
+    #   chat.with_response_format(:integer).ask("What is 2 + 2?").to_i
+    #   # => 4
+    #   chat.with_response_format(:string).ask("Say 'Hello World' and nothing else.").content
+    #   # => "Hello World"
+    #   chat.with_response_format(:array, items: { type: :string })
+    #   chat.ask('What are the 2 largest countries? Only respond with country names.').content
+    #   # => ["Russia", "Canada"]
+    #   chat.with_response_format(:object, properties: { age: { type: :integer } })
+    #   chat.ask('Provide sample customer age between 10 and 100.').content
+    #   # => { "age" => 42 }
+    #   chat.with_response_format(
+    #     :object,
+    #     properties: { hobbies: { type: :array, items: { type: :string, enum: %w[Soccer Golf Hockey] } } }
+    #   )
+    #   chat.ask('Provide at least 1 hobby.').content
+    #   # => { "hobbies" => ["Soccer"] }
+    #
+    # You can also provide the JSON schema you want directly to the method like this:
+    #   chat.with_response_format(type: :object, properties: { age: { type: :integer } })
+    #   # => { "age" => 31 }
+    #
+    # In this example the code is automatically switching to OpenAI's json_mode since no object
+    # properties are requested:
+    #   chat.with_response_format(:json) # Don't care about structure, just give me JSON
+    #   chat.ask('Provide a sample customer data object with name and email keys.').content
+    #   # => { "name" => "Tobias", "email" => "[email protected]" }
+    #   chat.ask('Provide a sample customer data object with name and email keys.').content
+    #   # => { "first_name" => "Michael", "email_address" => "[email protected]" }
+    #
+    # @param type [Symbol] (optional) This can be anything supported by the API JSON schema types (integer, object, etc)
+    # @param schema [Hash] The schema for the response format. It can be a JSON schema or a simple hash.
+    # @return [Chat] (self)
+    def with_response_format(type = nil, **schema)
+      schema_hash = if type.is_a?(Symbol) || type.is_a?(String)
+                      { type: type == :json ? :object : type }
+                    elsif type.is_a?(Hash)
+                      type
+                    else
+                      {}
+                    end.merge(schema)
+
+      @response_schema = Schema.new(schema_hash)
+
+      self
+    end
+    alias with_structured_response with_response_format
+
     def ask(message = nil, with: {}, &block)
       add_message role: :user, content: Content.new(message, with)
       complete(&block)
@@ -96,14 +146,16 @@ def each(&)
 
     def complete(&) # rubocop:disable Metrics/MethodLength
       @on[:new_message]&.call
-      response = @provider.complete(
-        messages,
-        tools: @tools,
-        temperature: @temperature,
-        model: @model.id,
-        connection: @connection,
-        &
-      )
+      response = @provider.with_response_schema(@response_schema) do
+        @provider.complete(
+          messages,
+          tools: @tools,
+          temperature: @temperature,
+          model: @model.id,
+          connection: @connection,
+          &
+        )
+      end
       @on[:end_message]&.call(response)
 
       add_message response

lib/ruby_llm/message.rb

@@ -7,7 +7,9 @@ module RubyLLM
   class Message
     ROLES = %i[system user assistant tool].freeze
 
-    attr_reader :role, :content, :tool_calls, :tool_call_id, :input_tokens, :output_tokens, :model_id
+    attr_reader :role, :tool_calls, :tool_call_id, :input_tokens, :output_tokens, :model_id
+
+    delegate :to_i, :to_a, :to_s, to: :content
 
     def initialize(options = {})
       @role = options[:role].to_sym
@@ -17,10 +19,29 @@ def initialize(options = {})
       @output_tokens = options[:output_tokens]
       @model_id = options[:model_id]
       @tool_call_id = options[:tool_call_id]
+      @schema = options[:schema]
 
       ensure_valid_role
     end
 
+    def content
+      return @content unless @schema.present?
+
+      if @schema[:type].to_s == :object.to_s && @schema[:properties].to_h.keys.none?
+        json_response
+      else
+        structured_content
+      end
+    end
+
+    def json_response
+      JSON.parse(@content)
+    end
+
+    def structured_content
+      json_response['result']
+    end
+
     def tool_call?
       !tool_calls.nil? && !tool_calls.empty?
     end

lib/ruby_llm/provider.rb

@@ -31,6 +31,29 @@ def list_models(connection:)
         parse_list_models_response response, slug, capabilities
       end
 
+      ##
+      # @return [::RubyLLM::Schema, NilClass]
+      def response_schema
+        Thread.current[:response_schema]
+      end
+
+      ##
+      # @param response_schema [::RubyLLM::Schema]
+      def with_response_schema(response_schema)
+        prev_response_schema = Thread.current[:response_schema]
+
+        result = nil
+        begin
+          Thread.current[:response_schema] = response_schema
+
+          result = yield
+        ensure
+          Thread.current[:response_schema] = prev_response_schema
+        end
+
+        result
+      end
+
       def embed(text, model:, connection:)
         payload = render_embedding_payload(text, model:)
         response = connection.post embedding_url, payload

lib/ruby_llm/providers/openai/chat.rb

@@ -22,6 +22,9 @@ def render_payload(messages, tools:, temperature:, model:, stream: false) # rubo
               payload[:tools] = tools.map { |_, tool| tool_for(tool) }
               payload[:tool_choice] = 'auto'
             end
+
+            add_response_schema_to_payload(payload) if response_schema.present?
+
             payload[:stream_options] = { include_usage: true } if stream
           end
         end
@@ -35,6 +38,7 @@ def parse_completion_response(response) # rubocop:disable Metrics/MethodLength
 
           Message.new(
             role: :assistant,
+            schema: response_schema,
             content: message_data['content'],
             tool_calls: parse_tool_calls(message_data['tool_calls']),
             input_tokens: data['usage']['prompt_tokens'],
@@ -62,6 +66,54 @@ def format_role(role)
             role.to_s
           end
         end
+
+        private
+
+        ##
+        # @param [Hash] payload
+        def add_response_schema_to_payload(payload)
+          payload[:response_format] = gen_response_format_request
+
+          return unless payload[:response_format][:type] == :json_object
+
+          # NOTE: this is required by the Open AI API when requesting arbitrary JSON.
+          payload[:messages].unshift({ role: :developer, content: <<~GUIDANCE
+            You must format your output as a valid JSON object.
+            Format your entire response as valid JSON.
+            Do not include explanations, markdown formatting, or any text outside the JSON.
+          GUIDANCE
+          })
+        end
+
+        ##
+        # @return [Hash]
+        def gen_response_format_request
+          if response_schema[:type].to_s == :object.to_s && response_schema[:properties].to_h.keys.none?
+            { type: :json_object } # Assume we just want json_mode
+          else
+            gen_json_schema_format_request
+          end
+        end
+
+        def gen_json_schema_format_request # rubocop:disable Metrics/MethodLength -- because it's mostly the standard hash
+          result_schema = response_schema.dup # so we don't modify the original in the thread
+          result_schema.add_to_each_object_type!(:additionalProperties, false)
+          result_schema.add_to_each_object_type!(:required, ->(schema) { schema[:properties].to_h.keys })
+
+          {
+            type: :json_schema,
+            json_schema: {
+              name: :response,
+              schema: {
+                type: :object,
+                properties: { result: result_schema.to_h },
+                additionalProperties: false,
+                required: [:result]
+              },
+              strict: true
+            }
+          }
+        end
       end
     end
   end

lib/ruby_llm/schema.rb

@@ -0,0 +1,76 @@
+# frozen_string_literal: true
+
+module RubyLLM
+  ##
+  # Schema class for defining the structure of data objects.
+  # Wraps the #Hash class
+  # @see #Hash
+  class Schema
+    delegate_missing_to :@schema
+
+    def initialize(schema = {})
+      @schema = deep_transform_keys_in_object(schema.to_h.dup, &:to_sym)
+    end
+
+    def [](key)
+      @schema[key.to_sym]
+    end
+
+    def []=(key, new_value)
+      @schema[key.to_sym] = deep_transform_keys_in_object(new_value, &:to_sym)
+    end
+
+    # Adds the new_value into the new_key key for every sub-schema that is of type: :object
+    # @param new_key [Symbol] The key to add to each object type.
+    # @param new_value [Boolean, String] The value to assign to the new key.
+    def add_to_each_object_type!(new_key, new_value)
+      add_to_each_object_type(new_key, new_value, @schema)
+    end
+
+    # @return [Boolean]
+    def present?
+      @schema.present? && @schema[:type].present?
+    end
+
+    private
+
+    def add_to_each_object_type(new_key, new_value, schema)
+      return schema unless schema.is_a?(Hash)
+
+      if schema[:type].to_s == :object.to_s
+        add_to_object_type(new_key, new_value, schema)
+      elsif schema[:type].to_s == :array.to_s && schema[:items]
+        schema[:items] = add_to_each_object_type(new_key, new_value, schema[:items])
+      end
+
+      schema
+    end
+
+    def add_to_object_type(new_key, new_value, schema)
+      if schema[new_key.to_sym].nil?
+        schema[new_key.to_sym] = new_value.is_a?(Proc) ? new_value.call(schema) : new_value
+      end
+
+      schema[:properties]&.transform_values! { |value| add_to_each_object_type(new_key, new_value, value) }
+    end
+
+    ##
+    # Recursively transforms keys in a hash or array to symbols.
+    # Borrowed from ActiveSupport's Hash#deep_transform_keys
+    # @param object [Object] The object to transform.
+    # @param block [Proc] The block to apply to each key.
+    # @return [Object] The transformed object.
+    def deep_transform_keys_in_object(object, &block)
+      case object
+      when Hash
+        object.each_with_object({}) do |(key, value), result|
+          result[yield(key)] = deep_transform_keys_in_object(value, &block)
+        end
+      when Array
+        object.map { |e| deep_transform_keys_in_object(e, &block) }
+      else
+        object
+      end
+    end
+  end
+end