From d7ed237f781486700e25d2ea2d782bfb870d067c Mon Sep 17 00:00:00 2001
From: pablohashescobar <118773738+pablohashescobar@users.noreply.github.com>
Date: Thu, 30 Mar 2023 01:23:56 +0530
Subject: [PATCH] feat: add API limit for AI assistance (#592)

---
 apiserver/plane/api/views/gpt.py       | 26 ++++++++++++++++++++++++++
 apiserver/plane/settings/local.py      |  4 ++++
 apiserver/plane/settings/production.py |  4 ++++
 apiserver/plane/settings/staging.py    |  4 ++++
 4 files changed, 38 insertions(+)

diff --git a/apiserver/plane/api/views/gpt.py b/apiserver/plane/api/views/gpt.py
index bdc060e4f..8a138fb89 100644
--- a/apiserver/plane/api/views/gpt.py
+++ b/apiserver/plane/api/views/gpt.py
@@ -1,3 +1,6 @@
+# Python imports
+import requests
+
 # Third party imports
 from rest_framework.response import Response
 from rest_framework import status
@@ -25,6 +28,28 @@ class GPTIntegrationEndpoint(BaseAPIView):
                     status=status.HTTP_400_BAD_REQUEST,
                 )
 
+            # If logger is enabled check for request limit
+            if settings.LOGGER_BASE_URL:
+                try:
+                    headers = {
+                        "Content-Type": "application/json",
+                    }
+
+                    response = requests.post(
+                        settings.LOGGER_BASE_URL,
+                        json={"user_id": str(request.user.id)},
+                        headers=headers,
+                    )
+                    if not response.json().get("success", False):
+                        return Response(
+                            {
+                                "error": "You have surpassed the monthly limit for AI assistance"
+                            },
+                            status=status.HTTP_429_TOO_MANY_REQUESTS,
+                        )
+                except Exception as e:
+                    capture_exception(e)
+
             prompt = request.data.get("prompt", False)
             task = request.data.get("task", False)
 
@@ -50,6 +75,7 @@ class GPTIntegrationEndpoint(BaseAPIView):
                 status=status.HTTP_200_OK,
             )
         except Exception as e:
+            print(e)
             capture_exception(e)
             return Response(
                 {"error": "Something went wrong please try again later"},
diff --git a/apiserver/plane/settings/local.py b/apiserver/plane/settings/local.py
index 21d4ec0ea..bf161568b 100644
--- a/apiserver/plane/settings/local.py
+++ b/apiserver/plane/settings/local.py
@@ -81,6 +81,10 @@ PROXY_BASE_URL = os.environ.get("PROXY_BASE_URL", False)
 
 ANALYTICS_SECRET_KEY = os.environ.get("ANALYTICS_SECRET_KEY", False)
 ANALYTICS_BASE_API = os.environ.get("ANALYTICS_BASE_API", False)
+
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", False)
 GPT_ENGINE = os.environ.get("GPT_ENGINE", "text-davinci-003")
+
 SLACK_BOT_TOKEN = os.environ.get("SLACK_BOT_TOKEN", False)
+
+LOGGER_BASE_URL = os.environ.get("LOGGER_BASE_URL", False)
diff --git a/apiserver/plane/settings/production.py b/apiserver/plane/settings/production.py
index ba8268b75..5569e1c09 100644
--- a/apiserver/plane/settings/production.py
+++ b/apiserver/plane/settings/production.py
@@ -229,6 +229,10 @@ PROXY_BASE_URL = os.environ.get("PROXY_BASE_URL", False)
 
 ANALYTICS_SECRET_KEY = os.environ.get("ANALYTICS_SECRET_KEY", False)
 ANALYTICS_BASE_API = os.environ.get("ANALYTICS_BASE_API", False)
+
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", False)
 GPT_ENGINE = os.environ.get("GPT_ENGINE", "text-davinci-003")
+
 SLACK_BOT_TOKEN = os.environ.get("SLACK_BOT_TOKEN", False)
+
+LOGGER_BASE_URL = os.environ.get("LOGGER_BASE_URL", False)
diff --git a/apiserver/plane/settings/staging.py b/apiserver/plane/settings/staging.py
index 0881d847f..9015ce03f 100644
--- a/apiserver/plane/settings/staging.py
+++ b/apiserver/plane/settings/staging.py
@@ -190,6 +190,10 @@ PROXY_BASE_URL = os.environ.get("PROXY_BASE_URL", False)
 
 ANALYTICS_SECRET_KEY = os.environ.get("ANALYTICS_SECRET_KEY", False)
 ANALYTICS_BASE_API = os.environ.get("ANALYTICS_BASE_API", False)
+
 OPENAI_API_KEY = os.environ.get("OPENAI_API_KEY", False)
 GPT_ENGINE = os.environ.get("GPT_ENGINE", "text-davinci-003")
+
 SLACK_BOT_TOKEN = os.environ.get("SLACK_BOT_TOKEN", False)
+
+LOGGER_BASE_URL = os.environ.get("LOGGER_BASE_URL", False)