diff --git a/fjerkroa_bot/ai_responder.py b/fjerkroa_bot/ai_responder.py index cbf0978..b57d4df 100644 --- a/fjerkroa_bot/ai_responder.py +++ b/fjerkroa_bot/ai_responder.py @@ -32,28 +32,21 @@ def parse_response(content: str) -> Dict: def parse_maybe_json(json_string): if json_string is None: return None - if isinstance(json_string, list): - return ' '.join([str(x) for x in json_string]) - if isinstance(json_string, dict): - return ' '.join([str(x) for x in json_string.values()]) + if isinstance(json_string, (list, dict)): + return ' '.join(map(str, (json_string.values() if isinstance(json_string, dict) else json_string))) json_string = str(json_string).strip() try: parsed_json = parse_response(json_string) except Exception: - if json_string.startswith('{') and json_string.endswith('}'): - return parse_maybe_json(json_string[1:-1]) + for b, e in [('{', '}'), ('[', ']')]: + if json_string.startswith(b) and json_string.endswith(e): + return parse_maybe_json(json_string[1:-1]) return json_string - + if isinstance(parsed_json, str): + return parsed_json if isinstance(parsed_json, (list, dict)): - concatenated_values = [] - for value in parsed_json.values() if isinstance(parsed_json, dict) else parsed_json: - concatenated_values.append(str(value)) - return '\n'.join(concatenated_values) - - result = str(parsed_json) - if result.lower() in ('', 'none', 'null', '"none"', '"null"', "'none'", "'null'"): - return None - return result + return '\n'.join(map(str, (parsed_json.values() if isinstance(parsed_json, dict) else parsed_json))) + return str(parsed_json) class AIMessageBase(object): diff --git a/tests/test_main.py b/tests/test_main.py index 73b4265..6e57881 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -77,6 +77,15 @@ class TestFunctionality(TestBotBase): json_struct = '{"This is a string."}' expected_output = 'This is a string.' self.assertEqual(parse_maybe_json(json_struct), expected_output) + json_struct = '["This is a string."]' + expected_output = 'This is a string.' + self.assertEqual(parse_maybe_json(json_struct), expected_output) + json_struct = '{This is a string.}' + expected_output = 'This is a string.' + self.assertEqual(parse_maybe_json(json_struct), expected_output) + json_struct = '[This is a string.]' + expected_output = 'This is a string.' + self.assertEqual(parse_maybe_json(json_struct), expected_output) async def test_message_lings(self) -> None: request = AIMessage('Lala', 'Hello there!', 'chat', False,)