From f2f713023d4c12658d3e0662c4db560883450cd5 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Wed, 4 Sep 2024 19:57:41 +0200
Subject: [PATCH] refac

---
 backend/open_webui/apps/webui/routers/functions.py | 4 ++--
 backend/open_webui/apps/webui/routers/tools.py     | 4 ++--
 backend/open_webui/apps/webui/utils.py             | 9 ++++-----
 3 files changed, 8 insertions(+), 9 deletions(-)

diff --git a/backend/open_webui/apps/webui/routers/functions.py b/backend/open_webui/apps/webui/routers/functions.py
index 26373ff93..130603462 100644
--- a/backend/open_webui/apps/webui/routers/functions.py
+++ b/backend/open_webui/apps/webui/routers/functions.py
@@ -56,11 +56,11 @@ async def create_new_function(
     function = Functions.get_function_by_id(form_data.id)
     if function is None:
         try:
+            form_data.content = replace_imports(form_data.content)
             function_module, function_type, frontmatter = load_function_module_by_id(
                 form_data.id,
                 content=form_data.content,
             )
-            form_data.content = replace_imports(form_data.content)
             form_data.meta.manifest = frontmatter
 
             FUNCTIONS = request.app.state.FUNCTIONS
@@ -173,10 +173,10 @@ async def update_function_by_id(
     request: Request, id: str, form_data: FunctionForm, user=Depends(get_admin_user)
 ):
     try:
+        form_data.content = replace_imports(form_data.content)
         function_module, function_type, frontmatter = load_function_module_by_id(
             id, content=form_data.content
         )
-        form_data.content = replace_imports(form_data.content)
         form_data.meta.manifest = frontmatter
 
         FUNCTIONS = request.app.state.FUNCTIONS
diff --git a/backend/open_webui/apps/webui/routers/tools.py b/backend/open_webui/apps/webui/routers/tools.py
index 2354a0429..0db21c895 100644
--- a/backend/open_webui/apps/webui/routers/tools.py
+++ b/backend/open_webui/apps/webui/routers/tools.py
@@ -60,10 +60,10 @@ async def create_new_toolkit(
     toolkit = Tools.get_tool_by_id(form_data.id)
     if toolkit is None:
         try:
+            form_data.content = replace_imports(form_data.content)
             toolkit_module, frontmatter = load_toolkit_module_by_id(
                 form_data.id, content=form_data.content
             )
-            form_data.content = replace_imports(form_data.content)
             form_data.meta.manifest = frontmatter
 
             TOOLS = request.app.state.TOOLS
@@ -126,10 +126,10 @@ async def update_toolkit_by_id(
     user=Depends(get_admin_user),
 ):
     try:
+        form_data.content = replace_imports(form_data.content)
         toolkit_module, frontmatter = load_toolkit_module_by_id(
             id, content=form_data.content
         )
-        form_data.content = replace_imports(form_data.content)
         form_data.meta.manifest = frontmatter
 
         TOOLS = request.app.state.TOOLS
diff --git a/backend/open_webui/apps/webui/utils.py b/backend/open_webui/apps/webui/utils.py
index 1f876a773..e63fc8b26 100644
--- a/backend/open_webui/apps/webui/utils.py
+++ b/backend/open_webui/apps/webui/utils.py
@@ -76,8 +76,8 @@ def load_toolkit_module_by_id(toolkit_id, content=None):
 
         content = tool.content
 
-    content = replace_imports(content)
-    Tools.update_tool_by_id(toolkit_id, {"content": content})
+        content = replace_imports(content)
+        Tools.update_tool_by_id(toolkit_id, {"content": content})
 
     module_name = f"{toolkit_id}"
     module = types.ModuleType(module_name)
@@ -114,9 +114,8 @@ def load_function_module_by_id(function_id, content=None):
             raise Exception(f"Function not found: {function_id}")
         content = function.content
 
-    # Replace the module paths in the function content
-    content = replace_imports(content)
-    Functions.update_function_by_id(function_id, {"content": content})
+        content = replace_imports(content)
+        Functions.update_function_by_id(function_id, {"content": content})
 
     module_name = f"{function_id}"
     module = types.ModuleType(module_name)