FIX: Structured output discrepancies. (#1340)

This change fixes two bugs and adds a safeguard.

The first issue is that the schema Gemini expected differed from the one sent, resulting in 400 errors when performing completions.

The second issue was that creating a new persona won't define a method
for `response_format`. This has to be explicitly defined when we wrap it inside the Persona class. Also, There was a mismatch between the default value and what we stored in the DB. Some parts of the code expected symbols as keys and others as strings.

Finally, we add a safeguard when, even if asked to, the model refuses to reply with a valid JSON. In this case, we are making a best-effort to recover and stream the raw response.

Author: romanrizzi

app/models/ai_persona.rb

@@ -266,6 +266,7 @@ def class_instance
       define_method(:top_p) { @ai_persona&.top_p }
       define_method(:system_prompt) { @ai_persona&.system_prompt || "You are a helpful bot." }
       define_method(:uploads) { @ai_persona&.uploads }
+      define_method(:response_format) { @ai_persona&.response_format }
       define_method(:examples) { @ai_persona&.examples }
     end
   end

lib/completions/endpoints/gemini.rb

@@ -87,9 +87,13 @@ def prepare_payload(prompt, model_params, dialect)
           if model_params.present?
             payload[:generationConfig].merge!(model_params.except(:response_format))
 
-            if model_params[:response_format].present?
-              # https://ai.google.dev/api/generate-content#generationconfig
-              payload[:generationConfig][:responseSchema] = model_params[:response_format]
+            # https://ai.google.dev/api/generate-content#generationconfig
+            gemini_schema = model_params[:response_format].dig(:json_schema, :schema)
+
+            if gemini_schema.present?
+              payload[:generationConfig][:responseSchema] = gemini_schema.except(
+                :additionalProperties,
+              )
               payload[:generationConfig][:responseMimeType] = "application/json"
             end
           end

lib/completions/json_streaming_tracker.rb

@@ -24,6 +24,10 @@ def initialize(stream_consumer)
         end
       end
 
+      def broken?
+        @broken
+      end
+
       def <<(json)
         # llm could send broken json
         # in that case just deal with it later

lib/completions/structured_output.rb

@@ -13,31 +13,40 @@ def initialize(json_schema_properties)
 
         @tracked = {}
 
+        @raw_response = +""
+        @raw_cursor = 0
+
         @partial_json_tracker = JsonStreamingTracker.new(self)
       end
 
       attr_reader :last_chunk_buffer
 
       def <<(raw)
+        @raw_response << raw
         @partial_json_tracker << raw
       end
 
-      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]..]
-
-              memo[pn] = unread if unread.present?
-              @property_cursors[pn] = @tracked[pn].length
-            else
-              # Ints and bools are always returned as is.
-              memo[pn] = @tracked[pn]
-            end
-          end
+      def read_buffered_property(prop_name)
+        # Safeguard: If the model is misbehaving and generating something that's not a JSON,
+        # treat response as a normal string.
+        # This is a best-effort to recover from an unexpected scenario.
+        if @partial_json_tracker.broken?
+          unread_chunk = @raw_response[@raw_cursor..]
+          @raw_cursor = @raw_response.length
+          return unread_chunk
+        end
 
-          memo
+        # Maybe we haven't read that part of the JSON yet.
+        return nil if @tracked[prop_name].blank?
+
+        # This means this property is a string and we want to return unread chunks.
+        if @property_cursors[prop_name].present?
+          unread = @tracked[prop_name][@property_cursors[prop_name]..]
+          @property_cursors[prop_name] = @tracked[prop_name].length
+          unread
+        else
+          # Ints and bools are always returned as is.
+          @tracked[prop_name]
         end
       end
 

lib/personas/bot.rb

@@ -316,7 +316,7 @@ def build_json_schema(response_format)
           response_format
             .to_a
             .reduce({}) do |memo, format|
-              memo[format[:key].to_sym] = { type: format[:type] }
+              memo[format["key"].to_sym] = { type: format["type"] }
               memo
             end
 

lib/personas/short_summarizer.rb

@@ -33,7 +33,7 @@ def system_prompt
       end
 
       def response_format
-        [{ key: "summary", type: "string" }]
+        [{ "key" => "summary", "type" => "string" }]
       end
     end
   end

lib/personas/summarizer.rb

@@ -34,7 +34,7 @@ def system_prompt
       end
 
       def response_format
-        [{ key: "summary", type: "string" }]
+        [{ "key" => "summary", "type" => "string" }]
       end
 
       def examples

lib/summarization/fold_content.rb

@@ -116,7 +116,7 @@ def fold(items, user, &on_partial_blk)
             if type == :structured_output
               json_summary_schema_key = bot.persona.response_format&.first.to_h
               partial_summary =
