FIX: Retry parsing escaped inner JSON to handle control chars.

lib/completions/json_streaming_tracker.rb

@@ -28,17 +28,21 @@ def broken?
         @broken
       end
 
-      def <<(json)
+      def <<(raw_json)
         # llm could send broken json
         # in that case just deal with it later
         # don't stream
         return if @broken
 
         begin
-          @parser << json
+          @parser << raw_json
         rescue DiscourseAi::Completions::ParserError
-          @broken = true
-          return
+          # Note: We're parsing JSON content that was itself embedded as a string inside another JSON object.
+          # During the outer JSON.parse, any escaped control characters (like "\\n") are unescaped to real characters ("\n"),
+          # which corrupts the inner JSON structure when passed to the parser here.
+          # To handle this, we retry parsing with the string JSON-escaped again (`.dump[1..-2]`) if the first attempt fails.
+          try_escape_and_parse(raw_json)
+          return if @broken
         end
 
         if @parser.state == :start_string && @current_key
@@ -48,6 +52,20 @@ def <<(json)
 
         @current_key = nil if @parser.state == :end_value
       end
+
+      private
+
+      def try_escape_and_parse(raw_json)
+        if raw_json.blank? || !raw_json.is_a?(String)
+          @broken = true
+          return
+        end
+        # Escape the string as JSON and remove surrounding quotes
+        escaped_json = raw_json.dump[1..-2]
+        @parser << escaped_json
+      rescue DiscourseAi::Completions::ParserError
+        @broken = true
+      end
     end
   end
 end

spec/lib/completions/endpoints/anthropic_spec.rb

@@ -809,6 +809,9 @@
       event: content_block_delta
       data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "Hello!"}}
 
+      event: content_block_delta
+      data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "\\n there"}}
+
       event: content_block_delta
       data: {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "\\"}"}}
 
@@ -845,7 +848,7 @@
         response_format: schema,
       ) { |partial, cancel| structured_output = partial }
 
-      expect(structured_output.read_buffered_property(:key)).to eq("Hello!")
+      expect(structured_output.read_buffered_property(:key)).to eq("Hello!\n there")
 
       expected_body = {
         model: "claude-3-opus-20240229",

spec/lib/completions/endpoints/aws_bedrock_spec.rb

@@ -574,6 +574,7 @@ def encode_message(message)
           { type: "content_block_delta", delta: { text: "key" } },
           { type: "content_block_delta", delta: { text: "\":\"" } },
           { type: "content_block_delta", delta: { text: "Hello!" } },
+          { type: "content_block_delta", delta: { text: "\n There" } },
           { type: "content_block_delta", delta: { text: "\"}" } },
           { type: "message_delta", delta: { usage: { output_tokens: 25 } } },
         ].map { |message| encode_message(message) }
@@ -607,7 +608,7 @@ def encode_message(message)
         }
         expect(JSON.parse(request.body)).to eq(expected)
 
-        expect(structured_output.read_buffered_property(:key)).to eq("Hello!")
+        expect(structured_output.read_buffered_property(:key)).to eq("Hello!\n There")
       end
     end
   end

spec/lib/completions/endpoints/cohere_spec.rb

@@ -334,8 +334,9 @@
       {"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!"}
+      {"is_finished":false,"event_type":"text-generation","text":"\\n there"}
       {"is_finished":false,"event_type":"text-generation","text":"\\"}"}|
-      {"is_finished":true,"event_type":"stream-end","response":{"response_id":"d235db17-8555-493b-8d91-e601f76de3f9","text":"{\\"key\\":\\"Hello!\\"}","generation_id":"eb889b0f-c27d-45ea-98cf-567bdb7fc8bf","chat_history":[{"role":"USER","message":"user1: hello"},{"role":"CHATBOT","message":"hi user"},{"role":"USER","message":"user1: thanks"},{"role":"CHATBOT","message":"You're welcome! Is there anything else I can help you with?"}],"token_count":{"prompt_tokens":29,"response_tokens":14,"total_tokens":43,"billed_tokens":28},"meta":{"api_version":{"version":"1"},"billed_units":{"input_tokens":14,"output_tokens":14}}},"finish_reason":"COMPLETE"}
+      {"is_finished":true,"event_type":"stream-end","response":{"response_id":"d235db17-8555-493b-8d91-e601f76de3f9","text":"{\\"key\\":\\"Hello! \\n there\\"}","generation_id":"eb889b0f-c27d-45ea-98cf-567bdb7fc8bf","chat_history":[{"role":"USER","message":"user1: hello"},{"role":"CHATBOT","message":"hi user"},{"role":"USER","message":"user1: thanks"},{"role":"CHATBOT","message":"You're welcome! Is there anything else I can help you with?"}],"token_count":{"prompt_tokens":29,"response_tokens":14,"total_tokens":43,"billed_tokens":28},"meta":{"api_version":{"version":"1"},"billed_units":{"input_tokens":14,"output_tokens":14}}},"finish_reason":"COMPLETE"}
     TEXT
 
     parsed_body = nil
