Merge pull request #35 from nervousapps/patch-1

Add VALIDATE_JSON_RETRY environment variable, and tests
clean-reqs
biobootloader 2023-04-26 17:26:43 -07:00 zatwierdzone przez GitHub
commit 756ddbdd61
Nie znaleziono w bazie danych klucza dla tego podpisu
ID klucza GPG: 4AEE18F83AFDEB23
7 zmienionych plików z 151 dodań i 56 usunięć

Wyświetl plik

@ -37,6 +37,13 @@ You can also use flag `--confirm=True` which will ask you `yes or no` before mak
python -m wolverine examples/buggy_script.py "subtract" 20 3 --confirm=True
## :label: Environement variables
| env name | description | default value |
| -------------------------------| ----------------| -------------------|
| OPENAI_API_KEY | OpenAI API key | None |
| DEFAULT_MODEL | GPT model to use | "gpt-4" |
| VALIDATE_JSON_RETRY | Number of retries when requesting OpenAI API (-1 means unlimites) | -1 |
## Future Plans
This is just a quick prototype I threw together in a few hours. There are many possible extensions and contributions are welcome:

Wyświetl plik

@ -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.

31
tests/conftest.py 100644
Wyświetl plik

@ -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

Wyświetl plik

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

Wyświetl plik

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

Wyświetl plik

@ -1,18 +1,12 @@
import os
import pytest
import tempfile
from wolverine import apply_changes
from wolverine import apply_changes, json_validated_response
@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)
from .conftest import (
mock_open_ai_response_object,
TEST_FILES_DIR
)
def test_apply_changes_replace(temp_file):
@ -53,3 +47,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

Wyświetl plik

@ -6,6 +6,9 @@ import subprocess
import sys
import openai
from typing import List, Dict
from termcolor import cprint
from dotenv import load_dotenv
from termcolor import cprint
@ -13,14 +16,18 @@ from termcolor import cprint
load_dotenv()
openai.api_key = os.getenv("OPENAI_API_KEY")
# Default model is GPT-4
DEFAULT_MODEL = os.environ.get("DEFAULT_MODEL", "gpt-4")
# Nb retries for json_validated_response, default to -1, infinite
VALIDATE_JSON_RETRY = int(os.getenv("VALIDATE_JSON_RETRY", -1))
with open("prompt.txt", encoding="utf-8") as file:
SYSTEM_PROMPT = file.read()
# Read the system prompt
with open(os.path.join(os.path.dirname(__file__), "..", "prompt.txt"), 'r') as f:
SYSTEM_PROMPT = f.read()
def run_script(script_name, script_args):
def run_script(script_name: str, script_args: List) -> str:
"""
If script_name.endswith(".py") then run with python
else run with node
@ -42,53 +49,55 @@ def run_script(script_name, script_args):
return result.decode("utf-8"), 0
def json_validated_response(model, messages):
def json_validated_response(model: str, messages: List[Dict], nb_retry: int = VALIDATE_JSON_RETRY) -> Dict:
"""
This function is needed because the API can return a non-json response.
This will run recursively until a valid json response is returned.
todo: might want to stop after a certain number of retries
This will run recursively VALIDATE_JSON_RETRY times.
If VALIDATE_JSON_RETRY is -1, it will run recursively until a valid json response is returned.
"""
response = openai.ChatCompletion.create(
model=model,
messages=messages,
temperature=0.5,
)
messages.append(response.choices[0].message)
content = response.choices[0].message.content
# see if json can be parsed
try:
json_start_index = content.index(
"["
) # find the starting position of the JSON data
json_data = content[
json_start_index:
] # extract the JSON data from the response string
json_response = json.loads(json_data)
except (json.decoder.JSONDecodeError, ValueError) as error:
cprint(f"{error}. Re-running the query.", "red")
# debug
cprint(f"\nGPT RESPONSE:\n\n{content}\n\n", "yellow")
# append a user message that says the json is invalid
messages.append(
{
"role": "user",
"content": (
"Your response could not be parsed by json.loads. "
"Please restate your last message as pure JSON."
),
}
json_response = {}
if nb_retry != 0:
response = openai.ChatCompletion.create(
model=model,
messages=messages,
temperature=0.5,
)
# rerun the api call
return json_validated_response(model, messages)
except Exception as error:
cprint(f"Unknown error: {error}", "red")
cprint(f"\nGPT RESPONSE:\n\n{content}\n\n", "yellow")
raise error
return json_response
messages.append(response.choices[0].message)
content = response.choices[0].message.content
# see if json can be parsed
try:
json_start_index = content.index(
"["
) # find the starting position of the JSON data
json_data = content[
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
cprint(f"\nGPT RESPONSE:\n\n{content}\n\n", "yellow")
# append a user message that says the json is invalid
messages.append(
{
"role": "user",
"content": "Your response could not be parsed by json.loads. Please restate your last message as pure JSON.",
}
)
# dec nb_retry
nb_retry-=1
# rerun the api call
return json_validated_response(model, messages, nb_retry)
except Exception as e:
cprint(f"Unknown error: {e}", "red")
cprint(f"\nGPT RESPONSE:\n\n{content}\n\n", "yellow")
raise e
raise Exception(f"No valid json response found after {VALIDATE_JSON_RETRY} tries. Exiting.")
def send_error_to_gpt(file_path, args, error_message, model=DEFAULT_MODEL):
with open(file_path) as f:
def send_error_to_gpt(file_path: str, args: List, error_message: str, model: str = DEFAULT_MODEL) -> Dict:
with open(file_path, "r") as f:
file_lines = f.readlines()
file_with_lines = []
@ -122,7 +131,7 @@ def send_error_to_gpt(file_path, args, error_message, model=DEFAULT_MODEL):
return json_validated_response(model, messages)
def apply_changes(file_path, changes: list, confirm=False):
def apply_changes(file_path: str, changes: List, confirm: bool = False):
"""
Pass changes as loaded json (list of dicts)
"""