From 94ec243021883dee5215f0e04b6fb1fddf1ff812 Mon Sep 17 00:00:00 2001 From: zerolab Date: Mon, 8 Jan 2024 19:17:51 +0000 Subject: [PATCH 1/7] Add tests for the prompt view --- tests/test_views.py | 61 ++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 60 insertions(+), 1 deletion(-) diff --git a/tests/test_views.py b/tests/test_views.py index 905d796..61d6f18 100644 --- a/tests/test_views.py +++ b/tests/test_views.py @@ -1,3 +1,5 @@ +import uuid + import pytest from django.urls import reverse from wagtail_ai.views import PromptEditForm, prompt_viewset @@ -35,4 +37,61 @@ def test_prompt_model_admin_viewset_edit_view(client, setup_users, setup_prompt_ assert setup_prompt_object.label in str(response.content) -# TODO add tests for process view +@pytest.mark.django_db +def test_process_view_get_request(client, setup_users): + url = reverse("wagtail_ai:process") + + superuser = setup_users + client.force_login(superuser) + + response = client.get(url) + assert response.status_code == 400 + assert response.json() == { + "error": "No text provided - please enter some text before using AI features." + } + + +@pytest.mark.django_db +def test_process_view_post_without_text(client, setup_users): + url = reverse("wagtail_ai:process") + + superuser = setup_users + client.force_login(superuser) + + response = client.post(url, data={}) + assert response.status_code == 400 + assert response.json() == { + "error": "No text provided - please enter some text before using AI features." + } + + +@pytest.mark.django_db +@pytest.mark.parametrize("prompt", [None, "NOT-A-UUID", str(uuid.uuid4())]) +def test_process_view_with_bad_prompt_id(client, setup_users, prompt): + url = reverse("wagtail_ai:process") + + superuser = setup_users + client.force_login(superuser) + + data = {"text": "test"} + if prompt is not None: + data["prompt"] = prompt + + response = client.post(url, data=data) + assert response.status_code == 400 + assert response.json() == {"error": "Invalid prompt provided"} + + +@pytest.mark.django_db +def test_process_view_with_correct_prompt(client, setup_users, setup_prompt_object): + url = reverse("wagtail_ai:process") + + superuser = setup_users + client.force_login(superuser) + + response = client.post( + url, data={"text": "test", "prompt": str(setup_prompt_object.uuid)} + ) + assert response.status_code == 200 + # correct, the tests default is the echo backend + assert response.json() == {"message": "This is an echo backend: test"} From 45b75e8ce4a28c0ac36e1e19b065f2ffab4420e9 Mon Sep 17 00:00:00 2001 From: zerolab Date: Mon, 8 Jan 2024 19:21:02 +0000 Subject: [PATCH 2/7] Fix no text error message to avoid lots of whitespace --- src/wagtail_ai/views.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/src/wagtail_ai/views.py b/src/wagtail_ai/views.py index 4dbeef5..253d22d 100644 --- a/src/wagtail_ai/views.py +++ b/src/wagtail_ai/views.py @@ -77,13 +77,8 @@ def process(request): text = request.POST.get("text") if not text: - return JsonResponse( - { - "error": "No text provided - please enter some text before using AI \ - features" - }, - status=400, - ) + error = "No text provided - please enter some text before using AI features." + return JsonResponse({"error": error}, status=400) prompt_id = request.POST.get("prompt") From 866b859f65b29012e666f19d9f468b3095be4088 Mon Sep 17 00:00:00 2001 From: zerolab Date: Mon, 8 Jan 2024 19:21:34 +0000 Subject: [PATCH 3/7] Fix prompt view error when the passed prompt is not a valid UUID --- src/wagtail_ai/views.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/wagtail_ai/views.py b/src/wagtail_ai/views.py index 253d22d..5c44ce2 100644 --- a/src/wagtail_ai/views.py +++ b/src/wagtail_ai/views.py @@ -2,6 +2,7 @@ import os from django import forms +from django.core.exceptions import ValidationError from django.http import JsonResponse from django.views.decorators.csrf import csrf_exempt from wagtail.admin.ui.tables import UpdatedAtColumn @@ -84,7 +85,7 @@ def process(request): try: prompt = Prompt.objects.get(uuid=prompt_id) - except Prompt.DoesNotExist: + except (Prompt.DoesNotExist, ValidationError): return JsonResponse({"error": "Invalid prompt provided"}, status=400) handlers = { From 0568fb5d05ac83249ac0fd1501955ac16fc8787d Mon Sep 17 00:00:00 2001 From: zerolab Date: Thu, 11 Jan 2024 18:17:25 +0000 Subject: [PATCH 4/7] Pre-validate prompt UUID before trying to fetch it --- src/wagtail_ai/views.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/src/wagtail_ai/views.py b/src/wagtail_ai/views.py index 5c44ce2..0d45553 100644 --- a/src/wagtail_ai/views.py +++ b/src/wagtail_ai/views.py @@ -1,8 +1,8 @@ import logging import os +import uuid from django import forms -from django.core.exceptions import ValidationError from django.http import JsonResponse from django.views.decorators.csrf import csrf_exempt from wagtail.admin.ui.tables import UpdatedAtColumn @@ -73,19 +73,31 @@ def _append_handler(*, prompt: Prompt, text: str) -> str: return message +def _is_prompt_id_valid(prompt_id: str) -> bool: + try: + uuid_object = uuid.UUID(prompt_id) + except ValueError: + return False + else: + return uuid_object.version == 4 + + @csrf_exempt -def process(request): - text = request.POST.get("text") +def process(request) -> JsonResponse: + text = request.POST.get("text", "").strip() if not text: error = "No text provided - please enter some text before using AI features." return JsonResponse({"error": error}, status=400) - prompt_id = request.POST.get("prompt") + prompt_id = request.POST.get("prompt", "").strip() + + if not _is_prompt_id_valid(prompt_id): + return JsonResponse({"error": "Invalid prompt provided"}, status=400) try: prompt = Prompt.objects.get(uuid=prompt_id) - except (Prompt.DoesNotExist, ValidationError): + except Prompt.DoesNotExist: return JsonResponse({"error": "Invalid prompt provided"}, status=400) handlers = { From eb58d28be4796ea7ee58095ac1cb40597432124c Mon Sep 17 00:00:00 2001 From: zerolab Date: Thu, 11 Jan 2024 18:19:05 +0000 Subject: [PATCH 5/7] Make the prompt view error messages translatable --- src/wagtail_ai/views.py | 10 ++++++---- tests/test_views.py | 2 +- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/wagtail_ai/views.py b/src/wagtail_ai/views.py index 0d45553..634cd72 100644 --- a/src/wagtail_ai/views.py +++ b/src/wagtail_ai/views.py @@ -4,6 +4,7 @@ from django import forms from django.http import JsonResponse +from django.utils.translation import gettext as _ from django.views.decorators.csrf import csrf_exempt from wagtail.admin.ui.tables import UpdatedAtColumn from wagtail.admin.viewsets.model import ModelViewSet @@ -87,18 +88,19 @@ def process(request) -> JsonResponse: text = request.POST.get("text", "").strip() if not text: - error = "No text provided - please enter some text before using AI features." + error = _("No text provided - please enter some text before using AI features.") return JsonResponse({"error": error}, status=400) + invalid_prompt_error = _("Invalid prompt provided.") prompt_id = request.POST.get("prompt", "").strip() if not _is_prompt_id_valid(prompt_id): - return JsonResponse({"error": "Invalid prompt provided"}, status=400) + return JsonResponse({"error": invalid_prompt_error}, status=400) try: prompt = Prompt.objects.get(uuid=prompt_id) except Prompt.DoesNotExist: - return JsonResponse({"error": "Invalid prompt provided"}, status=400) + return JsonResponse({"error": invalid_prompt_error}, status=400) handlers = { Prompt.Method.REPLACE: _replace_handler, @@ -113,7 +115,7 @@ def process(request) -> JsonResponse: return JsonResponse({"error": str(e)}, status=400) except Exception: logger.exception("An unexpected error occurred.") - return JsonResponse({"error": "An unexpected error occurred"}, status=500) + return JsonResponse({"error": _("An unexpected error occurred.")}, status=500) return JsonResponse({"message": response}) diff --git a/tests/test_views.py b/tests/test_views.py index 61d6f18..9fe1657 100644 --- a/tests/test_views.py +++ b/tests/test_views.py @@ -79,7 +79,7 @@ def test_process_view_with_bad_prompt_id(client, setup_users, prompt): response = client.post(url, data=data) assert response.status_code == 400 - assert response.json() == {"error": "Invalid prompt provided"} + assert response.json() == {"error": "Invalid prompt provided."} @pytest.mark.django_db From a2409169ef653b65b359219d04cb6f74b7462501 Mon Sep 17 00:00:00 2001 From: zerolab Date: Thu, 11 Jan 2024 18:38:39 +0000 Subject: [PATCH 6/7] Silence the RemovedInDjango60Warning URLField warnings `RemovedInDjango60Warning: The FORMS_URLFIELD_ASSUME_HTTPS transitional setting is deprecated.` --- tests/testapp/settings.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/testapp/settings.py b/tests/testapp/settings.py index abe3261..ef63a56 100644 --- a/tests/testapp/settings.py +++ b/tests/testapp/settings.py @@ -219,3 +219,5 @@ }, }, } + +FORMS_URLFIELD_ASSUME_HTTPS = True From ce57192e1b0fa52b4aeecb862af130edd8b64885 Mon Sep 17 00:00:00 2001 From: zerolab Date: Fri, 12 Jan 2024 13:05:17 +0000 Subject: [PATCH 7/7] Use a form for text/prompt validation --- src/wagtail_ai/forms.py | 40 ++++++++++++++++++++++++++++++++++++++++ src/wagtail_ai/views.py | 32 +++++++++----------------------- tests/test_views.py | 10 +++++++--- 3 files changed, 56 insertions(+), 26 deletions(-) create mode 100644 src/wagtail_ai/forms.py diff --git a/src/wagtail_ai/forms.py b/src/wagtail_ai/forms.py new file mode 100644 index 0000000..ddcc34b --- /dev/null +++ b/src/wagtail_ai/forms.py @@ -0,0 +1,40 @@ +from django import forms +from django.core.exceptions import ValidationError +from django.utils.translation import gettext_lazy as _ + + +class PromptTextField(forms.CharField): + default_error_messages = { + "required": _( + "No text provided - please enter some text before using AI features." + ), + } + + +class PromptUUIDField(forms.UUIDField): + default_error_messages = { + "required": _("Invalid prompt provided."), + "invalid": _("Invalid prompt provided."), + } + + +class PromptForm(forms.Form): + text = PromptTextField() + prompt = PromptUUIDField() + + def clean_prompt(self): + prompt_uuid = self.cleaned_data["prompt"] + if prompt_uuid.version != 4: + raise ValidationError( + self.fields["prompt"].error_messages["invalid"], code="invalid" + ) + + return prompt_uuid + + def errors_for_json_response(self) -> str: + errors_for_response = [] + for _field, errors in self.errors.get_json_data().items(): + for error in errors: + errors_for_response.append(error["message"]) + + return " \n".join(errors_for_response) diff --git a/src/wagtail_ai/views.py b/src/wagtail_ai/views.py index 634cd72..ae96b68 100644 --- a/src/wagtail_ai/views.py +++ b/src/wagtail_ai/views.py @@ -1,6 +1,5 @@ import logging import os -import uuid from django import forms from django.http import JsonResponse @@ -10,6 +9,7 @@ from wagtail.admin.viewsets.model import ModelViewSet from . import ai, types +from .forms import PromptForm from .models import Prompt logger = logging.getLogger(__name__) @@ -74,33 +74,19 @@ def _append_handler(*, prompt: Prompt, text: str) -> str: return message -def _is_prompt_id_valid(prompt_id: str) -> bool: - try: - uuid_object = uuid.UUID(prompt_id) - except ValueError: - return False - else: - return uuid_object.version == 4 - - @csrf_exempt def process(request) -> JsonResponse: - text = request.POST.get("text", "").strip() + prompt_form = PromptForm(request.POST) - if not text: - error = _("No text provided - please enter some text before using AI features.") - return JsonResponse({"error": error}, status=400) - - invalid_prompt_error = _("Invalid prompt provided.") - prompt_id = request.POST.get("prompt", "").strip() - - if not _is_prompt_id_valid(prompt_id): - return JsonResponse({"error": invalid_prompt_error}, status=400) + if not prompt_form.is_valid(): + return JsonResponse( + {"error": prompt_form.errors_for_json_response()}, status=400 + ) try: - prompt = Prompt.objects.get(uuid=prompt_id) + prompt = Prompt.objects.get(uuid=prompt_form.cleaned_data["prompt"]) except Prompt.DoesNotExist: - return JsonResponse({"error": invalid_prompt_error}, status=400) + return JsonResponse({"error": _("Invalid prompt provided.")}, status=400) handlers = { Prompt.Method.REPLACE: _replace_handler, @@ -110,7 +96,7 @@ def process(request) -> JsonResponse: handler = handlers[Prompt.Method(prompt.method)] try: - response = handler(prompt=prompt, text=text) + response = handler(prompt=prompt, text=prompt_form.cleaned_data["text"]) except AIHandlerException as e: return JsonResponse({"error": str(e)}, status=400) except Exception: diff --git a/tests/test_views.py b/tests/test_views.py index 9fe1657..7df7e62 100644 --- a/tests/test_views.py +++ b/tests/test_views.py @@ -47,7 +47,8 @@ def test_process_view_get_request(client, setup_users): response = client.get(url) assert response.status_code == 400 assert response.json() == { - "error": "No text provided - please enter some text before using AI features." + "error": "No text provided - please enter some text before using AI features. " + "\nInvalid prompt provided." } @@ -61,12 +62,15 @@ def test_process_view_post_without_text(client, setup_users): response = client.post(url, data={}) assert response.status_code == 400 assert response.json() == { - "error": "No text provided - please enter some text before using AI features." + "error": "No text provided - please enter some text before using AI features. " + "\nInvalid prompt provided." } @pytest.mark.django_db -@pytest.mark.parametrize("prompt", [None, "NOT-A-UUID", str(uuid.uuid4())]) +@pytest.mark.parametrize( + "prompt", [None, "NOT-A-UUID", str(uuid.uuid1()), str(uuid.uuid4())] +) def test_process_view_with_bad_prompt_id(client, setup_users, prompt): url = reverse("wagtail_ai:process")