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