-                partial.read_latest_buffered_chunk[json_summary_schema_key[:key].to_sym]
+                partial.read_buffered_property(json_summary_schema_key["key"]&.to_sym)
 
               if partial_summary.present?
                 summary << partial_summary

spec/lib/completions/endpoints/anthropic_spec.rb

@@ -845,7 +845,7 @@
         response_format: schema,
       ) { |partial, cancel| structured_output = partial }
 
-      expect(structured_output.read_latest_buffered_chunk).to eq({ key: "Hello!" })
+      expect(structured_output.read_buffered_property(:key)).to eq("Hello!")
 
       expected_body = {
         model: "claude-3-opus-20240229",

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.read_latest_buffered_chunk).to eq({ key: "Hello!" })
+        expect(structured_output.read_buffered_property(:key)).to eq("Hello!")
       end
     end
   end

spec/lib/completions/endpoints/cohere_spec.rb

@@ -366,6 +366,6 @@
     )
     expect(parsed_body[:message]).to eq("user1: thanks")
 
-    expect(structured_output.read_latest_buffered_chunk).to eq({ key: "Hello!" })
+    expect(structured_output.read_buffered_property(:key)).to eq("Hello!")
   end
 end

spec/lib/completions/endpoints/gemini_spec.rb

@@ -565,12 +565,14 @@ def tool_response
         structured_response = partial
       end
 
-      expect(structured_response.read_latest_buffered_chunk).to eq({ key: "Hello!" })
+      expect(structured_response.read_buffered_property(:key)).to eq("Hello!")
 
       parsed = JSON.parse(req_body, symbolize_names: true)
 
       # Verify that schema is passed following Gemini API specs.
-      expect(parsed.dig(:generationConfig, :responseSchema)).to eq(schema)
+      expect(parsed.dig(:generationConfig, :responseSchema)).to eq(
+        schema.dig(:json_schema, :schema).except(:additionalProperties),
+      )
       expect(parsed.dig(:generationConfig, :responseMimeType)).to eq("application/json")
     end
   end

spec/lib/completions/structured_output_spec.rb

@@ -34,36 +34,50 @@
       ]
 
       structured_output << chunks[0]
-      expect(structured_output.read_latest_buffered_chunk).to eq({ message: "Line 1\n" })
+      expect(structured_output.read_buffered_property(:message)).to eq("Line 1\n")
 
       structured_output << chunks[1]
-      expect(structured_output.read_latest_buffered_chunk).to eq({ message: "Line 2\n" })
+      expect(structured_output.read_buffered_property(:message)).to eq("Line 2\n")
 
       structured_output << chunks[2]
-      expect(structured_output.read_latest_buffered_chunk).to eq({ message: "Line 3" })
+      expect(structured_output.read_buffered_property(:message)).to eq("Line 3")
 
       structured_output << chunks[3]
-      expect(structured_output.read_latest_buffered_chunk).to eq({ bool: true })
+      expect(structured_output.read_buffered_property(:bool)).to eq(true)
 
       # Waiting for number to be fully buffered.
       structured_output << chunks[4]
-      expect(structured_output.read_latest_buffered_chunk).to eq({ bool: true })
+      expect(structured_output.read_buffered_property(:bool)).to eq(true)
+      expect(structured_output.read_buffered_property(:number)).to be_nil
 
       structured_output << chunks[5]
-      expect(structured_output.read_latest_buffered_chunk).to eq({ bool: true, number: 42 })
+      expect(structured_output.read_buffered_property(:number)).to eq(42)
 
       structured_output << chunks[6]
-      expect(structured_output.read_latest_buffered_chunk).to eq(
-        { bool: true, number: 42, status: "o" },
-      )
+      expect(structured_output.read_buffered_property(:number)).to eq(42)
+      expect(structured_output.read_buffered_property(:bool)).to eq(true)
+      expect(structured_output.read_buffered_property(:status)).to eq("o")
 
       structured_output << chunks[7]
-      expect(structured_output.read_latest_buffered_chunk).to eq(
-        { bool: true, number: 42, status: "\"k\"" },
-      )
+      expect(structured_output.read_buffered_property(:status)).to eq("\"k\"")
 
       # No partial string left to read.
-      expect(structured_output.read_latest_buffered_chunk).to eq({ bool: true, number: 42 })
+      expect(structured_output.read_buffered_property(:status)).to eq("")
+    end
+  end
+
+  describe "dealing with non-JSON responses" do
+    it "treat it as plain text once we determined it's invalid JSON" do
+      chunks = [+"I'm not", +"a", +"JSON :)"]
+
+      structured_output << chunks[0]
+      expect(structured_output.read_buffered_property(nil)).to eq("I'm not")
+
+      structured_output << chunks[1]
+      expect(structured_output.read_buffered_property(nil)).to eq("a")
+
+      structured_output << chunks[2]
+      expect(structured_output.read_buffered_property(nil)).to eq("JSON :)")
     end
   end
 end