Switch to JsonStreamingTracker for partial JSON parsing

romanrizzi · romanrizzi · commit 4db7773f640c · 2025-05-05T16:12:42.000-03:00
diff --git a/lib/completions/anthropic_message_processor.rb b/lib/completions/anthropic_message_processor.rb
@@ -10,7 +10,7 @@ def initialize(name, id, partial_tool_calls: false)
       @raw_json = +""
       @tool_call = DiscourseAi::Completions::ToolCall.new(id: id, name: name, parameters: {})
       @streaming_parser =
-        DiscourseAi::Completions::ToolCallProgressTracker.new(self) if partial_tool_calls
+        DiscourseAi::Completions::JsonStreamingTracker.new(self) if partial_tool_calls
     end
 
     def append(json)
diff --git a/lib/completions/endpoints/anthropic.rb b/lib/completions/endpoints/anthropic.rb
@@ -117,6 +117,7 @@ def prepare_payload(prompt, model_params, dialect)
           if model_params[:response_format].present?
             prefilled_message << " " if !prefilled_message.empty?
             prefilled_message << "{"
+            @forced_json_through_prefill = true
           end
 
           if !prefilled_message.empty?
diff --git a/lib/completions/endpoints/aws_bedrock.rb b/lib/completions/endpoints/aws_bedrock.rb
@@ -144,6 +144,7 @@ def prepare_payload(prompt, model_params, dialect)
             if model_params[:response_format].present?
               prefilled_message << " " if !prefilled_message.empty?
               prefilled_message << "{"
+              @forced_json_through_prefill = true
             end
 
             if !prefilled_message.empty?
