studio/docs
Archived
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

✨(backend) create ai endpoint

Browse Source 
We created 2 new action endpoints on the document
to perform AI operations:
- POST /api/v1.0/documents/{uuid}/ai-transform
- POST /api/v1.0/documents/{uuid}/ai-translate
This commit is contained in:
Anthony LC
2024-09-20 22:42:46 +02:00
committed by Samuel Paccoud
parent e8d95facdf
commit aff3b43c9d
17 changed files with 1440 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

33
src/backend/core/api/serializers.py
View File

@@ -9,7 +9,8 @@ from django.utils.translation import gettext_lazy as _
import magic
from rest_framework import exceptions, serializers
from core import models
from core import enums, models
from core.services.ai_services import AI_ACTIONS
class UserSerializer(serializers.ModelSerializer):
@@ -378,3 +379,33 @@ class VersionFilterSerializer(serializers.Serializer):
page_size = serializers.IntegerField(
required=False, min_value=1, max_value=50, default=20
)
class AITransformSerializer(serializers.Serializer):
"""Serializer for AI transform requests."""
action = serializers.ChoiceField(choices=AI_ACTIONS, required=True)
text = serializers.CharField(required=True)
def validate_text(self, value):
"""Ensure the text field is not empty."""
if len(value.strip()) == 0:
raise serializers.ValidationError("Text field cannot be empty.")
return value
class AITranslateSerializer(serializers.Serializer):
"""Serializer for AI translate requests."""
language = serializers.ChoiceField(
choices=tuple(enums.ALL_LANGUAGES.items()), required=True
)
text = serializers.CharField(required=True)
def validate_text(self, value):
"""Ensure the text field is not empty."""
if len(value.strip()) == 0:
raise serializers.ValidationError("Text field cannot be empty.")
return value

96
src/backend/core/api/utils.py
View File

@@ -1,8 +1,14 @@
"""Util to generate S3 authorization headers for object storage access control"""
import time
from abc import ABC, abstractmethod
from django.conf import settings
from django.core.cache import cache
from django.core.files.storage import default_storage
import botocore
from rest_framework.throttling import BaseThrottle
def generate_s3_authorization_headers(key):
@@ -31,3 +37,93 @@ def generate_s3_authorization_headers(key):
auth.add_auth(request)
return request
class AIBaseRateThrottle(BaseThrottle, ABC):
"""Base throttle class for AI-related rate limiting with backoff."""
def __init__(self, rates):
"""Initialize instance attributes with configurable rates."""
super().__init__()
self.rates = rates
self.cache_key = None
self.recent_requests_minute = 0
self.recent_requests_hour = 0
self.recent_requests_day = 0
@abstractmethod
def get_cache_key(self, request, view):
"""Abstract method to generate cache key for throttling."""
def allow_request(self, request, view):
"""Check if the request is allowed based on rate limits."""
self.cache_key = self.get_cache_key(request, view)
if not self.cache_key:
return True # Allow if no cache key is generated
now = time.time()
history = cache.get(self.cache_key, [])
# Keep requests within the last 24 hours
history = [req for req in history if req > now - 86400]
# Calculate recent requests
self.recent_requests_minute = len([req for req in history if req > now - 60])
self.recent_requests_hour = len([req for req in history if req > now - 3600])
self.recent_requests_day = len(history)
# Check rate limits
if self.recent_requests_minute >= self.rates["minute"]:
return False
if self.recent_requests_hour >= self.rates["hour"]:
return False
if self.recent_requests_day >= self.rates["day"]:
return False
# Log the request
history.append(now)
cache.set(self.cache_key, history, timeout=86400)
return True
def wait(self):
"""Implement a backoff strategy by increasing wait time based on limits hit."""
if self.recent_requests_day >= self.rates["day"]:
return 86400
if self.recent_requests_hour >= self.rates["hour"]:
return 3600
if self.recent_requests_minute >= self.rates["minute"]:
return 60
return None
class AIDocumentRateThrottle(AIBaseRateThrottle):
"""Throttle for limiting AI requests per document with backoff."""
def __init__(self, *args, **kwargs):
super().__init__(settings.AI_DOCUMENT_RATE_THROTTLE_RATES)
def get_cache_key(self, request, view):
"""Include document ID in the cache key."""
document_id = view.kwargs["pk"]
return f"document_{document_id}_throttle_ai"
class AIUserRateThrottle(AIBaseRateThrottle):
"""Throttle that limits requests per user or IP with backoff and rate limits."""
def __init__(self, *args, **kwargs):
super().__init__(settings.AI_USER_RATE_THROTTLE_RATES)
def get_cache_key(self, request, view=None):
"""Generate a cache key based on the user ID or IP for anonymous users."""
if request.user.is_authenticated:
return f"user_{request.user.id!s}_throttle_ai"
return f"anonymous_{self.get_ident(request)}_throttle_ai"
def get_ident(self, request):
"""Return the request IP address."""
x_forwarded_for = request.META.get("HTTP_X_FORWARDED_FOR")
return (
x_forwarded_for.split(",")[0]
if x_forwarded_for
else request.META.get("REMOTE_ADDR")
)

95
src/backend/core/api/viewsets.py
View File

