From 788b3015bc628130718e6db452da28fc3316ca5d Mon Sep 17 00:00:00 2001
From: pablonyx <pablo@danswer.ai>
Date: Mon, 16 Dec 2024 12:37:47 -0800
Subject: [PATCH] fix single quote block in llm answer (#3139)

---
 .../stream_processing/citation_processing.py  |  2 +-
 .../test_citation_processing.py               | 20 +++++++++++++++++++
 2 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/backend/onyx/chat/stream_processing/citation_processing.py b/backend/onyx/chat/stream_processing/citation_processing.py
index e66be3fb8..cd159dd15 100644
--- a/backend/onyx/chat/stream_processing/citation_processing.py
+++ b/backend/onyx/chat/stream_processing/citation_processing.py
@@ -65,7 +65,7 @@ class CitationProcessor:
         # Handle code blocks without language tags
         if "`" in self.curr_segment:
             if self.curr_segment.endswith("`"):
-                return
+                pass
             elif "```" in self.curr_segment:
                 piece_that_comes_after = self.curr_segment.split("```")[1][0]
                 if piece_that_comes_after == "\n" and in_code_block(self.llm_out):
diff --git a/backend/tests/unit/onyx/chat/stream_processing/test_citation_processing.py b/backend/tests/unit/onyx/chat/stream_processing/test_citation_processing.py
index a40f0392e..dcc960790 100644
--- a/backend/tests/unit/onyx/chat/stream_processing/test_citation_processing.py
+++ b/backend/tests/unit/onyx/chat/stream_processing/test_citation_processing.py
@@ -376,6 +376,26 @@ def process_text(
             "The code demonstrates variable assignment.",
             [],
         ),
+        (
+            "Long JSON string in code block",
+            [
+                "```json\n{",
+                '"name": "John Doe",',
+                '"age": 30,',
+                '"city": "New York",',
+                '"hobbies": ["reading", "swimming", "cycling"],',
+                '"education": {',
+                '    "degree": "Bachelor\'s",',
+                '    "major": "Computer Science",',
+                '    "university": "Example University"',
+                "}",
+                "}\n```",
+            ],
+            '```json\n{"name": "John Doe","age": 30,"city": "New York","hobbies": '
+            '["reading", "swimming", "cycling"],"education": {    '
+            '"degree": "Bachelor\'s",    "major": "Computer Science",    "university": "Example University"}}\n```',
+            [],
+        ),
         (
             "Citation as a single token",
             [