@@ -366,6 +367,6 @@
     )
     expect(parsed_body[:message]).to eq("user1: thanks")
 
-    expect(structured_output.read_buffered_property(:key)).to eq("Hello!")
+    expect(structured_output.read_buffered_property(:key)).to eq("Hello!\n there")
   end
 end

spec/lib/completions/endpoints/gemini_spec.rb

@@ -524,6 +524,9 @@ def tool_response
               key: {
                 type: "string",
               },
+              num: {
+                type: "integer",
+              },
             },
             required: ["key"],
             additionalProperties: false,
@@ -541,7 +544,19 @@ def tool_response
 
         data: {"candidates": [{"content": {"parts": [{"text": "Hello!"}],"role": "model"},"finishReason": "STOP"}],"usageMetadata": {"promptTokenCount": 399,"candidatesTokenCount": 191,"totalTokenCount": 590},"modelVersion": "gemini-1.5-pro-002"}
 
-        data: {"candidates": [{"content": {"parts": [{"text": "\\"}"}],"role": "model"},"finishReason": "STOP"}],"usageMetadata": {"promptTokenCount": 399,"candidatesTokenCount": 191,"totalTokenCount": 590},"modelVersion": "gemini-1.5-pro-002"}
+        data: {"candidates": [{"content": {"parts": [{"text": "\\n there"}],"role": "model"},"finishReason": "STOP"}],"usageMetadata": {"promptTokenCount": 399,"candidatesTokenCount": 191,"totalTokenCount": 590},"modelVersion": "gemini-1.5-pro-002"}
+
+        data: {"candidates": [{"content": {"parts": [{"text": "\\","}],"role": "model"},"finishReason": "STOP"}],"usageMetadata": {"promptTokenCount": 399,"candidatesTokenCount": 191,"totalTokenCount": 590},"modelVersion": "gemini-1.5-pro-002"}
+
+        data: {"candidates": [{"content": {"parts": [{"text": "\\""}],"role": "model"}}],"usageMetadata": {"promptTokenCount": 399,"totalTokenCount": 399},"modelVersion": "gemini-1.5-pro-002"}
+
+        data: {"candidates": [{"content": {"parts": [{"text": "num"}],"role": "model"},"finishReason": "STOP"}],"usageMetadata": {"promptTokenCount": 399,"candidatesTokenCount": 191,"totalTokenCount": 590},"modelVersion": "gemini-1.5-pro-002"}
+
+        data: {"candidates": [{"content": {"parts": [{"text": "\\":"}],"role": "model"},"safetyRatings": [{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"}]}],"usageMetadata": {"promptTokenCount": 399,"totalTokenCount": 399},"modelVersion": "gemini-1.5-pro-002"}
+
+        data: {"candidates": [{"content": {"parts": [{"text": "42"}],"role": "model"},"finishReason": "STOP"}],"usageMetadata": {"promptTokenCount": 399,"candidatesTokenCount": 191,"totalTokenCount": 590},"modelVersion": "gemini-1.5-pro-002"}
+
+        data: {"candidates": [{"content": {"parts": [{"text": "}"}],"role": "model"},"finishReason": "STOP"}],"usageMetadata": {"promptTokenCount": 399,"candidatesTokenCount": 191,"totalTokenCount": 590},"modelVersion": "gemini-1.5-pro-002"}
 
         data: {"candidates": [{"finishReason": "MALFORMED_FUNCTION_CALL"}],"usageMetadata": {"promptTokenCount": 399,"candidatesTokenCount": 191,"totalTokenCount": 590},"modelVersion": "gemini-1.5-pro-002"}
 
@@ -565,7 +580,8 @@ def tool_response
         structured_response = partial
       end
 
-      expect(structured_response.read_buffered_property(:key)).to eq("Hello!")
+      expect(structured_response.read_buffered_property(:key)).to eq("Hello!\n there")
+      expect(structured_response.read_buffered_property(:num)).to eq(42)
 
       parsed = JSON.parse(req_body, symbolize_names: true)