From 25f859a31e68f55eca68043b854f9f3a59438991 Mon Sep 17 00:00:00 2001
From: Mamadou DICKO <63923024+mamadoudicko@users.noreply.github.com>
Date: Mon, 24 Jul 2023 14:17:33 +0200
Subject: [PATCH] fix: make max_brain_size overwritable (#751)

---
 backend/core/models/brains.py        | 6 +-----
 backend/core/models/settings.py      | 4 ++--
 backend/core/routes/upload_routes.py | 5 ++---
 3 files changed, 5 insertions(+), 10 deletions(-)

diff --git a/backend/core/models/brains.py b/backend/core/models/brains.py
index 56f22d75f..362cf7728 100644
--- a/backend/core/models/brains.py
+++ b/backend/core/models/brains.py
@@ -19,15 +19,11 @@ class Brain(BaseModel):
     temperature: Optional[float] = 0.0
     max_tokens: Optional[int] = 256
     files: List[Any] = []
+    max_brain_size = BrainRateLimiting().max_brain_size
 
     class Config:
         arbitrary_types_allowed = True
 
-    @property
-    def max_brain_size(self) -> int:
-        brain_rate_limiting = BrainRateLimiting()
-        return brain_rate_limiting.max_brain_size
-
     @property
     def commons(self) -> CommonsDep:
         return common_dependencies()
diff --git a/backend/core/models/settings.py b/backend/core/models/settings.py
index 448b2c996..9a2f9edd7 100644
--- a/backend/core/models/settings.py
+++ b/backend/core/models/settings.py
@@ -8,8 +8,8 @@ from vectorstore.supabase import SupabaseVectorStore
 
 
 class BrainRateLimiting(BaseSettings):
-    max_brain_size = 52428800
-    max_brain_per_user = 5
+    max_brain_size: int = 52428800
+    max_brain_per_user: int = 5
 
 
 class BrainSettings(BaseSettings):
diff --git a/backend/core/routes/upload_routes.py b/backend/core/routes/upload_routes.py
index 94ac2f35e..d423300fc 100644
--- a/backend/core/routes/upload_routes.py
+++ b/backend/core/routes/upload_routes.py
@@ -46,9 +46,8 @@ async def upload_file(
     commons = common_dependencies()
 
     if request.headers.get("Openai-Api-Key"):
-        brain.max_brain_size = os.getenv(
-            "MAX_BRAIN_SIZE_WITH_KEY", 209715200
-        )  # pyright: ignore reportPrivateUsage=none
+        brain.max_brain_size = int(os.getenv("MAX_BRAIN_SIZE_WITH_KEY", 209715200))
+
     remaining_free_space = brain.remaining_brain_size
 
     file_size = get_file_size(uploadFile)