Merge pull request #12 from Kalgoc/AI

AI

Author: sofiagiagnoni

.gitignore

@@ -2,10 +2,10 @@
 __pycache__/
 *.py[cod]
 *$py.class
-
+*.env
 # C extensions
 *.so
-
+*.csv
 # Distribution / packaging
 .Python
 build/

AI/AI.py

@@ -0,0 +1,51 @@
+import openai
+import os
+
+
+def setup_api_key():
+    openai.api_key = os.environ["OPENAI_KEY"]
+
+
+def classify_text(texto):
+    response = openai.ChatCompletion.create(
+        model="gpt-4-turbo",
+        messages=[
+            {
+                "role": "system",
+                "content": "Eres un modelo que clasifica gastos en \
+                categorías predefinidas.",
+            },
+            {
+                "role": "user",
+                "content": f"Clasifica el siguiente gasto en una de las \
+              categorías: comida, transporte, vivienda, educación, salud, \
+                 entretenimiento,\
+             ahorro, inversión.\n\nTexto: {texto}\n\nCategoría:",
+            },
+        ],
+        max_tokens=15,
+        temperature=0,
+    )
+    category = response.choices[0].message["content"].strip()
+    return category
+
+
+def clean_category(category):
+    if category.lower().startswith("categoría: "):
+        category = category[11:]
+    if category.endswith("."):
+        category = category[:-1]
+    category = category.capitalize()
+    categories_allowed = [
+        "Comida",
+        "Transporte",
+        "Vivienda",
+        "Educación",
+        "Salud",
+        "Entretenimiento",
+        "Ahorro",
+        "Inversión",
+    ]
+    if category not in categories_allowed:
+        raise ValueError(f"Categoría no permitida: {category}")
+    return category

AI/__init__.py

@@ -0,0 +1,2 @@
+# flake8: noqa
+from .AI import setup_api_key, classify_text, clean_category

docker-compose.yml

@@ -12,6 +12,8 @@ services:
       - "5432:5432"
     networks:
       - app_network
+    environment:
+      POSTGRES_HOST_AUTH_METHOD: trust
 
   web:
     build: