diff --git a/fjerkroa_bot/ai_responder.py b/fjerkroa_bot/ai_responder.py index ae3b4bd..5fdbcd8 100644 --- a/fjerkroa_bot/ai_responder.py +++ b/fjerkroa_bot/ai_responder.py @@ -29,6 +29,25 @@ def parse_response(content: str) -> Dict: raise err +def parse_maybe_json(json_string): + try: + # Attempt to parse the string as JSON + parsed_json = parse_response(json_string) + except Exception: + # If it fails, return the original string unchanged + return json_string + + # If the parsed JSON is a list or dictionary, concatenate its values as a string + 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) + + # If the parsed JSON is a string, return it + return str(parsed_json) + + class AIMessageBase(object): def __init__(self) -> None: pass @@ -108,8 +127,8 @@ class AIResponder(object): response['answer'] = str(response['answer']) return AIResponse(response['answer'], response['answer_needed'], - response['staff'], - response['picture'], + parse_maybe_json(response['staff']), + parse_maybe_json(response['picture']), response['hack']) def short_path(self, message: AIMessage, limit: int) -> bool: diff --git a/pyproject.toml b/pyproject.toml index 02c8575..c960e3e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,6 @@ [build-system] -requires = ["setuptools", "wheel"] +requires = ["poetry-core>=1.0.0"] +build-backend = "poetry.core.masonry.api" [tool.mypy] files = ["fjerkroa_bot", "tests"] @@ -23,3 +24,24 @@ exclude = [ "dist", "venv", ] + +[tool.poetry] +name = "fjerkroa_bot" +version = "2.0" +description = "" +authors = ["Oleksandr Kozachuk "] + +[tool.poetry.dependencies] +python = "^3.8" +discordpy = "*" +openai = "*" +aiohttp = "*" +mypy = "*" +flake8 = "*" +pre-commit = "*" +pytest = "*" +setuptools = "*" +wheel = "*" +watchdog = "*" +toml = "*" +types-toml = "*" diff --git a/tests/test_main.py b/tests/test_main.py index 1b34f17..54cb5e4 100644 --- a/tests/test_main.py +++ b/tests/test_main.py @@ -8,6 +8,7 @@ import logging import pytest from unittest.mock import Mock, PropertyMock, MagicMock, AsyncMock, patch, mock_open, ANY from fjerkroa_bot import FjerkroaBot +from fjerkroa_bot.ai_responder import parse_maybe_json from discord import User, Message, TextChannel @@ -61,6 +62,20 @@ class TestFunctionality(TestBotBase): result = FjerkroaBot.load_config('config.toml') self.assertEqual(result, self.config_data) + def test_json_strings(self) -> None: + json_string = '{"key1": "value1", "key2": "value2"}' + expected_output = "value1\nvalue2" + self.assertEqual(parse_maybe_json(json_string), expected_output) + non_json_string = "This is not a JSON string." + self.assertEqual(parse_maybe_json(non_json_string), non_json_string) + json_array = '["value1", "value2", "value3"]' + + expected_output = "value1\nvalue2\nvalue3" + self.assertEqual(parse_maybe_json(json_array), expected_output) + json_string = '"value1"' + expected_output = 'value1' + self.assertEqual(parse_maybe_json(json_string), expected_output) + async def test_on_message_event(self) -> None: async def acreate(*a, **kw): answer = {'answer': 'Hello!',