From d72ad84f8f57d1575e6d3611a885fbaab84a265a Mon Sep 17 00:00:00 2001 From: Roman Rizzi Date: Wed, 21 May 2025 11:25:59 -0300 Subject: [PATCH] FIX: Retry parsing escaped inner JSON to handle control chars. (#1357) The structured output JSON comes embedded inside the API response, which is also a JSON. Since we have to parse the response to process it, any control characters inside the structured output are unescaped into regular characters, leading to invalid JSON and breaking during parsing. This change adds a retry mechanism that escapes the string again if parsing fails, preventing the parser from breaking on malformed input and working around this issue. For example: ``` original = '{ "a": "{\\"key\\":\\"value with \\n newline\\"}" }' JSON.parse(original) => { "a" => "{\"key\":\"value with \n newline\"}" } # At this point, the inner JSON string contains an actual newline. ``` --- lib/completions/json_streaming_tracker.rb | 26 ++++++++++++++++--- .../completions/endpoints/anthropic_spec.rb | 5 +++- .../completions/endpoints/aws_bedrock_spec.rb | 3 ++- spec/lib/completions/endpoints/cohere_spec.rb | 5 ++-- spec/lib/completions/endpoints/gemini_spec.rb | 20 ++++++++++++-- 5 files changed, 49 insertions(+), 10 deletions(-) diff --git a/lib/completions/json_streaming_tracker.rb b/lib/completions/json_streaming_tracker.rb index 849fd8ac..a5bdfff4 100644 --- a/lib/completions/json_streaming_tracker.rb +++ b/lib/completions/json_streaming_tracker.rb @@ -28,17 +28,21 @@ module DiscourseAi @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 @@ module DiscourseAi @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 diff --git a/spec/lib/completions/endpoints/anthropic_spec.rb b/spec/lib/completions/endpoints/anthropic_spec.rb index 96a4ae07..f0e53bea 100644 --- a/spec/lib/completions/endpoints/anthropic_spec.rb +++ b/spec/lib/completions/endpoints/anthropic_spec.rb @@ -809,6 +809,9 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Anthropic do 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 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Anthropic do 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", diff --git a/spec/lib/completions/endpoints/aws_bedrock_spec.rb b/spec/lib/completions/endpoints/aws_bedrock_spec.rb index aa86548e..bd60f988 100644 --- a/spec/lib/completions/endpoints/aws_bedrock_spec.rb +++ b/spec/lib/completions/endpoints/aws_bedrock_spec.rb @@ -574,6 +574,7 @@ RSpec.describe DiscourseAi::Completions::Endpoints::AwsBedrock do { 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 @@ RSpec.describe DiscourseAi::Completions::Endpoints::AwsBedrock do } 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 diff --git a/spec/lib/completions/endpoints/cohere_spec.rb b/spec/lib/completions/endpoints/cohere_spec.rb index 546922b3..c4fb06b6 100644 --- a/spec/lib/completions/endpoints/cohere_spec.rb +++ b/spec/lib/completions/endpoints/cohere_spec.rb @@ -334,8 +334,9 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Cohere do {"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 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Cohere do ) 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 diff --git a/spec/lib/completions/endpoints/gemini_spec.rb b/spec/lib/completions/endpoints/gemini_spec.rb index aea97bb8..59d76cac 100644 --- a/spec/lib/completions/endpoints/gemini_spec.rb +++ b/spec/lib/completions/endpoints/gemini_spec.rb @@ -524,6 +524,9 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Gemini do key: { type: "string", }, + num: { + type: "integer", + }, }, required: ["key"], additionalProperties: false, @@ -541,7 +544,19 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Gemini do 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 @@ RSpec.describe DiscourseAi::Completions::Endpoints::Gemini do 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)