diff --git a/lib/completions/endpoints/base.rb b/lib/completions/endpoints/base.rb
@@ -73,6 +73,7 @@ def perform_completion!(
           LlmQuota.check_quotas!(@llm_model, user)
           start_time = Time.now
 
+          @forced_json_through_prefill = false
           @partial_tool_calls = partial_tool_calls
           @output_thinking = output_thinking
 
@@ -109,12 +110,12 @@ def perform_completion!(
           structured_output = nil
 
           if model_params[:response_format].present?
-            response_structure =
-              model_params[:response_format].dig(:json_schema, :schema, :required)
+            schema_properties =
+              model_params[:response_format].dig(:json_schema, :schema, :properties)
 
-            if response_structure.present?
+            if schema_properties.present?
               structured_output =
-                DiscourseAi::Completions::StructuredOutput.new(response_structure.map(&:to_sym))
+                DiscourseAi::Completions::StructuredOutput.new(schema_properties)
             end
           end
 
@@ -135,6 +136,10 @@ def perform_completion!(
 
             request = prepare_request(request_body)
 
+            # Some providers rely on prefill to return structured outputs, so the start
+            # of the JSON won't be included in the response. Supply it to keep JSON valid.
+            structured_output << +"{" if structured_output && @forced_json_through_prefill
+
             http.request(request) do |response|
               if response.code.to_i != 200
                 Rails.logger.error(
diff --git a/lib/completions/endpoints/canned_response.rb b/lib/completions/endpoints/canned_response.rb
@@ -90,11 +90,11 @@ def is_structured_output?(response)
         end
 
         def as_structured_output(response)
-          keys = model_params[:response_format].dig(:json_schema, :schema, :properties)&.keys
-          return response if keys.blank?
+          schema_properties = model_params[:response_format].dig(:json_schema, :schema, :properties)
+          return response if schema_properties.blank?
 
-          output = DiscourseAi::Completions::StructuredOutput.new(keys)
-          output << { keys.first => response }.to_json
+          output = DiscourseAi::Completions::StructuredOutput.new(schema_properties)
+          output << { schema_properties.keys.first => response }.to_json
 
           output
         end
diff --git a/lib/completions/json_streaming_tracker.rb b/lib/completions/json_streaming_tracker.rb
@@ -2,11 +2,11 @@
 
 module DiscourseAi
   module Completions
-    class ToolCallProgressTracker
-      attr_reader :current_key, :current_value, :tool_call
+    class JsonStreamingTracker
+      attr_reader :current_key, :current_value, :stream_consumer
 
-      def initialize(tool_call)
-        @tool_call = tool_call
+      def initialize(stream_consumer)
+        @stream_consumer = stream_consumer
         @current_key = nil
         @current_value = nil
         @parser = DiscourseAi::Completions::JsonStreamingParser.new
@@ -18,7 +18,7 @@ def initialize(tool_call)
 
         @parser.value do |v|
           if @current_key
-            tool_call.notify_progress(@current_key, v)
+            stream_consumer.notify_progress(@current_key, v)
             @current_key = nil
           end
         end
@@ -39,7 +39,7 @@ def <<(json)
 
         if @parser.state == :start_string && @current_key
           # this is is worth notifying
-          tool_call.notify_progress(@current_key, @parser.buf)
+          stream_consumer.notify_progress(@current_key, @parser.buf)
         end
 
         @current_key = nil if @parser.state == :end_value
diff --git a/lib/completions/nova_message_processor.rb b/lib/completions/nova_message_processor.rb
@@ -10,7 +10,7 @@ def initialize(name, id, partial_tool_calls: false)
       @raw_json = +""
       @tool_call = DiscourseAi::Completions::ToolCall.new(id: id, name: name, parameters: {})
       @streaming_parser =
-        DiscourseAi::Completions::ToolCallProgressTracker.new(self) if partial_tool_calls
+        DiscourseAi::Completions::JsonStreamingTracker.new(self) if partial_tool_calls
     end
 
     def append(json)
diff --git a/lib/completions/open_ai_message_processor.rb b/lib/completions/open_ai_message_processor.rb
@@ -59,7 +59,7 @@ def process_streamed_message(json)
         if id.present? && name.present?
           @tool_arguments = +""
           @tool = ToolCall.new(id: id, name: name)
-          @streaming_parser = ToolCallProgressTracker.new(self) if @partial_tool_calls
+          @streaming_parser = JsonStreamingTracker.new(self) if @partial_tool_calls
         end
 
         @tool_arguments << arguments.to_s
diff --git a/lib/completions/structured_output.rb b/lib/completions/structured_output.rb
@@ -3,91 +3,49 @@
 module DiscourseAi
   module Completions
     class StructuredOutput
-      def initialize(property_names)
-        @raw_response = +""
-        @state = :awaiting_key
-        @current_key = +""
-        @escape = false
-
-        @full_output =
-          property_names.reduce({}) do |memo, pn|
-            memo[pn.to_sym] = +""
-            memo
+      def initialize(json_schema_properties)
+        @property_names = json_schema_properties.keys.map(&:to_sym)
+        @property_cursors =
+          json_schema_properties.reduce({}) do |m, (k, prop)|
+            m[k.to_sym] = 0 if prop[:type] == "string"
+            m
           end
 
-        # Partial output is what we processed in the last chunk.
-        @partial_output_proto = @full_output.deep_dup
-        @last_chunk_output = @full_output.deep_dup
+        @tracked = {}
+
+        @partial_json_tracker = JsonStreamingTracker.new(self)
       end
 
-      attr_reader :full_output, :last_chunk_output
+      attr_reader :last_chunk_buffer
 
       def <<(raw)
-        @raw_response << raw
-
-        @last_chunk_output = @partial_output_proto.deep_dup
+        @partial_json_tracker << raw
+      end
 
-        raw.each_char do |char|
-          case @state
-          when :awaiting_key
-            if char == "\""
-              @current_key = +""
-              @state = :parsing_key
-              @escape = false
-            end
-          when :parsing_key
-            if char == "\""
-              @state = :awaiting_colon
-            else
-              @current_key << char
-            end
-          when :awaiting_colon
-            @state = :awaiting_value if char == ":"
-          when :awaiting_value
-            if char == '"'
-              @escape = false
-              @state = :parsing_value
-            end
-          when :parsing_value
-            if @escape
-              # Don't add escape sequence until we know what it is
-              unescaped = unescape_char(char)
-              @full_output[@current_key.to_sym] << unescaped
-              @last_chunk_output[@current_key.to_sym] << unescaped
+      def read_latest_buffered_chunk
+        @property_names.reduce({}) do |memo, pn|
+          if @tracked[pn].present?
+            # This means this property is a string and we want to return unread chunks.
+            if @property_cursors[pn].present?
+              unread = @tracked[pn][@property_cursors[pn]..]
 
-              @escape = false
-            elsif char == "\\"
-              @escape = true
-            elsif char == "\""
-              @state = :awaiting_key_or_end
+              memo[pn] = unread if unread.present?
+              @property_cursors[pn] = @tracked[pn].length
             else
-              @full_output[@current_key.to_sym] << char
-              @last_chunk_output[@current_key.to_sym] << char
+              # Ints and bools are always returned as is.
+              memo[pn] = @tracked[pn]
             end
-          when :awaiting_key_or_end
-            @state = :awaiting_key if char == ","
-            # End of object or whitespace ignored here
-          else
-            next
           end
+
+          memo
         end
       end
 
-      private
-
-      def unescape_char(char)
-        chars = {
-          '"' => '"',
-          '\\' => '\\',
-          "/" => "/",
-          "b" => "\b",
-          "f" => "\f",
-          "n" => "\n",
-          "r" => "\r",
-          "t" => "\t",
-        }
+      def notify_progress(key, value)
+        key_sym = key.to_sym
+        return if !@property_names.include?(key_sym)
 
-        chars[char] || char
+        @tracked[key_sym] = value
       end
     end
   end
diff --git a/lib/personas/bot.rb b/lib/personas/bot.rb
@@ -152,7 +152,7 @@ def reply(context, llm_args: {}, &update_blk)
                       current_thinking << partial
                     end
                   elsif partial.is_a?(DiscourseAi::Completions::StructuredOutput)
-                    update_blk.call(partial.last_chunk_output, cancel, nil, :structured_output)
+                    update_blk.call(partial, cancel, nil, :structured_output)
                   else
                     update_blk.call(partial, cancel)
                   end
diff --git a/lib/summarization/fold_content.rb b/lib/summarization/fold_content.rb
@@ -115,7 +115,8 @@ def fold(items, user, &on_partial_blk)
           Proc.new do |partial, cancel, _, type|
             if type == :structured_output
               json_summary_schema_key = bot.persona.response_format&.first.to_h
-              partial_summary = partial[json_summary_schema_key[:key].to_sym]
+              partial_summary =
+                partial.read_latest_buffered_chunk[json_summary_schema_key[:key].to_sym]
 
               summary << partial_summary
               on_partial_blk.call(partial_summary, cancel) if on_partial_blk
diff --git a/spec/lib/completions/endpoints/anthropic_spec.rb b/spec/lib/completions/endpoints/anthropic_spec.rb
@@ -845,7 +845,7 @@
         response_format: schema,
       ) { |partial, cancel| structured_output = partial }
 
-      expect(structured_output.full_output).to eq({ key: "Hello!" })
+      expect(structured_output.read_latest_buffered_chunk).to eq({ key: "Hello!" })
 
       expected_body = {
         model: "claude-3-opus-20240229",
diff --git a/spec/lib/completions/endpoints/aws_bedrock_spec.rb b/spec/lib/completions/endpoints/aws_bedrock_spec.rb
@@ -607,7 +607,7 @@ def encode_message(message)
         }
         expect(JSON.parse(request.body)).to eq(expected)
 
-        expect(structured_output.full_output).to eq({ key: "Hello!" })
+        expect(structured_output.read_latest_buffered_chunk).to eq({ key: "Hello!" })
       end
     end
   end
diff --git a/spec/lib/completions/endpoints/cohere_spec.rb b/spec/lib/completions/endpoints/cohere_spec.rb
@@ -330,7 +330,7 @@
 
     body = <<~TEXT
       {"is_finished":false,"event_type":"stream-start","generation_id":"eb889b0f-c27d-45ea-98cf-567bdb7fc8bf"}
-      {"is_finished":false,"event_type":"text-generation","text":"\\""}
+      {"is_finished":false,"event_type":"text-generation","text":"{\\""}
       {"is_finished":false,"event_type":"text-generation","text":"key"}
       {"is_finished":false,"event_type":"text-generation","text":"\\":\\""}
       {"is_finished":false,"event_type":"text-generation","text":"Hello!"}
@@ -366,6 +366,6 @@
     )
     expect(parsed_body[:message]).to eq("user1: thanks")
 
-    expect(structured_output.full_output).to eq({ key: "Hello!" })
+    expect(structured_output.read_latest_buffered_chunk).to eq({ key: "Hello!" })
   end
 end
diff --git a/spec/lib/completions/endpoints/gemini_spec.rb b/spec/lib/completions/endpoints/gemini_spec.rb
@@ -487,7 +487,7 @@ def tool_response
         structured_response = partial
       end
 
-      expect(structured_response.full_output).to eq({ key: "Hello!" })
+      expect(structured_response.read_latest_buffered_chunk).to eq({ key: "Hello!" })
 
       parsed = JSON.parse(req_body, symbolize_names: true)
 
diff --git a/spec/lib/completions/structured_output_spec.rb b/spec/lib/completions/structured_output_spec.rb
@@ -1,32 +1,69 @@
 # frozen_string_literal: true
 
 RSpec.describe DiscourseAi::Completions::StructuredOutput do
-  subject(:structured_output) { described_class.new(%i[message status]) }
+  subject(:structured_output) do
+    described_class.new(
+      {
+        message: {
+          type: "string",
+        },
+        bool: {
+          type: "boolean",
+        },
+        number: {
+          type: "integer",
+        },
+        status: {
+          type: "string",
+        },
+      },
+    )
+  end
 
   describe "Parsing structured output on the fly" do
     it "acts as a buffer for an streamed JSON" do
-      chunks = ["{\"message\": \"Line 1\\", "nLine 2\", \"status\": \"o", "\\\"k\\\"\"}"]
+      chunks = [
+        +"{\"message\": \"Line 1\\n",
+        +"Line 2\\n",
+        +"Line 3\", ",
+        +"\"bool\": true,",
+        +"\"number\": 4",
+        +"2,",
+        +"\"status\": \"o",
+        +"\\\"k\\\"\"}",
+      ]
 
       structured_output << chunks[0]
-
-      expect(structured_output.full_output[:status]).to be_empty
-      expect(structured_output.full_output[:message]).to eq("Line 1")
-      expect(structured_output.last_chunk_output[:status]).to be_empty
-      expect(structured_output.last_chunk_output[:message]).to eq("Line 1")
+      expect(structured_output.read_latest_buffered_chunk).to eq({ message: "Line 1\n" })
 
       structured_output << chunks[1]
-
-      expect(structured_output.full_output[:status]).to eq("o")
-      expect(structured_output.full_output[:message]).to eq("Line 1\nLine 2")
-      expect(structured_output.last_chunk_output[:status]).to eq("o")
-      expect(structured_output.last_chunk_output[:message]).to eq("\nLine 2")
+      expect(structured_output.read_latest_buffered_chunk).to eq({ message: "Line 2\n" })
 
       structured_output << chunks[2]
+      expect(structured_output.read_latest_buffered_chunk).to eq({ message: "Line 3" })
+
+      structured_output << chunks[3]
+      expect(structured_output.read_latest_buffered_chunk).to eq({ bool: true })
+
+      # Waiting for number to be fully buffered.
+      structured_output << chunks[4]
+      expect(structured_output.read_latest_buffered_chunk).to eq({ bool: true })
+
+      structured_output << chunks[5]
+      expect(structured_output.read_latest_buffered_chunk).to eq({ bool: true, number: 42 })
+
+      structured_output << chunks[6]
+      expect(structured_output.read_latest_buffered_chunk).to eq(
+        { bool: true, number: 42, status: "o" },
+      )
+
+      structured_output << chunks[7]
+      expect(structured_output.read_latest_buffered_chunk).to eq(
+        { bool: true, number: 42, status: "\"k\"" },
+      )
 
-      expect(structured_output.full_output[:status]).to eq("o\"k\"")
-      expect(structured_output.full_output[:message]).to eq("Line 1\nLine 2")
-      expect(structured_output.last_chunk_output[:status]).to eq("\"k\"")
-      expect(structured_output.last_chunk_output[:message]).to be_empty
+      # No partial string left to read.
+      expect(structured_output.read_latest_buffered_chunk).to eq({ bool: true, number: 42 })
     end
   end
 end
