diff --git a/prompt.txt b/prompt.txt index ec0582a..8b4af94 100644 --- a/prompt.txt +++ b/prompt.txt @@ -14,3 +14,5 @@ example response: {"operation": "Replace", "line": 18, "content": " x += 1"}, {"operation": "Delete", "line": 20, "content": ""} ] + +From now, your response must be only the json object, no talking, no comments. diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..1c93f2d --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,31 @@ +""" +Conftest +""" +import os +import pytest +import tempfile + + +TEST_FILES_DIR = os.path.join(os.path.dirname(__file__), "test_files") + + +@pytest.fixture(scope='function') +def temp_file(): + # Create a temporary file + with tempfile.NamedTemporaryFile(mode="w", delete=False) as f: + f.write("first line\nsecond line\nthird line") + file_path = f.name + yield file_path + # Clean up the temporary file + os.remove(file_path) + + +def mock_open_ai_response_object(mocker, content: str): + """ + Mocks the response object from the openai api. + """ + mock_generator_object = mocker.MagicMock() + mock_message_object = mocker.MagicMock() + mock_message_object.configure_mock(**{"message.content": content}) + mock_generator_object.configure_mock(**{"choices": [mock_message_object]}) + return mock_generator_object \ No newline at end of file diff --git a/tests/test_files/cc_resp.txt b/tests/test_files/cc_resp.txt new file mode 100644 index 0000000..83e3155 --- /dev/null +++ b/tests/test_files/cc_resp.txt @@ -0,0 +1,8 @@ +Explanation: The function `subtract_numbers` is never defined in the script, causing a `NameError` when it is called in the `calculate` function. + +[ + {"explanation": "The 'subtract_numbers' function is never defined in the script."}, + {"operation": "InsertAfter", "line": 12, "content": "\n# Define subtract_numbers function\ndef subtract_numbers(a, b):\n return a - b\n"}, + {"operation": "Replace", "line": 18, "content": " if operation == \"add\":\n result = add_numbers(num1, num2)\n elif operation == \"subtract\":\n result = subtract_numbers(num1, num2)\n elif operation == \"multiply\":\n result = multiply_numbers(num1, num2)\n elif operation == \"divide\":\n result = divide_numbers(num1, num2)\n else:\n print(\"Invalid operation\")\n"}, + {"operation": "Replace", "line": 30, "content": " return result\n"} +] \ No newline at end of file diff --git a/tests/test_files/cc_resp_fail.txt b/tests/test_files/cc_resp_fail.txt new file mode 100644 index 0000000..452a594 --- /dev/null +++ b/tests/test_files/cc_resp_fail.txt @@ -0,0 +1,7 @@ +Explanation: The function `subtract_numbers` is never defined in the script, causing a `NameError` when it is called in the `calculate` function. + +[ + {"explanation": "The 'subtract_numbers' function is never defined in the script."}, + {"operation": "InsertAfter", "line": 12, "content": "\n# Define subtract_numbers function\ndef subtract_numbers(a, b):\n return a - b\n"}, + {"operation": "Replace", "line": 18, "content": " if operation == \"add\":\n result = add_numbers(num1, num2)\n elif operation == \"subtract\":\n result = subtract_numbers(num1, num2)\n elif operation == \"multiply\":\n result = multiply_numbers(num1, num2)\n elif operation == \"divide\":\n result = divide_numbers(num1, num2)\n else:\n print(\"Invalid operation\")\n"}, + {"operation": "Replace", "line": 30, "content": " return result\n"} diff --git a/tests/test_wolverine.py b/tests/test_wolverine.py index 1f82d76..e2d4415 100644 --- a/tests/test_wolverine.py +++ b/tests/test_wolverine.py @@ -1,19 +1,15 @@ import os -import json import pytest -import tempfile -from wolverine import apply_changes, json_validated_response +from wolverine import ( + apply_changes, + json_validated_response, +) +from .conftest import ( + mock_open_ai_response_object, + TEST_FILES_DIR +) -@pytest.fixture(scope='function') -def temp_file(): - # Create a temporary file - with tempfile.NamedTemporaryFile(mode="w", delete=False) as f: - f.write("first line\nsecond line\nthird line") - file_path = f.name - yield file_path - # Clean up the temporary file - os.remove(file_path) def test_apply_changes_replace(temp_file): @@ -54,3 +50,40 @@ def test_apply_changes_insert(temp_file): content = f.read() assert content == 'first line\nsecond line\ninserted line\nthird line' + +@pytest.mark.parametrize("chat_completion_response, nb_retry, fail", [ + (os.path.join(TEST_FILES_DIR, "cc_resp.txt"), 3, False), + (os.path.join(TEST_FILES_DIR, "cc_resp_fail.txt"), 3, True), + (os.path.join(TEST_FILES_DIR, "cc_resp_fail.txt"), 10, True), +]) +def test_json_validated_response(mocker, chat_completion_response, nb_retry, fail): + # Open the test file + with open(chat_completion_response, 'r') as file: + response = file.read() + # Mock the openAi chat completion API call + mocker.patch( + "openai.ChatCompletion.create", + return_value=mock_open_ai_response_object(mocker=mocker, content=response)) + # ChatCompletion returned an invalid response + if fail: + with pytest.raises(Exception) as err: + json_response = json_validated_response("gpt-4", [ + { + "role": "user", + "content": "prompt" + } + ], + nb_retry=nb_retry + ) + # Check that the exception is raised after nb_retry time + assert err.value == f"No valid json response found after 3 tries. Exiting." + else: + json_response = json_validated_response("gpt-4", [ + { + "role": "user", + "content": "prompt" + } + ], + nb_retry=nb_retry + ) + assert json_response \ No newline at end of file diff --git a/wolverine/wolverine.py b/wolverine/wolverine.py index 8e21c22..30f3960 100644 --- a/wolverine/wolverine.py +++ b/wolverine/wolverine.py @@ -68,6 +68,7 @@ def json_validated_response(model: str, messages: List[Dict], nb_retry: int = VA json_start_index: ] # extract the JSON data from the response string json_response = json.loads(json_data) + return json_response except (json.decoder.JSONDecodeError, ValueError) as e: cprint(f"{e}. Re-running the query.", "red") # debug @@ -87,8 +88,7 @@ def json_validated_response(model: str, messages: List[Dict], nb_retry: int = VA cprint(f"Unknown error: {e}", "red") cprint(f"\nGPT RESPONSE:\n\n{content}\n\n", "yellow") raise e - # If not valid after VALIDATE_JSON_RETRY retries, return an empty object / or raise an exception and exit - return json_response + raise Exception(f"No valid json response found after {VALIDATE_JSON_RETRY} tries. Exiting.") def send_error_to_gpt(file_path: str, args: List, error_message: str, model: str = DEFAULT_MODEL) -> Dict: