diff --git a/cookbook/helper/ai_helper.py b/cookbook/helper/ai_helper.py index 5ac52cb5c..b614faa5c 100644 --- a/cookbook/helper/ai_helper.py +++ b/cookbook/helper/ai_helper.py @@ -8,7 +8,10 @@ def get_monthly_token_usage(space): """ returns the number of credits the space has used in the current month """ - return AiLog.objects.filter(space=space, credits_from_balance=False, created_at__month=timezone.now().month).aggregate(Sum('credit_cost'))['credit_cost__sum'] + token_usage = AiLog.objects.filter(space=space, credits_from_balance=False, created_at__month=timezone.now().month).aggregate(Sum('credit_cost'))['credit_cost__sum'] + if token_usage is None: + token_usage = 0 + return token_usage def has_monthly_token(space): @@ -16,3 +19,7 @@ def has_monthly_token(space): checks if the monthly credit limit has been exceeded """ return get_monthly_token_usage(space) < space.ai_credits_monthly + + +def can_perform_ai_request(space): + return has_monthly_token(space) and space.ai_enabled \ No newline at end of file diff --git a/cookbook/migrations/0224_space_ai_credits_balance_space_ai_credits_monthly_and_more.py b/cookbook/migrations/0224_space_ai_credits_balance_space_ai_credits_monthly_and_more.py index 518afca49..916aa77bd 100644 --- a/cookbook/migrations/0224_space_ai_credits_balance_space_ai_credits_monthly_and_more.py +++ b/cookbook/migrations/0224_space_ai_credits_balance_space_ai_credits_monthly_and_more.py @@ -22,7 +22,7 @@ class Migration(migrations.Migration): migrations.AddField( model_name='space', name='ai_credits_monthly', - field=models.IntegerField(default=0), + field=models.IntegerField(default=100), ), migrations.CreateModel( name='AiProvider', diff --git a/cookbook/models.py b/cookbook/models.py index e36e3e570..4b870a844 100644 --- a/cookbook/models.py +++ b/cookbook/models.py @@ -345,6 +345,9 @@ class Space(ExportModelOperationsMixin('space'), models.Model): BookmarkletImport.objects.filter(space=self).delete() CustomFilter.objects.filter(space=self).delete() + AiLog.objects.filter(space=self).delete() + AiProvider.objects.filter(space=self).delete() + Property.objects.filter(space=self).delete() PropertyType.objects.filter(space=self).delete() diff --git a/cookbook/serializer.py b/cookbook/serializer.py index 2caecda6b..70cce2a21 100644 --- a/cookbook/serializer.py +++ b/cookbook/serializer.py @@ -404,21 +404,33 @@ class SpacedModelSerializer(serializers.ModelSerializer): class AiProviderSerializer(serializers.ModelSerializer): + api_key = serializers.CharField(required=False, write_only=True) def create(self, validated_data): - if not self.context['request'].user.is_superuser: - validated_data['space'] = self.context['request'].space + validated_data = self.handle_global_space_logic(validated_data) + return super().create(validated_data) + def update(self, instance, validated_data): + validated_data = self.handle_global_space_logic(validated_data) + return super().update(instance, validated_data) + + def handle_global_space_logic(self, validated_data): + """ + allow superusers to create AI providers without a space but make sure everyone else only uses their own space + """ + if ('space' not in validated_data or not validated_data['space']) and self.context['request'].user.is_superuser: + validated_data['space'] = None + else: + validated_data['space'] = self.context['request'].space + + return validated_data + class Meta: model = AiProvider fields = ('id', 'name', 'description', 'api_key', 'model_name', 'url', 'log_credit_cost', 'space', 'created_at', 'updated_at') read_only_fields = ('created_at', 'updated_at',) - extra_kwargs = { - 'api_key': {'write_only': True}, - } - class AiLogSerializer(serializers.ModelSerializer): ai_provider = AiProviderSerializer(read_only=True) diff --git a/cookbook/views/api.py b/cookbook/views/api.py index b7f4f26b0..ae30b15eb 100644 --- a/cookbook/views/api.py +++ b/cookbook/views/api.py @@ -65,7 +65,7 @@ from cookbook.connectors.connector_manager import ConnectorManager, ActionType from cookbook.forms import ImportForm, ImportExportBase from cookbook.helper import recipe_url_import as helper from cookbook.helper.HelperFunctions import str2bool, validate_import_url -from cookbook.helper.ai_helper import has_monthly_token +from cookbook.helper.ai_helper import has_monthly_token, can_perform_ai_request from cookbook.helper.image_processing import handle_image from cookbook.helper.ingredient_parser import IngredientParser from cookbook.helper.open_data_importer import OpenDataImporter @@ -2032,10 +2032,10 @@ class AiImportView(APIView): } return Response(RecipeFromSourceResponseSerializer(context={'request': request}).to_representation(response), status=status.HTTP_400_BAD_REQUEST) - if not has_monthly_token(request.space): + if not can_perform_ai_request(request.space): response = { 'error': True, - 'msg': _("You don't have any credits remaining to use AI."), + 'msg': _("You don't have any credits remaining to use AI or AI features are not enabled for your space."), } return Response(RecipeFromSourceResponseSerializer(context={'request': request}).to_representation(response), status=status.HTTP_400_BAD_REQUEST) @@ -2129,9 +2129,15 @@ class AiImportView(APIView): return Response(RecipeFromSourceResponseSerializer(context={'request': request}).to_representation(response), status=status.HTTP_400_BAD_REQUEST) try: - ai_response = completion(api_key=ai_provider.api_key, - model=ai_provider.model_name, - response_format={"type": "json_object"}, messages=messages, ) + ai_request = { + 'api_key': ai_provider.api_key, + 'model': ai_provider.model_name, + 'response_format': {"type": "json_object"}, + 'messages': messages, + } + if ai_provider.url: + ai_request['api_base'] = ai_provider.url + ai_response = completion(**ai_request) except BadRequestError as err: response = { 'error': True, diff --git a/vue3/src/components/inputs/ModelSelect.vue b/vue3/src/components/inputs/ModelSelect.vue index 7decee941..3be13e2e8 100644 --- a/vue3/src/components/inputs/ModelSelect.vue +++ b/vue3/src/components/inputs/ModelSelect.vue @@ -58,7 +58,6 @@ -