Browse Source

refac: function cache

Timothy Jaeryang Baek 1 month ago
parent
commit
b944acd3ff
2 changed files with 27 additions and 3 deletions
  1. 4 0
      backend/open_webui/main.py
  2. 23 3
      backend/open_webui/utils/plugin.py

+ 4 - 0
backend/open_webui/main.py

@@ -637,8 +637,12 @@ app.state.WEBUI_AUTH_SIGNOUT_REDIRECT_URL = WEBUI_AUTH_SIGNOUT_REDIRECT_URL
 app.state.EXTERNAL_PWA_MANIFEST_URL = EXTERNAL_PWA_MANIFEST_URL
 app.state.EXTERNAL_PWA_MANIFEST_URL = EXTERNAL_PWA_MANIFEST_URL
 
 
 app.state.USER_COUNT = None
 app.state.USER_COUNT = None
+
 app.state.TOOLS = {}
 app.state.TOOLS = {}
+app.state.TOOL_CONTENTS = {}
+
 app.state.FUNCTIONS = {}
 app.state.FUNCTIONS = {}
+app.state.FUNCTION_CONTENTS = {}
 
 
 ########################################
 ########################################
 #
 #

+ 23 - 3
backend/open_webui/utils/plugin.py

@@ -115,7 +115,7 @@ def load_tool_module_by_id(tool_id, content=None):
         os.unlink(temp_file.name)
         os.unlink(temp_file.name)
 
 
 
 
-def load_function_module_by_id(function_id, content=None):
+def load_function_module_by_id(function_id: str, content: str | None = None):
     if content is None:
     if content is None:
         function = Functions.get_function_by_id(function_id)
         function = Functions.get_function_by_id(function_id)
         if not function:
         if not function:
@@ -167,20 +167,40 @@ def load_function_module_by_id(function_id, content=None):
 
 
 
 
 def get_function_module_from_cache(request, function_id):
 def get_function_module_from_cache(request, function_id):
+    function = Functions.get_function_by_id(function_id)
+    if not function:
+        raise Exception(f"Function not found: {function_id}")
+    content = function.content
+
+    new_content = replace_imports(content)
+    if new_content != content:
+        content = new_content
+        # Update the function content in the database
+        Functions.update_function_by_id(function_id, {"content": content})
+
     if (
     if (
+        hasattr(request.app.state, "FUNCTION_CONTENTS")
+        and function_id in request.app.state.FUNCTION_CONTENTS
+    ) and (
         hasattr(request.app.state, "FUNCTIONS")
         hasattr(request.app.state, "FUNCTIONS")
         and function_id in request.app.state.FUNCTIONS
         and function_id in request.app.state.FUNCTIONS
     ):
     ):
-        return request.app.state.FUNCTIONS[function_id], None, None
+        if request.app.state.FUNCTION_CONTENTS[function_id] == content:
+            return request.app.state.FUNCTIONS[function_id], None, None
 
 
     function_module, function_type, frontmatter = load_function_module_by_id(
     function_module, function_type, frontmatter = load_function_module_by_id(
-        function_id
+        function_id, content
     )
     )
 
 
     if not hasattr(request.app.state, "FUNCTIONS"):
     if not hasattr(request.app.state, "FUNCTIONS"):
         request.app.state.FUNCTIONS = {}
         request.app.state.FUNCTIONS = {}
 
 
+    if not hasattr(request.app.state, "FUNCTION_CONTENTS"):
+        request.app.state.FUNCTION_CONTENTS = {}
+
     request.app.state.FUNCTIONS[function_id] = function_module
     request.app.state.FUNCTIONS[function_id] = function_module
+    request.app.state.FUNCTION_CONTENTS[function_id] = content
+
     return function_module, function_type, frontmatter
     return function_module, function_type, frontmatter