@@ -21,6 +21,7 @@ from rest_framework import (
decorators,
exceptions,
filters,
metadata,
mixins,
pagination,
status,
@@ -30,7 +31,8 @@ from rest_framework import (
response as drf_response,
)
from core import models
from core import enums, models
from core.services.ai_services import AIService
from . import permissions, serializers, utils
@@ -302,6 +304,23 @@ class ResourceAccessViewsetMixin:
serializer.save()
class DocumentMetadata(metadata.SimpleMetadata):
"""Custom metadata class to add information"""
def determine_metadata(self, request, view):
"""Add language choices only for the list endpoint."""
simple_metadata = super().determine_metadata(request, view)
if request.path.endswith("/documents/"):
simple_metadata["actions"]["POST"]["language"] = {
"choices": [
{"value": code, "display_name": name}
for code, name in enums.ALL_LANGUAGES.items()
]
}
return simple_metadata
class DocumentViewSet(
ResourceViewsetMixin,
mixins.CreateModelMixin,
@@ -319,6 +338,7 @@ class DocumentViewSet(
resource_field_name = "document"
queryset = models.Document.objects.all()
ordering = ["-updated_at"]
metadata_class = DocumentMetadata
def list(self, request, *args, **kwargs):
"""Restrict resources returned by the list endpoint"""
@@ -455,10 +475,7 @@ class DocumentViewSet(
serializer = serializers.LinkDocumentSerializer(
document, data=request.data, partial=True
)
if not serializer.is_valid():
return drf_response.Response(
serializer.errors, status=status.HTTP_400_BAD_REQUEST
)
serializer.is_valid(raise_exception=True)
serializer.save()
return drf_response.Response(serializer.data, status=status.HTTP_200_OK)
@@ -471,10 +488,7 @@ class DocumentViewSet(
# Validate metadata in payload
serializer = serializers.FileUploadSerializer(data=request.data)
if not serializer.is_valid():
return drf_response.Response(
serializer.errors, status=status.HTTP_400_BAD_REQUEST
)
serializer.is_valid(raise_exception=True)
# Generate a generic yet unique filename to store the image in object storage
file_id = uuid.uuid4()
@@ -482,13 +496,13 @@ class DocumentViewSet(
key = f"{document.key_base}/{ATTACHMENTS_FOLDER:s}/{file_id!s}.{extension:s}"
# Prepare metadata for storage
metadata = {"Metadata": {"owner": str(request.user.id)}}
extra_args = {"Metadata": {"owner": str(request.user.id)}}
if serializer.validated_data["is_unsafe"]:
metadata["Metadata"]["is_unsafe"] = "true"
extra_args["Metadata"]["is_unsafe"] = "true"
file = serializer.validated_data["file"]
default_storage.connection.meta.client.upload_fileobj(
file, default_storage.bucket_name, key, ExtraArgs=metadata
file, default_storage.bucket_name, key, ExtraArgs=extra_args
)
return drf_response.Response(
@@ -537,6 +551,63 @@ class DocumentViewSet(
request = utils.generate_s3_authorization_headers(f"{pk:s}/{attachment_key:s}")
return drf_response.Response("authorized", headers=request.headers, status=200)
@decorators.action(
detail=True,
methods=["post"],
name="Apply a transformation action on a piece of text with AI",
url_path="ai-transform",
throttle_classes=[utils.AIDocumentRateThrottle, utils.AIUserRateThrottle],
)
def ai_transform(self, request, *args, **kwargs):
"""
POST /api/v1.0/documents/<resource_id>/ai-transform
with expected data:
- text: str
- action: str [prompt, correct, rephrase, summarize]
Return JSON response with the processed text.
"""
# Check permissions first
self.get_object()
serializer = serializers.AITransformSerializer(data=request.data)
serializer.is_valid(raise_exception=True)
text = serializer.validated_data["text"]
action = serializer.validated_data["action"]
response = AIService().transform(text, action)
return drf_response.Response(response, status=status.HTTP_200_OK)
@decorators.action(
detail=True,
methods=["post"],
name="Translate a piece of text with AI",
serializer_class=serializers.AITranslateSerializer,
url_path="ai-translate",
throttle_classes=[utils.AIDocumentRateThrottle, utils.AIUserRateThrottle],
)
def ai_translate(self, request, *args, **kwargs):
"""
POST /api/v1.0/documents/<resource_id>/ai-translate
with expected data:
- text: str
- language: str [settings.LANGUAGES]
Return JSON response with the translated text.
"""
# Check permissions first
self.get_object()
serializer = self.get_serializer(data=request.data)
serializer.is_valid(raise_exception=True)
text = serializer.validated_data["text"]
language = serializer.validated_data["language"]
response = AIService().translate(text, language)
return drf_response.Response(response, status=status.HTTP_200_OK)
class DocumentAccessViewSet(
ResourceAccessViewsetMixin,

14
src/backend/core/enums.py
View File

@@ -2,15 +2,11 @@
Core application enums declaration
"""
from django.conf import global_settings, settings
from django.conf import global_settings
from django.utils.translation import gettext_lazy as _
# Django sets `LANGUAGES` by default with all supported languages. We can use it for
# the choice of languages which should not be limited to the few languages active in
# the app.
# In Django's code base, `LANGUAGES` is set by default with all supported languages.
# We can use it for the choice of languages which should not be limited to the few languages
# active in the app.
# pylint: disable=no-member
ALL_LANGUAGES = getattr(
settings,
"ALL_LANGUAGES",
[(language, _(name)) for language, name in global_settings.LANGUAGES],
)
ALL_LANGUAGES = {language: _(name) for language, name in global_settings.LANGUAGES}

2
src/backend/core/models.py
View File

@@ -508,6 +508,8 @@ class Document(BaseModel):
can_get = bool(roles)
return {
"ai_transform": is_owner_or_admin or is_editor,
"ai_translate": is_owner_or_admin or is_editor,
"attachment_upload": is_owner_or_admin or is_editor,
"destroy": RoleChoices.OWNER in roles,
"link_configuration": is_owner_or_admin,

0
src/backend/core/services/__init__.py Normal file
View File

89
src/backend/core/services/ai_services.py Normal file
View File

@@ -0,0 +1,89 @@
"""AI services."""
import json
import re
from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
from openai import OpenAI
from core import enums
AI_ACTIONS = {
"prompt": (
"Answer the prompt in markdown format. Return JSON: "
'{"answer": "Your markdown answer"}. '
"Do not provide any other information."
),
"correct": (
"Correct grammar and spelling of the markdown text, "
"preserving language and markdown formatting. "
'Return JSON: {"answer": "your corrected markdown text"}. '
"Do not provide any other information."
),
"rephrase": (
"Rephrase the given markdown text, "
"preserving language and markdown formatting. "
'Return JSON: {"answer": "your rephrased markdown text"}. '
"Do not provide any other information."
),
"summarize": (
"Summarize the markdown text, preserving language and markdown formatting. "
'Return JSON: {"answer": "your markdown summary"}. '
"Do not provide any other information."
),
}
AI_TRANSLATE = (
"Translate the markdown text to {language:s}, preserving markdown formatting. "
'Return JSON: {{"answer": "your translated markdown text in {language:s}"}}. '
"Do not provide any other information."
)
class AIService:
"""Service class for AI-related operations."""
def __init__(self):
"""Ensure that the AI configuration is set properly."""
if (
settings.AI_BASE_URL is None
or settings.AI_API_KEY is None
or settings.AI_MODEL is None
):
raise ImproperlyConfigured("AI configuration not set")
self.client = OpenAI(base_url=settings.AI_BASE_URL, api_key=settings.AI_API_KEY)
def call_ai_api(self, system_content, text):
"""Helper method to call the OpenAI API and process the response."""
response = self.client.chat.completions.create(
model=settings.AI_MODEL,
response_format={"type": "json_object"},
messages=[
{"role": "system", "content": system_content},
{"role": "user", "content": json.dumps({"markdown_input": text})},
],
)
content = response.choices[0].message.content
sanitized_content = re.sub(r"(?<!\\)\n", "\\\\n", content)
sanitized_content = re.sub(r"(?<!\\)\t", "\\\\t", sanitized_content)
json_response = json.loads(sanitized_content)
if "answer" not in json_response:
raise RuntimeError("AI response does not contain an answer")
return json_response
def transform(self, text, action):
"""Transform text based on specified action."""
system_content = AI_ACTIONS[action]
return self.call_ai_api(system_content, text)
def translate(self, text, language):
"""Translate text to a specified language."""
language_display = enums.ALL_LANGUAGES.get(language, language)
system_content = AI_TRANSLATE.format(language=language_display)
return self.call_ai_api(system_content, text)

346
src/backend/core/tests/documents/test_api_documents_ai_transform.py Normal file
View File

@@ -0,0 +1,346 @@
"""
Test AI transform API endpoint for users in impress's core app.
"""
from unittest.mock import MagicMock, patch
from django.core.cache import cache
from django.test import override_settings
import pytest
from rest_framework.test import APIClient
from core import factories
from core.tests.conftest import TEAM, USER, VIA
pytestmark = pytest.mark.django_db
@pytest.fixture(autouse=True)
def clear_cache():
"""Fixture to clear the cache before each test."""
cache.clear()
@pytest.fixture
def ai_settings():
"""Fixture to set AI settings."""
with override_settings(
AI_BASE_URL="http://example.com", AI_API_KEY="test-key", AI_MODEL="llama"
):
yield
@pytest.mark.parametrize(
"reach, role",
[
("restricted", "reader"),
("restricted", "editor"),
("authenticated", "reader"),
("authenticated", "editor"),
("public", "reader"),
],
)
def test_api_documents_ai_transform_anonymous_forbidden(reach, role):
"""
Anonymous users should not be able to request AI transform if the link reach
and role don't allow it.
"""
document = factories.DocumentFactory(link_reach=reach, link_role=role)
url = f"/api/v1.0/documents/{document.id!s}/ai-transform/"
response = APIClient().post(url, {"text": "hello", "action": "prompt"})
assert response.status_code == 401
assert response.json() == {
"detail": "Authentication credentials were not provided."
}
@pytest.mark.usefixtures("ai_settings")
@patch("openai.resources.chat.completions.Completions.create")
def test_api_documents_ai_transform_anonymous_success(mock_create):
"""
Anonymous users should be able to request AI transform to a document
if the link reach and role permit it.
"""
document = factories.DocumentFactory(link_reach="public", link_role="editor")
answer = '{"answer": "Salut"}'
mock_create.return_value = MagicMock(
choices=[MagicMock(message=MagicMock(content=answer))]
)
url = f"/api/v1.0/documents/{document.id!s}/ai-transform/"
response = APIClient().post(url, {"text": "Hello", "action": "summarize"})
assert response.status_code == 200
assert response.json() == {"answer": "Salut"}
mock_create.assert_called_once_with(
model="llama",
response_format={"type": "json_object"},
messages=[
{
"role": "system",
"content": (
"Summarize the markdown text, preserving language and markdown formatting. "
'Return JSON: {"answer": "your markdown summary"}. Do not provide any other '
"information."
),
},
{"role": "user", "content": '{"markdown_input": "Hello"}'},
],
)
@pytest.mark.parametrize(
"reach, role",
[
("restricted", "reader"),
("restricted", "editor"),
("authenticated", "reader"),
("public", "reader"),
],
)
def test_api_documents_ai_transform_authenticated_forbidden(reach, role):
"""
Users who are not related to a document can't request AI transform if the
link reach and role don't allow it.
"""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
document = factories.DocumentFactory(link_reach=reach, link_role=role)
url = f"/api/v1.0/documents/{document.id!s}/ai-transform/"
response = client.post(url, {"text": "Hello", "action": "prompt"})
assert response.status_code == 403
assert response.json() == {
"detail": "You do not have permission to perform this action."
}
@pytest.mark.parametrize(
"reach, role",
[
("authenticated", "editor"),
("public", "editor"),
],
)
@pytest.mark.usefixtures("ai_settings")
@patch("openai.resources.chat.completions.Completions.create")
def test_api_documents_ai_transform_authenticated_success(mock_create, reach, role):
"""
Autenticated who are not related to a document should be able to request AI transform
if the link reach and role permit it.
"""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
document = factories.DocumentFactory(link_reach=reach, link_role=role)
answer = '{"answer": "Salut"}'
mock_create.return_value = MagicMock(
choices=[MagicMock(message=MagicMock(content=answer))]
)
url = f"/api/v1.0/documents/{document.id!s}/ai-transform/"
response = client.post(url, {"text": "Hello", "action": "prompt"})
assert response.status_code == 200
assert response.json() == {"answer": "Salut"}
mock_create.assert_called_once_with(
model="llama",
response_format={"type": "json_object"},
messages=[
{
"role": "system",
"content": (
'Answer the prompt in markdown format. Return JSON: {"answer": '
'"Your markdown answer"}. Do not provide any other information.'
),
},
{"role": "user", "content": '{"markdown_input": "Hello"}'},
],
)
@pytest.mark.parametrize("via", VIA)
def test_api_documents_ai_transform_reader(via, mock_user_teams):
"""
Users who are simple readers on a document should not be allowed to request AI transform.
"""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
document = factories.DocumentFactory(link_role="reader")
if via == USER:
factories.UserDocumentAccessFactory(document=document, user=user, role="reader")
elif via == TEAM:
mock_user_teams.return_value = ["lasuite", "unknown"]
factories.TeamDocumentAccessFactory(
document=document, team="lasuite", role="reader"
)
url = f"/api/v1.0/documents/{document.id!s}/ai-transform/"
response = client.post(url, {"text": "Hello", "action": "prompt"})
assert response.status_code == 403
assert response.json() == {
"detail": "You do not have permission to perform this action."
}
@pytest.mark.parametrize("role", ["editor", "administrator", "owner"])
@pytest.mark.parametrize("via", VIA)
@pytest.mark.usefixtures("ai_settings")
@patch("openai.resources.chat.completions.Completions.create")
def test_api_documents_ai_transform_success(mock_create, via, role, mock_user_teams):
"""
Editors, administrators and owners of a document should be able to request AI transform.
"""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
document = factories.DocumentFactory()
if via == USER:
factories.UserDocumentAccessFactory(document=document, user=user, role=role)
elif via == TEAM:
mock_user_teams.return_value = ["lasuite", "unknown"]
factories.TeamDocumentAccessFactory(
document=document, team="lasuite", role=role
)
answer = '{"answer": "Salut"}'
mock_create.return_value = MagicMock(
choices=[MagicMock(message=MagicMock(content=answer))]
)
url = f"/api/v1.0/documents/{document.id!s}/ai-transform/"
response = client.post(url, {"text": "Hello", "action": "prompt"})
assert response.status_code == 200
assert response.json() == {"answer": "Salut"}
mock_create.assert_called_once_with(
model="llama",
response_format={"type": "json_object"},
messages=[
{
"role": "system",
"content": (
'Answer the prompt in markdown format. Return JSON: {"answer": '
'"Your markdown answer"}. Do not provide any other information.'
),
},
{"role": "user", "content": '{"markdown_input": "Hello"}'},
],
)
def test_api_documents_ai_transform_empty_text():
"""The text should not be empty when requesting AI transform."""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
document = factories.DocumentFactory(link_reach="public", link_role="editor")
url = f"/api/v1.0/documents/{document.id!s}/ai-transform/"
response = client.post(url, {"text": " ", "action": "prompt"})
assert response.status_code == 400
assert response.json() == {"text": ["This field may not be blank."]}
def test_api_documents_ai_transform_invalid_action():
"""The action should valid when requesting AI transform."""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
document = factories.DocumentFactory(link_reach="public", link_role="editor")
url = f"/api/v1.0/documents/{document.id!s}/ai-transform/"
response = client.post(url, {"text": "Hello", "action": "invalid"})
assert response.status_code == 400
assert response.json() == {"action": ['"invalid" is not a valid choice.']}
@override_settings(AI_DOCUMENT_RATE_THROTTLE_RATES={"minute": 3, "hour": 6, "day": 10})
@pytest.mark.usefixtures("ai_settings")
@patch("openai.resources.chat.completions.Completions.create")
def test_api_documents_ai_transform_throttling_document(mock_create):
"""
Throttling per document should be triggered on the AI transform endpoint.
For full throttle class test see: `test_api_utils_ai_document_rate_throttles`
"""
client = APIClient()
document = factories.DocumentFactory(link_reach="public", link_role="editor")
answer = '{"answer": "Salut"}'
mock_create.return_value = MagicMock(
choices=[MagicMock(message=MagicMock(content=answer))]
)
url = f"/api/v1.0/documents/{document.id!s}/ai-transform/"
for _ in range(3):
user = factories.UserFactory()
client.force_login(user)
response = client.post(url, {"text": "Hello", "action": "summarize"})
assert response.status_code == 200
assert response.json() == {"answer": "Salut"}
user = factories.UserFactory()
client.force_login(user)
response = client.post(url, {"text": "Hello", "action": "summarize"})
assert response.status_code == 429
assert response.json() == {
"detail": "Request was throttled. Expected available in 60 seconds."
}
@override_settings(AI_USER_RATE_THROTTLE_RATES={"minute": 3, "hour": 6, "day": 10})
@pytest.mark.usefixtures("ai_settings")
@patch("openai.resources.chat.completions.Completions.create")
def test_api_documents_ai_transform_throttling_user(mock_create):
"""
Throttling per user should be triggered on the AI transform endpoint.
For full throttle class test see: `test_api_utils_ai_user_rate_throttles`
"""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
answer = '{"answer": "Salut"}'
mock_create.return_value = MagicMock(
choices=[MagicMock(message=MagicMock(content=answer))]
)
for _ in range(3):
document = factories.DocumentFactory(link_reach="public", link_role="editor")
url = f"/api/v1.0/documents/{document.id!s}/ai-transform/"
response = client.post(url, {"text": "Hello", "action": "summarize"})
assert response.status_code == 200
assert response.json() == {"answer": "Salut"}
document = factories.DocumentFactory(link_reach="public", link_role="editor")
url = f"/api/v1.0/documents/{document.id!s}/ai-transform/"
response = client.post(url, {"text": "Hello", "action": "summarize"})
assert response.status_code == 429
assert response.json() == {
"detail": "Request was throttled. Expected available in 60 seconds."
}

370
src/backend/core/tests/documents/test_api_documents_ai_translate.py Normal file
View File

@@ -0,0 +1,370 @@
"""
Test AI translate API endpoint for users in impress's core app.
"""
from unittest.mock import MagicMock, patch
from django.core.cache import cache
from django.test import override_settings
import pytest
from rest_framework.test import APIClient
from core import factories
from core.tests.conftest import TEAM, USER, VIA
pytestmark = pytest.mark.django_db
@pytest.fixture(autouse=True)
def clear_cache():
"""Fixture to clear the cache before each test."""
cache.clear()
@pytest.fixture
def ai_settings():
"""Fixture to set AI settings."""
with override_settings(
AI_BASE_URL="http://example.com", AI_API_KEY="test-key", AI_MODEL="llama"
):
yield
def test_api_documents_ai_translate_viewset_options_metadata():
"""The documents endpoint should give us the list of available languages."""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
factories.DocumentFactory(link_reach="public", link_role="editor")
response = APIClient().options("/api/v1.0/documents/")
assert response.status_code == 200
metadata = response.json()
assert metadata["name"] == "Document List"
assert metadata["actions"]["POST"]["language"]["choices"][0] == {
"value": "af",
"display_name": "Afrikaans",
}
@pytest.mark.parametrize(
"reach, role",
[
("restricted", "reader"),
("restricted", "editor"),
("authenticated", "reader"),
("authenticated", "editor"),
("public", "reader"),
],
)
def test_api_documents_ai_translate_anonymous_forbidden(reach, role):
"""
Anonymous users should not be able to request AI translate if the link reach
and role don't allow it.
"""
document = factories.DocumentFactory(link_reach=reach, link_role=role)
url = f"/api/v1.0/documents/{document.id!s}/ai-translate/"
response = APIClient().post(url, {"text": "hello", "language": "es"})
assert response.status_code == 401
assert response.json() == {
"detail": "Authentication credentials were not provided."
}
@pytest.mark.usefixtures("ai_settings")
@patch("openai.resources.chat.completions.Completions.create")
def test_api_documents_ai_translate_anonymous_success(mock_create):
"""
Anonymous users should be able to request AI translate to a document
if the link reach and role permit it.
"""
document = factories.DocumentFactory(link_reach="public", link_role="editor")
answer = '{"answer": "Salut"}'
mock_create.return_value = MagicMock(
choices=[MagicMock(message=MagicMock(content=answer))]
)
url = f"/api/v1.0/documents/{document.id!s}/ai-translate/"
response = APIClient().post(url, {"text": "Hello", "language": "es"})
assert response.status_code == 200
assert response.json() == {"answer": "Salut"}
mock_create.assert_called_once_with(
model="llama",
response_format={"type": "json_object"},
messages=[
{
"role": "system",
"content": (
"Translate the markdown text to Spanish, preserving markdown formatting. "
'Return JSON: {"answer": "your translated markdown text in Spanish"}. '
"Do not provide any other information."
),
},
{"role": "user", "content": '{"markdown_input": "Hello"}'},
],
)
@pytest.mark.parametrize(
"reach, role",
[
("restricted", "reader"),
("restricted", "editor"),
("authenticated", "reader"),
("public", "reader"),
],
)
def test_api_documents_ai_translate_authenticated_forbidden(reach, role):
"""
Users who are not related to a document can't request AI translate if the
link reach and role don't allow it.
"""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
document = factories.DocumentFactory(link_reach=reach, link_role=role)
url = f"/api/v1.0/documents/{document.id!s}/ai-translate/"
response = client.post(url, {"text": "Hello", "language": "es"})
assert response.status_code == 403
assert response.json() == {
"detail": "You do not have permission to perform this action."
}
@pytest.mark.parametrize(
"reach, role",
[
("authenticated", "editor"),
("public", "editor"),
],
)
@pytest.mark.usefixtures("ai_settings")
@patch("openai.resources.chat.completions.Completions.create")
def test_api_documents_ai_translate_authenticated_success(mock_create, reach, role):
"""
Autenticated who are not related to a document should be able to request AI translate
if the link reach and role permit it.
"""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
document = factories.DocumentFactory(link_reach=reach, link_role=role)
answer = '{"answer": "Salut"}'
mock_create.return_value = MagicMock(
choices=[MagicMock(message=MagicMock(content=answer))]
)
url = f"/api/v1.0/documents/{document.id!s}/ai-translate/"
response = client.post(url, {"text": "Hello", "language": "es-co"})
assert response.status_code == 200
assert response.json() == {"answer": "Salut"}
mock_create.assert_called_once_with(
model="llama",
response_format={"type": "json_object"},
messages=[
{
"role": "system",
"content": (
"Translate the markdown text to Colombian Spanish, "
"preserving markdown formatting. Return JSON: "
'{"answer": "your translated markdown text in Colombian Spanish"}. '
"Do not provide any other information."
),
},
{"role": "user", "content": '{"markdown_input": "Hello"}'},
],
)
@pytest.mark.parametrize("via", VIA)
def test_api_documents_ai_translate_reader(via, mock_user_teams):
"""
Users who are simple readers on a document should not be allowed to request AI translate.
"""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
document = factories.DocumentFactory(link_role="reader")
if via == USER:
factories.UserDocumentAccessFactory(document=document, user=user, role="reader")
elif via == TEAM:
mock_user_teams.return_value = ["lasuite", "unknown"]
factories.TeamDocumentAccessFactory(
document=document, team="lasuite", role="reader"
)
url = f"/api/v1.0/documents/{document.id!s}/ai-translate/"
response = client.post(url, {"text": "Hello", "language": "es"})
assert response.status_code == 403
assert response.json() == {
"detail": "You do not have permission to perform this action."
}
@pytest.mark.parametrize("role", ["editor", "administrator", "owner"])
@pytest.mark.parametrize("via", VIA)
@pytest.mark.usefixtures("ai_settings")
@patch("openai.resources.chat.completions.Completions.create")
def test_api_documents_ai_translate_success(mock_create, via, role, mock_user_teams):
"""
Editors, administrators and owners of a document should be able to request AI translate.
"""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
document = factories.DocumentFactory()
if via == USER:
factories.UserDocumentAccessFactory(document=document, user=user, role=role)
elif via == TEAM:
mock_user_teams.return_value = ["lasuite", "unknown"]
factories.TeamDocumentAccessFactory(
document=document, team="lasuite", role=role
)
answer = '{"answer": "Salut"}'
mock_create.return_value = MagicMock(
choices=[MagicMock(message=MagicMock(content=answer))]
)
url = f"/api/v1.0/documents/{document.id!s}/ai-translate/"
response = client.post(url, {"text": "Hello", "language": "es-co"})
assert response.status_code == 200
assert response.json() == {"answer": "Salut"}
mock_create.assert_called_once_with(
model="llama",
response_format={"type": "json_object"},
messages=[
{
"role": "system",
"content": (
"Translate the markdown text to Colombian Spanish, "
"preserving markdown formatting. Return JSON: "
'{"answer": "your translated markdown text in Colombian Spanish"}. '
"Do not provide any other information."
),
},
{"role": "user", "content": '{"markdown_input": "Hello"}'},
],
)
def test_api_documents_ai_translate_empty_text():
"""The text should not be empty when requesting AI translate."""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
document = factories.DocumentFactory(link_reach="public", link_role="editor")
url = f"/api/v1.0/documents/{document.id!s}/ai-translate/"
response = client.post(url, {"text": " ", "language": "es"})
assert response.status_code == 400
assert response.json() == {"text": ["This field may not be blank."]}
def test_api_documents_ai_translate_invalid_action():
"""The action should valid when requesting AI translate."""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
document = factories.DocumentFactory(link_reach="public", link_role="editor")
url = f"/api/v1.0/documents/{document.id!s}/ai-translate/"
response = client.post(url, {"text": "Hello", "language": "invalid"})
assert response.status_code == 400
assert response.json() == {"language": ['"invalid" is not a valid choice.']}
@override_settings(AI_DOCUMENT_RATE_THROTTLE_RATES={"minute": 3, "hour": 6, "day": 10})
@pytest.mark.usefixtures("ai_settings")
@patch("openai.resources.chat.completions.Completions.create")
def test_api_documents_ai_translate_throttling_document(mock_create):
"""
Throttling per document should be triggered on the AI translate endpoint.
For full throttle class test see: `test_api_utils_ai_document_rate_throttles`
"""
client = APIClient()
document = factories.DocumentFactory(link_reach="public", link_role="editor")
answer = '{"answer": "Salut"}'
mock_create.return_value = MagicMock(
choices=[MagicMock(message=MagicMock(content=answer))]
)
url = f"/api/v1.0/documents/{document.id!s}/ai-translate/"
for _ in range(3):
user = factories.UserFactory()
client.force_login(user)
response = client.post(url, {"text": "Hello", "language": "es"})
assert response.status_code == 200
assert response.json() == {"answer": "Salut"}
user = factories.UserFactory()
client.force_login(user)
response = client.post(url, {"text": "Hello", "language": "es"})
assert response.status_code == 429
assert response.json() == {
"detail": "Request was throttled. Expected available in 60 seconds."
}
@override_settings(AI_USER_RATE_THROTTLE_RATES={"minute": 3, "hour": 6, "day": 10})
@pytest.mark.usefixtures("ai_settings")
@patch("openai.resources.chat.completions.Completions.create")
def test_api_documents_ai_translate_throttling_user(mock_create):
"""
Throttling per user should be triggered on the AI translate endpoint.
For full throttle class test see: `test_api_utils_ai_user_rate_throttles`
"""
user = factories.UserFactory()
client = APIClient()
client.force_login(user)
answer = '{"answer": "Salut"}'
mock_create.return_value = MagicMock(
choices=[MagicMock(message=MagicMock(content=answer))]
)
for _ in range(3):
document = factories.DocumentFactory(link_reach="public", link_role="editor")
url = f"/api/v1.0/documents/{document.id!s}/ai-translate/"
response = client.post(url, {"text": "Hello", "language": "es"})
assert response.status_code == 200
assert response.json() == {"answer": "Salut"}
document = factories.DocumentFactory(link_reach="public", link_role="editor")
url = f"/api/v1.0/documents/{document.id!s}/ai-translate/"
response = client.post(url, {"text": "Hello", "language": "es"})
assert response.status_code == 429
assert response.json() == {
"detail": "Request was throttled. Expected available in 60 seconds."
}

4
src/backend/core/tests/documents/test_api_documents_retrieve.py
View File

@@ -21,6 +21,8 @@ def test_api_documents_retrieve_anonymous_public():
assert response.json() == {
"id": str(document.id),
"abilities": {
"ai_transform": document.link_role == "editor",
"ai_translate": document.link_role == "editor",
"attachment_upload": document.link_role == "editor",
"destroy": False,
"link_configuration": False,
@@ -75,6 +77,8 @@ def test_api_documents_retrieve_authenticated_unrelated_public_or_authenticated(
assert response.json() == {
"id": str(document.id),
"abilities": {
"ai_transform": document.link_role == "editor",
"ai_translate": document.link_role == "editor",
"attachment_upload": document.link_role == "editor",
"link_configuration": False,
"destroy": False,

127
src/backend/core/tests/test_api_utils_ai_document_rate_throttles.py Normal file
View File

@@ -0,0 +1,127 @@
"""
Test throttling on documents for the AI endpoint.
"""
from unittest.mock import patch
from django.core.cache import cache
from django.test import override_settings
import pytest
from rest_framework.response import Response
from rest_framework.test import APIRequestFactory
from rest_framework.views import APIView
from core.api.utils import AIDocumentRateThrottle
class DocumentAPIView(APIView):
"""A simple view to test the throttle"""
throttle_classes = [AIDocumentRateThrottle]
def get(self, request, *args, **kwargs):
"""Minimal get method for testing purposes."""
return Response({"message": "Success"})
@pytest.fixture(autouse=True)
def clear_cache():
"""Fixture to clear the cache before each test."""
cache.clear()
@override_settings(AI_DOCUMENT_RATE_THROTTLE_RATES={"minute": 3, "hour": 6, "day": 10})
@patch("time.time")
def test_api_utils_ai_document_rate_throttle_minute_limit(mock_time):
"""Test that minute limit is enforced."""
api_rf = APIRequestFactory()
mock_time.return_value = 1000000
# Simulate requests to the document API
for _i in range(3): # 3 first requests should be allowed
request = api_rf.get("/documents/1/")
response = DocumentAPIView.as_view()(request, pk=1)
assert response.status_code == 200
# Simulate passage of time
mock_time.return_value += 59
# 4th request should be throttled
request = api_rf.get("/documents/1/")
response = DocumentAPIView.as_view()(request, pk=1)
assert response.status_code == 429
# After the 60s backoff wait time has passed, we can make a request again
mock_time.return_value += 1
request = api_rf.get("/documents/1/")
response = DocumentAPIView.as_view()(request, pk=1)
assert response.status_code == 200
@override_settings(
AI_DOCUMENT_RATE_THROTTLE_RATES={"minute": 100000, "hour": 6, "day": 10}
)
@patch("time.time")
def test_ai_document_rate_throttle_hour_limit(mock_time):
"""Test that the hour limit is enforced without hitting the minute limit."""
api_rf = APIRequestFactory()
mock_time.return_value = 1000000
# Make requests to the document API, one per 21 seconds to avoid hitting the minute limit
for _i in range(6):
request = api_rf.get("/documents/1/")
response = DocumentAPIView.as_view()(request, pk=1)
assert response.status_code == 200
# Simulate passage of time
mock_time.return_value += 21
# Simulate passage of time
mock_time.return_value += 3600 - 6 * 21 - 1
# 7th request should be throttled
request = api_rf.get("/documents/1/")
response = DocumentAPIView.as_view()(request, pk=1)
assert response.status_code == 429
# After the 1h backoff wait time has passed, we can make a request again
mock_time.return_value += 1
request = api_rf.get("/documents/1/")
response = DocumentAPIView.as_view()(request, pk=1)
assert response.status_code == 200
@override_settings(AI_DOCUMENT_RATE_THROTTLE_RATES={"minute": 3, "hour": 6, "day": 10})
@patch("time.time")
def test_api_utils_ai_document_rate_throttle_day_limit(mock_time):
"""Test that day limit is enforced."""
api_rf = APIRequestFactory()
mock_time.return_value = 1000000
# Make requests to the document API, one per 10 minutes to avoid hitting
# the minute and hour limits
for _i in range(10): # 10 requests should be allowed
request = api_rf.get("/documents/1/")
response = DocumentAPIView.as_view()(request, pk=1)
assert response.status_code == 200
# Simulate passage of time
mock_time.return_value += 60 * 10
# Simulate passage of time
mock_time.return_value += 24 * 3600 - 10 * 60 * 10 - 1
# 11th request should be throttled
request = api_rf.get("/documents/1/")
response = DocumentAPIView.as_view()(request, pk=1)
assert response.status_code == 429
# After the 24h backoff wait time has passed we can make a request again
mock_time.return_value += 1
request = api_rf.get("/documents/1/")
response = DocumentAPIView.as_view()(request, pk=1)
assert response.status_code == 200

146
src/backend/core/tests/test_api_utils_ai_user_rate_throttles.py Normal file
View File

@@ -0,0 +1,146 @@
"""
Test throttling on users for the AI endpoint.
"""
from unittest.mock import patch
from uuid import uuid4
from django.core.cache import cache
from django.test import override_settings
import pytest
from rest_framework.response import Response
from rest_framework.test import APIRequestFactory
from rest_framework.views import APIView
from core.api.utils import AIUserRateThrottle
from core.factories import UserFactory
pytestmark = pytest.mark.django_db
class DocumentAPIView(APIView):
"""A simple view to test the throttle"""
throttle_classes = [AIUserRateThrottle]
def get(self, request, *args, **kwargs):
"""Minimal get method for testing purposes."""
return Response({"message": "Success"})
@pytest.fixture(autouse=True)
def clear_cache():
"""Fixture to clear the cache before each test."""
cache.clear()
@override_settings(AI_USER_RATE_THROTTLE_RATES={"minute": 3, "hour": 6, "day": 10})
@patch("time.time")
def test_api_utils_ai_user_rate_throttle_minute_limit(mock_time):
"""Test that minute limit is enforced."""
user = UserFactory()
api_rf = APIRequestFactory()
mock_time.return_value = 1000000
# Simulate requests to the document API
for _i in range(3): # 3 first requests should be allowed
document_id = str(uuid4())
request = api_rf.get(f"/documents/{document_id:s}/")
request.user = user
response = DocumentAPIView.as_view()(request, pk=document_id)
assert response.status_code == 200
# Simulate passage of time
mock_time.return_value += 59
# 4th request should be throttled
document_id = str(uuid4())
request = api_rf.get(f"/documents/{document_id:s}/")
request.user = user
response = DocumentAPIView.as_view()(request, pk=document_id)
assert response.status_code == 429
# After the 60s backoff wait time has passed, we can make a request again
mock_time.return_value += 1
document_id = str(uuid4())
request = api_rf.get(f"/documents/{document_id:s}/")
request.user = user
response = DocumentAPIView.as_view()(request, pk=document_id)
assert response.status_code == 200
@override_settings(AI_USER_RATE_THROTTLE_RATES={"minute": 100000, "hour": 6, "day": 10})
@patch("time.time")
def test_ai_user_rate_throttle_hour_limit(mock_time):
"""Test that the hour limit is enforced without hitting the minute limit."""
user = UserFactory()
api_rf = APIRequestFactory()
mock_time.return_value = 1000000
# Make requests to the document API, one per 21 seconds to avoid hitting the minute limit
for _i in range(6):
document_id = str(uuid4())
request = api_rf.get(f"/documents/{document_id:s}/")
request.user = user
response = DocumentAPIView.as_view()(request, pk=document_id)
assert response.status_code == 200
# Simulate passage of time
mock_time.return_value += 21
# Simulate passage of time
mock_time.return_value += 3600 - 6 * 21 - 1
# 7th request should be throttled
request = api_rf.get(f"/documents/{document_id:s}/")
request.user = user
response = DocumentAPIView.as_view()(request, pk=document_id)
assert response.status_code == 429
# After the 1h backoff wait time has passed, we can make a request again
mock_time.return_value += 1
request = api_rf.get(f"/documents/{document_id:s}/")
request.user = user
response = DocumentAPIView.as_view()(request, pk=document_id)
assert response.status_code == 200
@override_settings(AI_USER_RATE_THROTTLE_RATES={"minute": 3, "hour": 6, "day": 10})
@patch("time.time")
def test_api_utils_ai_user_rate_throttle_day_limit(mock_time):
"""Test that day limit is enforced."""
user = UserFactory()
api_rf = APIRequestFactory()
mock_time.return_value = 1000000
# Make requests to the document API, one per 10 minutes to avoid hitting
# the minute and hour limits
for _i in range(10): # 10 requests should be allowed
document_id = str(uuid4())
request = api_rf.get(f"/documents/{document_id:s}/")
request.user = user
response = DocumentAPIView.as_view()(request, pk=document_id)
assert response.status_code == 200
# Simulate passage of time
mock_time.return_value += 60 * 10
# Simulate passage of time
mock_time.return_value += 24 * 3600 - 10 * 60 * 10 - 1
# 11th request should be throttled
request = api_rf.get(f"/documents/{document_id:s}/")
request.user = user
response = DocumentAPIView.as_view()(request, pk=document_id)
assert response.status_code == 429
# After the 24h backoff wait time has passed we can make a request again
mock_time.return_value += 1
request = api_rf.get(f"/documents/{document_id:s}/")
request.user = user
response = DocumentAPIView.as_view()(request, pk=document_id)
assert response.status_code == 200

16
src/backend/core/tests/test_models_documents.py
View File

@@ -83,6 +83,8 @@ def test_models_documents_get_abilities_forbidden(is_authenticated, reach, role)
user = factories.UserFactory() if is_authenticated else AnonymousUser()
abilities = document.get_abilities(user)
assert abilities == {
"ai_transform": False,
"ai_translate": False,
"attachment_upload": False,
"link_configuration": False,
"destroy": False,
@@ -113,6 +115,8 @@ def test_models_documents_get_abilities_reader(is_authenticated, reach):
user = factories.UserFactory() if is_authenticated else AnonymousUser()
abilities = document.get_abilities(user)
assert abilities == {
"ai_transform": False,
"ai_translate": False,
"attachment_upload": False,
"destroy": False,
"link_configuration": False,
@@ -143,6 +147,8 @@ def test_models_documents_get_abilities_editor(is_authenticated, reach):
user = factories.UserFactory() if is_authenticated else AnonymousUser()
abilities = document.get_abilities(user)
assert abilities == {
"ai_transform": True,
"ai_translate": True,
"attachment_upload": True,
"destroy": False,
"link_configuration": False,
@@ -162,6 +168,8 @@ def test_models_documents_get_abilities_owner():
access = factories.UserDocumentAccessFactory(role="owner", user=user)
abilities = access.document.get_abilities(access.user)
assert abilities == {
"ai_transform": True,
"ai_translate": True,
"attachment_upload": True,
"destroy": True,
"link_configuration": True,
@@ -180,6 +188,8 @@ def test_models_documents_get_abilities_administrator():
access = factories.UserDocumentAccessFactory(role="administrator")
abilities = access.document.get_abilities(access.user)
assert abilities == {
"ai_transform": True,
"ai_translate": True,
"attachment_upload": True,
"destroy": False,
"link_configuration": True,
@@ -201,6 +211,8 @@ def test_models_documents_get_abilities_editor_user(django_assert_num_queries):
abilities = access.document.get_abilities(access.user)
assert abilities == {
"ai_transform": True,
"ai_translate": True,
"attachment_upload": True,
"destroy": False,
"link_configuration": False,
@@ -224,6 +236,8 @@ def test_models_documents_get_abilities_reader_user(django_assert_num_queries):
abilities = access.document.get_abilities(access.user)
assert abilities == {
"ai_transform": False,
"ai_translate": False,
"attachment_upload": False,
"destroy": False,
"link_configuration": False,
@@ -248,6 +262,8 @@ def test_models_documents_get_abilities_preset_role(django_assert_num_queries):
abilities = access.document.get_abilities(access.user)
assert abilities == {
"ai_transform": False,
"ai_translate": False,
"attachment_upload": False,
"destroy": False,
"link_configuration": False,

104
src/backend/core/tests/test_services_ai_services.py Normal file
View File

@@ -0,0 +1,104 @@
"""
Test ai API endpoints in the impress core app.
"""
import json
from unittest.mock import MagicMock, patch
from django.core.exceptions import ImproperlyConfigured
from django.test.utils import override_settings
import pytest
from openai import OpenAIError
from core.services.ai_services import AIService
pytestmark = pytest.mark.django_db
@pytest.mark.parametrize(
"setting_name, setting_value",
[
("AI_BASE_URL", None),
("AI_API_KEY", None),
("AI_MODEL", None),
],
)
def test_api_ai_setting_missing(setting_name, setting_value):
"""Setting should be set"""
with override_settings(**{setting_name: setting_value}):
with pytest.raises(
ImproperlyConfigured,
match="AI configuration not set",
):
AIService()
@override_settings(
AI_BASE_URL="http://example.com", AI_API_KEY="test-key", AI_MODEL="test-model"
)
@patch("openai.resources.chat.completions.Completions.create")
def test_api_ai__client_error(mock_create):
"""Fail when the client raises an error"""
mock_create.side_effect = OpenAIError("Mocked client error")
with pytest.raises(
OpenAIError,
match="Mocked client error",
):
AIService().transform("hello", "prompt")
@override_settings(
AI_BASE_URL="http://example.com", AI_API_KEY="test-key", AI_MODEL="test-model"
)
@patch("openai.resources.chat.completions.Completions.create")
def test_api_ai__client_invalid_response(mock_create):
"""Fail when the client response is invalid"""
answer = {"no_answer": "This is an invalid response"}
mock_create.return_value = MagicMock(
choices=[MagicMock(message=MagicMock(content=json.dumps(answer)))]
)
with pytest.raises(
RuntimeError,
match="AI response does not contain an answer",
):
AIService().transform("hello", "prompt")
@override_settings(
AI_BASE_URL="http://example.com", AI_API_KEY="test-key", AI_MODEL="test-model"
)
@patch("openai.resources.chat.completions.Completions.create")
def test_api_ai__success(mock_create):
"""The AI request should work as expect when called with valid arguments."""
answer = '{"answer": "Salut"}'
mock_create.return_value = MagicMock(
choices=[MagicMock(message=MagicMock(content=answer))]
)
response = AIService().transform("hello", "prompt")
assert response == {"answer": "Salut"}
@override_settings(
AI_BASE_URL="http://example.com", AI_API_KEY="test-key", AI_MODEL="test-model"
)
@patch("openai.resources.chat.completions.Completions.create")
def test_api_ai__success_sanitize(mock_create):
"""The AI response should be sanitized"""
answer = '{"answer": "Salut\\n \tle \nmonde"}'
mock_create.return_value = MagicMock(
choices=[MagicMock(message=MagicMock(content=answer))]
)
response = AIService().transform("hello", "prompt")
assert response == {"answer": "Salut\n \tle \nmonde"}