From ff1cd306d8898e8a50f29adff3154df2bcd5ede6 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Mon, 10 Jun 2024 22:38:48 -0700
Subject: [PATCH] refac

---
 backend/apps/webui/main.py          |  3 ++-
 backend/apps/webui/routers/tools.py | 21 +++++++++++++++------
 2 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/backend/apps/webui/main.py b/backend/apps/webui/main.py
index 339fe8a83..62a0a7a7b 100644
--- a/backend/apps/webui/main.py
+++ b/backend/apps/webui/main.py
@@ -39,6 +39,7 @@ app.state.config = AppConfig()
 
 app.state.config.ENABLE_SIGNUP = ENABLE_SIGNUP
 app.state.config.JWT_EXPIRES_IN = JWT_EXPIRES_IN
+app.state.AUTH_TRUSTED_EMAIL_HEADER = WEBUI_AUTH_TRUSTED_EMAIL_HEADER
 
 
 app.state.config.SHOW_ADMIN_DETAILS = SHOW_ADMIN_DETAILS
@@ -55,7 +56,7 @@ app.state.config.BANNERS = WEBUI_BANNERS
 app.state.config.ENABLE_COMMUNITY_SHARING = ENABLE_COMMUNITY_SHARING
 
 app.state.MODELS = {}
-app.state.AUTH_TRUSTED_EMAIL_HEADER = WEBUI_AUTH_TRUSTED_EMAIL_HEADER
+app.state.TOOLS = {}
 
 
 app.add_middleware(
diff --git a/backend/apps/webui/routers/tools.py b/backend/apps/webui/routers/tools.py
index 67e391ddd..60e8319e6 100644
--- a/backend/apps/webui/routers/tools.py
+++ b/backend/apps/webui/routers/tools.py
@@ -1,4 +1,4 @@
-from fastapi import Depends, FastAPI, HTTPException, status
+from fastapi import Depends, FastAPI, HTTPException, status, Request
 from datetime import datetime, timedelta
 from typing import List, Union, Optional
 
@@ -20,8 +20,6 @@ from config import DATA_DIR
 TOOLS_DIR = f"{DATA_DIR}/tools"
 os.makedirs(TOOLS_DIR, exist_ok=True)
 
-TOOLS = {}
-
 
 router = APIRouter()
 
@@ -73,7 +71,9 @@ async def get_toolkits(user=Depends(get_admin_user)):
 
 
 @router.post("/create", response_model=Optional[ToolResponse])
-async def create_new_toolkit(form_data: ToolForm, user=Depends(get_admin_user)):
+async def create_new_toolkit(
+    request: Request, form_data: ToolForm, user=Depends(get_admin_user)
+):
     if not form_data.id.isidentifier():
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
@@ -90,6 +90,8 @@ async def create_new_toolkit(form_data: ToolForm, user=Depends(get_admin_user)):
                 tool_file.write(form_data.content)
 
             toolkit_module = load_toolkit_module_from_path(form_data.id, toolkit_path)
+
+            TOOLS = request.app.state.TOOLS
             TOOLS[form_data.id] = toolkit_module
 
             specs = get_tools_specs(TOOLS[form_data.id])
@@ -139,7 +141,7 @@ async def get_toolkit_by_id(id: str, user=Depends(get_admin_user)):
 
 @router.post("/id/{id}/update", response_model=Optional[ToolModel])
 async def update_toolkit_by_id(
-    id: str, form_data: ToolForm, user=Depends(get_admin_user)
+    request: Request, id: str, form_data: ToolForm, user=Depends(get_admin_user)
 ):
     toolkit_path = os.path.join(TOOLS_DIR, f"{id}.py")
 
@@ -148,6 +150,8 @@ async def update_toolkit_by_id(
             tool_file.write(form_data.content)
 
         toolkit_module = load_toolkit_module_from_path(id, toolkit_path)
+
+        TOOLS = request.app.state.TOOLS
         TOOLS[id] = toolkit_module
 
         specs = get_tools_specs(TOOLS[id])
@@ -181,6 +185,11 @@ async def update_toolkit_by_id(
 
 
 @router.delete("/id/{id}/delete", response_model=bool)
-async def delete_toolkit_by_id(id: str, user=Depends(get_admin_user)):
+async def delete_toolkit_by_id(request: Request, id: str, user=Depends(get_admin_user)):
     result = Tools.delete_tool_by_id(id)
+
+    if result:
+        TOOLS = request.app.state.TOOLS
+        del TOOLS[id]
+
     return result