diff --git a/cookbook/helper/ai_helper.py b/cookbook/helper/ai_helper.py index e8cbde3d4..584f9a11a 100644 --- a/cookbook/helper/ai_helper.py +++ b/cookbook/helper/ai_helper.py @@ -1,3 +1,5 @@ +from decimal import Decimal + from django.utils import timezone from django.db.models import Sum from litellm import CustomLogger @@ -55,19 +57,15 @@ class AiCallbackHandler(CustomLogger): if self.ai_provider.log_credit_cost: credit_cost = kwargs.get("response_cost", 0) * 100 - print(not has_monthly_token(self.space) , self.space.ai_credits_balance > 0, not has_monthly_token(self.space) and self.space.ai_credits_balance > 0) if (not has_monthly_token(self.space)) and self.space.ai_credits_balance > 0: - print('taking credits from balance') - self.space.ai_credits_balance = max(0, self.space.ai_credits_balance - credit_cost) - print('setting from balance to true') - credits_from_balance = True - print('saving space') - self.space.save() - print('done') - else: - print('not taking credits from balance') + remaining_balance = self.space.ai_credits_balance - Decimal(str(credit_cost)) + if remaining_balance < 0: + remaining_balance = 0 + + self.space.ai_credits_balance = remaining_balance + credits_from_balance = True + self.space.save() - print('creating AI log with credit cost ', credit_cost , ' from balance: ', credits_from_balance) AiLog.objects.create( created_by=self.user, space=self.space, diff --git a/cookbook/helper/permission_helper.py b/cookbook/helper/permission_helper.py index 4ddcb5a60..5a131bffa 100644 --- a/cookbook/helper/permission_helper.py +++ b/cookbook/helper/permission_helper.py @@ -333,6 +333,9 @@ class CustomRecipePermission(permissions.BasePermission): class CustomAiProviderPermission(permissions.BasePermission): """ Custom permission class for the AiProvider api endpoint + users: can read all + admins: can read and write + superusers: can read and write + write providers without a space """ message = _('You do not have the required permissions to view this page!') diff --git a/cookbook/migrations/0226_aiprovider_log_credit_cost_and_more.py b/cookbook/migrations/0226_aiprovider_log_credit_cost_and_more.py index 2f1d1dfa8..7b8ca2204 100644 --- a/cookbook/migrations/0226_aiprovider_log_credit_cost_and_more.py +++ b/cookbook/migrations/0226_aiprovider_log_credit_cost_and_more.py @@ -18,6 +18,6 @@ class Migration(migrations.Migration): migrations.AlterField( model_name='space', name='ai_credits_monthly', - field=models.IntegerField(default=100), + field=models.IntegerField(default=10000), ), ] diff --git a/cookbook/tests/api/test_api_ai_provider.py b/cookbook/tests/api/test_api_ai_provider.py new file mode 100644 index 000000000..173d986e9 --- /dev/null +++ b/cookbook/tests/api/test_api_ai_provider.py @@ -0,0 +1,162 @@ +import json + +import pytest +from django.urls import reverse +from django_scopes import scopes_disabled + +from cookbook.models import MealType, PropertyType, AiProvider + +LIST_URL = 'api:aiprovider-list' +DETAIL_URL = 'api:aiprovider-detail' + + +@pytest.fixture() +def obj_1(space_1, a1_s1): + return AiProvider.objects.get_or_create(name='test_1', space=space_1)[0] + + +@pytest.fixture +def obj_2(space_1, a1_s1): + return AiProvider.objects.get_or_create(name='test_2', space=None)[0] + + +@pytest.mark.parametrize("arg", [ + ['a_u', 403], + ['g1_s1', 403], + ['u1_s1', 200], + ['a1_s1', 200], +]) +def test_list_permission(arg, request): + c = request.getfixturevalue(arg[0]) + assert c.get(reverse(LIST_URL)).status_code == arg[1] + + +def test_list_space(obj_1, obj_2, u1_s1, u1_s2, space_2): + assert json.loads(u1_s1.get(reverse(LIST_URL)).content)['count'] == 2 + assert json.loads(u1_s2.get(reverse(LIST_URL)).content)['count'] == 1 + + obj_1.space = space_2 + obj_1.save() + + assert json.loads(u1_s1.get(reverse(LIST_URL)).content)['count'] == 1 + assert json.loads(u1_s2.get(reverse(LIST_URL)).content)['count'] == 2 + + +@pytest.mark.parametrize("arg", [ + ['a_u', 403], + ['g1_s1', 403], + ['u1_s1', 403], + ['a1_s1', 200], + ['g1_s2', 403], + ['u1_s2', 404], + ['a1_s2', 403], +]) +def test_update(arg, request, obj_1): + c = request.getfixturevalue(arg[0]) + r = c.patch( + reverse( + DETAIL_URL, + args={obj_1.id} + ), + {'name': 'new'}, + content_type='application/json' + ) + response = json.loads(r.content) + assert r.status_code == arg[1] + if r.status_code == 200: + assert response['name'] == 'new' + + +@pytest.mark.parametrize("arg", [ + ['a_u', 403], + ['g1_s1', 403], + ['u1_s1', 403], + ['a1_s1', 403], + ['g1_s2', 403], + ['u1_s2', 403], + ['a1_s2', 403], + ['s1_s1', 200], +]) +def test_update_global(arg, request, obj_2): + c = request.getfixturevalue(arg[0]) + r = c.patch( + reverse( + DETAIL_URL, + args={obj_2.id} + ), + {'name': 'new'}, + content_type='application/json' + ) + response = json.loads(r.content) + assert r.status_code == arg[1] + if r.status_code == 200: + assert response['name'] == 'new' + + +@pytest.mark.parametrize("arg", [ + ['a_u', 403], + ['g1_s1', 403], + ['u1_s1', 403], + ['a1_s1', 201], +]) +def test_add(arg, request, u1_s2): + c = request.getfixturevalue(arg[0]) + r = c.post( + reverse(LIST_URL), + {'name': 'test', 'api_key': 'test', 'model_name': 'test'}, + content_type='application/json' + ) + response = json.loads(r.content) + assert r.status_code == arg[1] + if r.status_code == 201: + assert response['name'] == 'test' + r = c.get(reverse(DETAIL_URL, args={response['id']})) + assert r.status_code == 200 + r = u1_s2.get(reverse(DETAIL_URL, args={response['id']})) + assert r.status_code == 404 + + +def test_delete(a1_s1, a1_s2, obj_1): + # admins cannot delete foreign space providers + r = a1_s2.delete( + reverse( + DETAIL_URL, + args={obj_1.id} + ) + ) + assert r.status_code == 404 + + # admins can delete their space providers + r = a1_s1.delete( + reverse( + DETAIL_URL, + args={obj_1.id} + ) + ) + + assert r.status_code == 204 + with scopes_disabled(): + assert AiProvider.objects.count() == 0 + + +def test_delete_global(a1_s1, s1_s1, obj_2): + # admins cant delete global providers + r = a1_s1.delete( + reverse( + DETAIL_URL, + args={obj_2.id} + ) + ) + assert r.status_code == 403 + + # superusers can delete global providers + r = s1_s1.delete( + reverse( + DETAIL_URL, + args={obj_2.id} + ) + ) + + assert r.status_code == 204 + with scopes_disabled(): + assert AiProvider.objects.count() == 0 diff --git a/cookbook/tests/conftest.py b/cookbook/tests/conftest.py index d57c303c6..6e1be2620 100644 --- a/cookbook/tests/conftest.py +++ b/cookbook/tests/conftest.py @@ -298,3 +298,11 @@ def a1_s2(client, space_2): @pytest.fixture() def a2_s2(client, space_2): return create_user(client, space_2, group='admin') + +@pytest.fixture() +def s1_s1(client, space_1): + client = create_user(client, space_1, group='admin') + user = auth.get_user(client) + user.is_superuser = True + user.save() + return client diff --git a/recipes/settings.py b/recipes/settings.py index fadef2ace..7ee5bfd51 100644 --- a/recipes/settings.py +++ b/recipes/settings.py @@ -60,7 +60,7 @@ SPACE_DEFAULT_MAX_USERS = int(os.getenv('SPACE_DEFAULT_MAX_USERS', 0)) SPACE_DEFAULT_MAX_FILES = int(os.getenv('SPACE_DEFAULT_MAX_FILES', 0)) SPACE_DEFAULT_ALLOW_SHARING = extract_bool('SPACE_DEFAULT_ALLOW_SHARING', True) SPACE_AI_ENABLED = extract_bool('SPACE_AI_ENABLED', True) -SPACE_AI_CREDITS_MONTHLY = int(os.getenv('SPACE_AI_CREDITS_MONTHLY', 100)) +SPACE_AI_CREDITS_MONTHLY = int(os.getenv('SPACE_AI_CREDITS_MONTHLY', 10000)) INTERNAL_IPS = extract_comma_list('INTERNAL_IPS', '127.0.0.1') diff --git a/vue3/src/components/settings/SpaceSettings.vue b/vue3/src/components/settings/SpaceSettings.vue index 80c55decd..c7b3fee30 100644 --- a/vue3/src/components/settings/SpaceSettings.vue +++ b/vue3/src/components/settings/SpaceSettings.vue @@ -96,8 +96,8 @@ {{ $t('Save') }} diff --git a/vue3/src/pages/ModelListPage.vue b/vue3/src/pages/ModelListPage.vue index efb474b0a..ad8639acc 100644 --- a/vue3/src/pages/ModelListPage.vue +++ b/vue3/src/pages/ModelListPage.vue @@ -35,7 +35,8 @@ {{$t('MonthlyCreditsUsed')}} ({{ useUserPreferenceStore().activeSpace.aiMonthlyCreditsUsed }} / {{ useUserPreferenceStore().activeSpace.aiCreditsMonthly }}) - + {{$t('AiCreditsBalance')}} : {{useUserPreferenceStore().activeSpace.aiCreditsBalance}} +