mirror of
https://github.com/TandoorRecipes/recipes.git
synced 2025-12-24 02:39:20 -05:00
ai system improvements
This commit is contained in:
@@ -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,
|
||||
|
||||
@@ -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!')
|
||||
|
||||
|
||||
@@ -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),
|
||||
),
|
||||
]
|
||||
|
||||
162
cookbook/tests/api/test_api_ai_provider.py
Normal file
162
cookbook/tests/api/test_api_ai_provider.py
Normal file
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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')
|
||||
|
||||
|
||||
@@ -96,8 +96,8 @@
|
||||
<template v-if="space.aiEnabled">
|
||||
<model-select model="AiProvider" :label="$t('Default')" v-model="space.aiDefaultProvider"></model-select>
|
||||
|
||||
<v-number-input v-model="space.aiCreditsMonthly" :label="$t('MonthlyCredits')" :disabled="!useUserPreferenceStore().userSettings.user.isSuperuser"></v-number-input>
|
||||
<v-number-input v-model="space.aiCreditsBalance" :label="$t('AiCreditsBalance')" :disabled="!useUserPreferenceStore().userSettings.user.isSuperuser"></v-number-input>
|
||||
<v-number-input v-model="space.aiCreditsMonthly" :precision="2" :label="$t('MonthlyCredits')" :disabled="!useUserPreferenceStore().userSettings.user.isSuperuser"></v-number-input>
|
||||
<v-number-input v-model="space.aiCreditsBalance" :precision="4" :label="$t('AiCreditsBalance')" :disabled="!useUserPreferenceStore().userSettings.user.isSuperuser"></v-number-input>
|
||||
|
||||
</template>
|
||||
<v-btn color="success" @click="updateSpace()" prepend-icon="$save">{{ $t('Save') }}</v-btn>
|
||||
|
||||
@@ -35,7 +35,8 @@
|
||||
|
||||
<v-card-text v-if="genericModel.model.name == 'AiLog'">
|
||||
{{$t('MonthlyCreditsUsed')}} ({{ useUserPreferenceStore().activeSpace.aiMonthlyCreditsUsed }} / {{ useUserPreferenceStore().activeSpace.aiCreditsMonthly }})
|
||||
<v-progress-linear :model-value="useUserPreferenceStore().activeSpace.aiMonthlyCreditsUsed"></v-progress-linear>
|
||||
{{$t('AiCreditsBalance')}} : {{useUserPreferenceStore().activeSpace.aiCreditsBalance}}
|
||||
<v-progress-linear :model-value="useUserPreferenceStore().activeSpace.aiMonthlyCreditsUsed" :max="useUserPreferenceStore().activeSpace.aiCreditsMonthly"></v-progress-linear>
|
||||
</v-card-text>
|
||||
</v-card>
|
||||
</v-col>
|
||||
|
||||
Reference in New Issue
Block a user