Improve JSON parsing.

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)}")

requirements.txt

@@ -10,3 +10,4 @@ wheel
 watchdog
 toml
 types-toml
+multiline

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):