Rebase, add test, add last line in prompt.txt to be sure to have only a json object in response.

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.

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

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"}
+]

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

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

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: