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

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