From ddc44bb9da7b856e66d673ed46e419147e9a8264 Mon Sep 17 00:00:00 2001
From: Oleksandr Kozachuk <ddeus.lp@mailnull.com>
Date: Wed, 29 Mar 2023 19:22:54 +0200
Subject: [PATCH] Improve JSON parsing.

---
 fjerkroa_bot/ai_responder.py | 20 ++++++++++++++++----
 requirements.txt             |  1 +
 tests/test_main.py           |  5 +++--
 3 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/fjerkroa_bot/ai_responder.py b/fjerkroa_bot/ai_responder.py
index fe51e37..ec74fc8 100644
--- a/fjerkroa_bot/ai_responder.py
+++ b/fjerkroa_bot/ai_responder.py
@@ -1,4 +1,5 @@
 import json
+import multiline
 import openai
 import aiohttp
 import logging
@@ -17,6 +18,17 @@ def pp(*args, **kw):
     return pformat(*args, **kw)
 
 
+def parse_response(content: str) -> Dict:
+    content = content.strip()
+    try:
+        return json.loads(content)
+    except Exception:
+        try:
+            return multiline.loads(content, multiline=True)
+        except Exception as err:
+            raise err
+
+
 class AIMessageBase(object):
     def __init__(self) -> None:
         pass
@@ -180,15 +192,15 @@ class AIResponder(object):
             if answer is None:
                 continue
             try:
-                response = json.loads(answer['content'])
+                response = parse_response(answer['content'])
             except Exception as err:
                 logging.warning(f"failed to parse the answer: {pp(err)}\n{repr(answer['content'])}")
                 answer['content'] = await self.fix(answer['content'])
                 try:
-                    response = json.loads(answer['content'])
+                    response = parse_response(answer['content'])
                 except Exception as err:
-                    logging.error(f"failed to parse the answer: {pp(err)}\n{repr(answer['content'])}")
-                    return AIResponse(None, False, f"ERROR: I could not parse this answer: {repr(answer['content'])}", None, False)
+                    logging.error(f"failed to parse the fixed answer: {pp(err)}\n{repr(answer['content'])}")
+                    continue
             if 'hack' not in response or type(response.get('picture', None)) not in (type(None), str):
                 continue
             logging.info(f"got this answer:\n{pp(response)}")
diff --git a/requirements.txt b/requirements.txt
index 7568841..64854f0 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,3 +10,4 @@ wheel
 watchdog
 toml
 types-toml
+multiline
diff --git a/tests/test_main.py b/tests/test_main.py
index 6ac3ca0..af878f8 100644
--- a/tests/test_main.py
+++ b/tests/test_main.py
@@ -5,6 +5,7 @@ import json
 import toml
 import openai
 import logging
+import pytest
 from unittest.mock import Mock, PropertyMock, MagicMock, AsyncMock, patch, mock_open, ANY
 from fjerkroa_bot import FjerkroaBot
 from discord import User, Message, TextChannel
@@ -113,8 +114,8 @@ class TestFunctionality(TestBotBase):
             return {'choices': [{'message': {'content': '{ "test": 3 ]'}}]}
         message = self.create_message("Hello there! How are you?")
         with patch.object(openai.ChatCompletion, 'acreate', new=acreate):
-            await self.bot.on_message(message)
-        self.bot.staff_channel.send.assert_called_once_with("ERROR: I could not parse this answer: '{ \"test\": 3 ]'", suppress_embeds=True)
+            with pytest.raises(RuntimeError, match="Failed to generate answer after multiple retries"):
+                await self.bot.on_message(message)
 
     async def test_on_message_event4(self) -> None:
         async def acreate(*a, **kw):