From cec49614c21d74e5ef19062e8c572d3051974aad Mon Sep 17 00:00:00 2001
From: Ori Yonay <oriyonay12@gmail.com>
Date: Sat, 21 Jan 2023 15:37:42 -0600
Subject: [PATCH] added message size limit of 2000 characters (issue 822)
 (#880)

* added message size limit (issue 822)
---
 backend/oasst_backend/api/v1/admin.py                   | 1 +
 backend/oasst_backend/config.py                         | 1 +
 backend/oasst_backend/tree_manager.py                   | 7 +++++++
 oasst-shared/oasst_shared/exceptions/oasst_api_error.py | 1 +
 4 files changed, 10 insertions(+)

diff --git a/backend/oasst_backend/api/v1/admin.py b/backend/oasst_backend/api/v1/admin.py
index 59c5efda..584c5fb7 100644
--- a/backend/oasst_backend/api/v1/admin.py
+++ b/backend/oasst_backend/api/v1/admin.py
@@ -54,6 +54,7 @@ class PublicSettings(pydantic.BaseModel):
 
     PROJECT_NAME: str
     API_V1_STR: str
+    MESSAGE_SIZE_LIMIT: int
     DEBUG_USE_SEED_DATA: bool
     DEBUG_ALLOW_SELF_LABELING: bool
     DEBUG_SKIP_EMBEDDING_COMPUTATION: bool
diff --git a/backend/oasst_backend/config.py b/backend/oasst_backend/config.py
index 99b10cb4..8ca2a413 100644
--- a/backend/oasst_backend/config.py
+++ b/backend/oasst_backend/config.py
@@ -72,6 +72,7 @@ class Settings(BaseSettings):
     DATABASE_MAX_TX_RETRY_COUNT: int = 3
 
     RATE_LIMIT: bool = True
+    MESSAGE_SIZE_LIMIT: int = 2000
     REDIS_HOST: str = "localhost"
     REDIS_PORT: str = "6379"
 
diff --git a/backend/oasst_backend/tree_manager.py b/backend/oasst_backend/tree_manager.py
index 026bb564..1224af4a 100644
--- a/backend/oasst_backend/tree_manager.py
+++ b/backend/oasst_backend/tree_manager.py
@@ -465,6 +465,13 @@ class TreeManager:
                     f"Frontend reports text reply to {interaction.message_id=} with {interaction.text=} by {interaction.user=}."
                 )
 
+                # ensure message size is below the predefined limit
+                if len(interaction.text) > settings.MESSAGE_SIZE_LIMIT:
+                    logger.error(
+                        f"Message size {len(interaction.text)=} exceeds size limit of {settings.MESSAGE_SIZE_LIMIT=}."
+                    )
+                    raise OasstError("Message size too long.", OasstErrorCode.TASK_MESSAGE_TOO_LONG)
+
                 # here we store the text reply in the database
                 message = pr.store_text_reply(
                     text=interaction.text,
diff --git a/oasst-shared/oasst_shared/exceptions/oasst_api_error.py b/oasst-shared/oasst_shared/exceptions/oasst_api_error.py
index bed6f942..0a548ebb 100644
--- a/oasst-shared/oasst_shared/exceptions/oasst_api_error.py
+++ b/oasst-shared/oasst_shared/exceptions/oasst_api_error.py
@@ -37,6 +37,7 @@ class OasstErrorCode(IntEnum):
     TASK_GENERATION_FAILED = 1005
     TASK_REQUESTED_TYPE_NOT_AVAILABLE = 1006
     TASK_AVAILABILITY_QUERY_FAILED = 1007
+    TASK_MESSAGE_TOO_LONG = 1008
 
     # 2000-3000: prompt_repository
     INVALID_FRONTEND_MESSAGE_ID = 2000