Merge pull request #32 from Kalgoc/fix/expenses

anibalcontreras · web-flow · commit 316cb588d817 · 2024-07-01T19:26:16.000-04:00
Fix/expenses
diff --git a/.flake8 b/.flake8
@@ -1,5 +1,5 @@
 [flake8]
 exclude = .git,__pycache__,old,build,dist,.venv, */migrations/*
-ignore = E226,E302,E41,F401,F403,F405
+ignore = E226,E302,E41,F401,F403,F405,W503
 max-line-length = 120
 max-complexity = 10
diff --git a/AI/AI.py b/AI/AI.py
@@ -3,17 +3,28 @@
 import os
 from textwrap import dedent
 from dotenv import load_dotenv
+from categories.models import Category
 
 
-def setup_api_key():
-    env_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), ".env")
-    load_dotenv(dotenv_path=env_path)
-    openai.api_key = os.getenv("OPENAI_KEY")
+GPT_MODEL = "gpt-4-turbo"
+
+
+def get_category_name_from_description(description):
+    classified_category = classify_text(description)
+    cleaned_category = category_matched_with_id(classified_category)
+
+    if not Category.objects.filter(name=cleaned_category).exists():
+        raise ValueError("Categoría no encontrada.")
+    return cleaned_category
 
 
 def classify_text(texto):
+    setup_api_key()
+    if not openai.api_key:
+        raise ValueError("OpenAI API key is not set. Please check your .env file.")
+
     response = openai.ChatCompletion.create(
-        model="gpt-4-turbo",
+        model=GPT_MODEL,
         messages=[
             {
                 "role": "system",
@@ -43,7 +54,13 @@ def classify_text(texto):
     return category
 
 
-def clean_category(category):
+def setup_api_key():
+    env_path = os.path.join(os.path.dirname(os.path.dirname(__file__)), ".env")
+    load_dotenv(dotenv_path=env_path)
+    openai.api_key = os.getenv("OPENAI_KEY")
+
+
+def category_matched_with_id(category):
     if category.lower().startswith("categoría: "):
         category = category[11:]
     if category.endswith("."):
diff --git a/AI/__init__.py b/AI/__init__.py
@@ -1,2 +1,2 @@
 # flake8: noqa
-from .AI import setup_api_key, classify_text, clean_category
+from .AI import setup_api_key, classify_text, category_matched_with_id
diff --git a/AI/tests.py b/AI/tests.py
@@ -1,12 +1,11 @@
 # flake8: noqa
 import unittest
 from unittest.mock import patch
-from AI import classify_text, clean_category
+from AI import classify_text, category_matched_with_id
 from textwrap import dedent
 
 
 class TestGastosClassifier(unittest.TestCase):
-
     @patch("openai.ChatCompletion.create")
     def test_classify_text(self, mock_create):
         # Configurar el mock para devolver una respuesta simulada
@@ -50,13 +49,13 @@ def test_classify_text(self, mock_create):
 
     def test_clean_category(self):
         # Pruebas con diferentes casos
-        self.assertEqual(clean_category("Categoría: Comida."), "Comida")
-        self.assertEqual(clean_category("categoría: transporte"), "Transporte")
-        self.assertEqual(clean_category("vivienda."), "Vivienda")
+        self.assertEqual(category_matched_with_id("Categoría: Comida."), "Comida")
+        self.assertEqual(category_matched_with_id("categoría: transporte"), "Transporte")
+        self.assertEqual(category_matched_with_id("vivienda."), "Vivienda")
 
         # Prueba de una categoría no permitida
         with self.assertRaises(ValueError):
-            clean_category("Categoría: Viajes.")
+            category_matched_with_id("Categoría: Viajes.")
 
 
 if __name__ == "__main__":
diff --git a/authentication/services/cognito_service.py b/authentication/services/cognito_service.py
@@ -4,6 +4,7 @@
 import hashlib
 import base64
 from django.contrib.auth import get_user_model
+from user_expense_type.models import UserExpenseType
 
 
 class CognitoService:
@@ -38,6 +39,10 @@ def register_user(self, name, phone, email, password):
             User.objects.create_user(
                 username=email, email=email, password=password, first_name=name, phone=phone, user_id=cognito_uuid
             )
+            user_expense_type = UserExpenseType.objects.create(
+                username=cognito_uuid, set_by_user=False, name="Personal"
+            )
+            user_expense_type.save()
             return response
         except ClientError as e:
             raise e
diff --git a/categories/signals.py b/categories/signals.py
@@ -1,5 +1,3 @@
-# categories/signals.py
-
 from django.db.models.signals import post_migrate
 from django.dispatch import receiver
 from .models import Category
diff --git a/docker-compose.yml b/docker-compose.yml
@@ -32,6 +32,7 @@ services:
       POSTGRES_PASSWORD: ${DB_PASSWORD}
       POSTGRES_PORT: ${DB_PORT}
       POSTGRES_HOST: ${DB_HOST}
+      OPENAI_API_KEY: ${OPENAI_API_KEY}
       DJANGO_SETTINGS_MODULE: piggywallet.settings.dev
 
 networks:
diff --git a/expenses/migrations/0002_expense_description.py b/expenses/migrations/0002_expense_description.py
@@ -0,0 +1,17 @@
+# Generated by Django 5.0.6 on 2024-06-25 02:31
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("expenses", "0001_initial"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="expense",
+            name="description",
+            field=models.CharField(blank=True, max_length=255, null=True),
+        ),
+    ]
diff --git a/expenses/migrations/0003_merge_20240701_2317.py b/expenses/migrations/0003_merge_20240701_2317.py
@@ -0,0 +1,12 @@
+# Generated by Django 5.0.6 on 2024-07-01 23:17
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("expenses", "0002_alter_expense_bankcard_id_alter_expense_username"),
+        ("expenses", "0002_expense_description"),
+    ]
+
+    operations = []
diff --git a/expenses/models.py b/expenses/models.py
@@ -9,6 +9,7 @@
 class Expense(models.Model):
     id = models.AutoField(primary_key=True)
     username = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.CASCADE, to_field="user_id")
+    description = models.CharField(max_length=255, blank=True, null=True)
     user_expense_type = models.ForeignKey(UserExpenseType, on_delete=models.CASCADE)
     category = models.ForeignKey(Category, on_delete=models.CASCADE)
     bankcard_id = models.ForeignKey(BankCard, on_delete=models.CASCADE)
diff --git a/expenses/serializers.py b/expenses/serializers.py
@@ -9,15 +9,42 @@
 class ExpenseSerializer(serializers.Serializer):
     id = serializers.IntegerField(read_only=True)
     username = serializers.PrimaryKeyRelatedField(queryset=get_user_model().objects.all())
-    user_expense_type = serializers.PrimaryKeyRelatedField(queryset=UserExpenseType.objects.all())
+    user_expense_type = serializers.PrimaryKeyRelatedField(queryset=UserExpenseType.objects.all(), required=False)
     category = serializers.PrimaryKeyRelatedField(queryset=Category.objects.all())
     bankcard_id = serializers.PrimaryKeyRelatedField(queryset=BankCard.objects.all())
     amount = serializers.IntegerField()
+    description = serializers.CharField(max_length=255, allow_blank=True, allow_null=True)
+
+    def validate(self, data):
+        user_expense_type = data.get("user_expense_type")
+        username = data.get("username")
+        if (
+            user_expense_type
+            and not UserExpenseType.objects.filter(id=user_expense_type.id, username=username).exists()
+        ):
+            raise serializers.ValidationError(
+                {"user_expense_type": "This user_expense_type does not belong to the specified user."}
+            )
+
+        return data
 
     def create(self, validated_data):
+        user_id = validated_data.get("username")
+        if "user_expense_type" not in validated_data:
+            default_expense_type = UserExpenseType.objects.filter(
+                username=user_id, name="Personal", set_by_user=False
+            ).first()
+            if not default_expense_type:
+                raise serializers.ValidationError("Default personal expense type not found for the user.")
+            validated_data["user_expense_type"] = default_expense_type
+
         return Expense.objects.create(**validated_data)
 
     def update(self, instance, validated_data):
         instance.amount = validated_data.get("amount", instance.amount)
+        instance.description = validated_data.get("description", instance.description)
+        instance.bankcard_id = validated_data.get("bankcard_id", instance.bankcard_id)
+        instance.user_expense_type = validated_data.get("user_expense_type", instance.user_expense_type)
+        instance.category = validated_data.get("category", instance.category)
         instance.save()
         return instance
diff --git a/expenses/services/categorize_expense.py b/expenses/services/categorize_expense.py
@@ -0,0 +1,17 @@
+from rest_framework.response import Response
+from rest_framework import status
+from categories.models import Category
+from AI.AI import get_category_name_from_description
+
+
+def categorize_expense_description(data):
+    if "description" in data:
+        try:
+            category_name = get_category_name_from_description(data["description"])
+            category_obj = Category.objects.get(name=category_name)
+            data["category"] = category_obj.id
+            return data, None
+        except ValueError as e:
+            return None, Response({"error": str(e)}, status=status.HTTP_400_BAD_REQUEST)
+    else:
+        return None, Response({"error": "Description is required"}, status=status.HTTP_400_BAD_REQUEST)
diff --git a/expenses/views.py b/expenses/views.py
@@ -1,7 +1,5 @@
 from rest_framework import viewsets, status
 from rest_framework.response import Response
-
-# from .models import Expense
 from .models import Expense
 from user_expense_type.models import UserExpenseType
 from categories.models import Category
@@ -10,6 +8,8 @@
 import jwt
 from django.db.models import Sum
 from django.utils.timezone import now
+from authentication.utils import get_user_id_from_token
+from expenses.services.categorize_expense import categorize_expense_description
 
 
 class ExpenseViewSet(viewsets.ViewSet):
@@ -35,11 +35,9 @@ def get_user_id_from_token(self, request):
     @cognito_authenticated
     def list(self, request):
         try:
-            username = self.get_user_id_from_token(request)
+            username = get_user_id_from_token(request)
             expenses = Expense.objects.filter(username=username)
             serializer = ExpenseSerializer(expenses, many=True)
-            for ser in serializer.data:
-                ser.pop("username")
             return Response(serializer.data)
         except Expense.DoesNotExist:
             return Response({"error": "Expenses not found"}, status=status.HTTP_404_NOT_FOUND)
@@ -49,16 +47,10 @@ def list(self, request):
     @cognito_authenticated
     def retrieve(self, request, pk=None):
         try:
-            username = self.get_user_id_from_token(request)
+            username = get_user_id_from_token(request)
             expense = Expense.objects.get(id=pk, username=username)
             serializer = ExpenseSerializer(expense)
-            response = {
-                "user_expense_type": serializer.data["user_expense_type"],
-                "category": serializer.data["category"],
-                "bankcard_id": serializer.data["bankcard_id"],
-                "amount": serializer.data["amount"],
-            }
-            return Response(data=response)
+            return Response(serializer.data)
         except Expense.DoesNotExist:
             return Response({"error": "Expense not found"}, status=status.HTTP_404_NOT_FOUND)
         except Exception as e:
@@ -67,28 +59,26 @@ def retrieve(self, request, pk=None):
     @cognito_authenticated
     def create(self, request):
         try:
-            username = self.get_user_id_from_token(request)
+            username = get_user_id_from_token(request)
             data = request.data.copy()
             data["username"] = username
 
+            data, error_response = categorize_expense_description(data)
+            if error_response:
+                return error_response
+
             serializer = ExpenseSerializer(data=data)
             if serializer.is_valid():
                 serializer.save()
-                response = {
-                    "user_expense_type": serializer.data["user_expense_type"],
-                    "category": serializer.data["category"],
-                    "bankcard_id": serializer.data["bankcard_id"],
-                    "amount": serializer.data["amount"],
-                }
-                return Response(data=response, status=status.HTTP_201_CREATED)
+                return Response(serializer.data, status=status.HTTP_201_CREATED)
             return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
         except Exception as e:
             return Response({"error": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
 
     @cognito_authenticated
     def destroy(self, request, pk=None):
         try:
-            username = self.get_user_id_from_token(request)
+            username = get_user_id_from_token(request)
             expense = Expense.objects.get(id=pk, username=username)
             expense.delete()
             return Response(status=status.HTTP_204_NO_CONTENT)
@@ -100,49 +90,24 @@ def destroy(self, request, pk=None):
     @cognito_authenticated
     def partial_update(self, request, pk=None):
         try:
-            username = self.get_user_id_from_token(request)
+            username = get_user_id_from_token(request)
             expense = Expense.objects.get(id=pk, username=username)
             serializer = ExpenseSerializer(expense, data=request.data, partial=True)
-            if serializer.is_valid():
-                serializer.save()
-                response = {
-                    "user_expense_type": serializer.data["user_expense_type"],
-                    "category": serializer.data["category"],
-                    "bankcard_id": serializer.data["bankcard_id"],
-                    "amount": serializer.data["amount"],
-                }
-                return Response(data=response)
-            return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)
+            serializer.is_valid(raise_exception=True)
+            serializer.save()
+            return Response(serializer.data, status=status.HTTP_200_OK)
+
         except Expense.DoesNotExist:
             return Response({"error": "Expense not found"}, status=status.HTTP_404_NOT_FOUND)
         except Exception as e:
             return Response({"error": str(e)}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
 
 
 class ExpenseGroupedByTypeAndCategoryViewSet(viewsets.ViewSet):
-    def get_user_id_from_token(self, request):
-        try:
-            authorization_header = request.headers.get("Authorization")
-            if not authorization_header:
-                raise Exception("Authorization header not found")
-
-            token = authorization_header.split()[1]
-            decoded_token = jwt.decode(token, options={"verify_signature": False})
-            username = decoded_token.get("username")
-            if not username:
-                raise Exception("User ID not found in token")
-            return username
-        except jwt.DecodeError:
-            raise Exception("Invalid token")
-        except jwt.ExpiredSignatureError:
-            raise Exception("Expired token")
-        except Exception as e:
-            raise Exception(f"Error decoding token: {e}")
-
     @cognito_authenticated
     def list(self, request):
         try:
-            username = self.get_user_id_from_token(request)
+            username = get_user_id_from_token(request)
             expenses_grouped = {}
             today = now()
             expenses = (
diff --git a/user_expense_type/migrations/0002_alter_userexpensetype_set_by_user.py b/user_expense_type/migrations/0002_alter_userexpensetype_set_by_user.py
@@ -0,0 +1,17 @@
+# Generated by Django 5.0.6 on 2024-06-25 04:06
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("user_expense_type", "0001_initial"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="userexpensetype",
+            name="set_by_user",
+            field=models.BooleanField(default=True),
+        ),
+    ]
diff --git a/user_expense_type/models.py b/user_expense_type/models.py
@@ -7,6 +7,6 @@ class UserExpenseType(models.Model):
     username = models.UUIDField()
     name = models.CharField(max_length=70, default="Personal")
     description = models.CharField(max_length=255, blank=True, null=True)
-    set_by_user = models.BooleanField(default=False)
+    set_by_user = models.BooleanField(default=True)
     created_at = models.DateTimeField(auto_now_add=True)
     updated_at = models.DateTimeField(auto_now=True)
diff --git a/user_expense_type/serializers.py b/user_expense_type/serializers.py
diff --git a/user_expense_type/views.py b/user_expense_type/views.py

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`	`1`	`# flake8: noqa`
`2`		`-from .AI import setup_api_key, classify_text, clean_category`
	`2`	`+from .AI import setup_api_key, classify_text, category_matched_with_id`
Original file line number	Diff line number	Diff line change
`@@ -1,5 +1,3 @@`
`1`		`-# categories/signals.py`
`2`		`-`
`3`	`1`	`from django.db.models.signals import post_migrate`
`4`	`2`	`from django.dispatch import receiver`
`5`	`3`	`from .models import Category`