Kaynağa Gözat

Merge pull request #13358 from feddersen-group/fix/load_ext_dependencies_on_startup

fix: install external requirements of tools and functions to prevent deactivation
Tim Jaeryang Baek 3 ay önce
ebeveyn
işleme
e822984162
2 değiştirilmiş dosya ile 40 ekleme ve 2 silme
  1. 9 1
      backend/open_webui/main.py
  2. 31 1
      backend/open_webui/utils/plugin.py

+ 9 - 1
backend/open_webui/main.py

@@ -383,6 +383,7 @@ from open_webui.utils.auth import (
     get_admin_user,
     get_verified_user,
 )
+from open_webui.utils.plugin import install_admin_tool_and_function_dependencies
 from open_webui.utils.oauth import OAuthManager
 from open_webui.utils.security_headers import SecurityHeadersMiddleware
 
@@ -451,6 +452,12 @@ async def lifespan(app: FastAPI):
         limiter.total_tokens = pool_size
 
     asyncio.create_task(periodic_usage_pool_cleanup())
+
+    # This should be blocking (sync) so functions are not deactivated on first /get_models calls
+    # when the first user lands on the / route.
+    log.info("Installing external dependencies of functions and tools...")
+    install_admin_tool_and_function_dependencies()
+
     yield
 
 
@@ -895,7 +902,8 @@ class RedirectMiddleware(BaseHTTPMiddleware):
 
             # Check for the specific watch path and the presence of 'v' parameter
             if path.endswith("/watch") and "v" in query_params:
-                video_id = query_params["v"][0]  # Extract the first 'v' parameter
+                # Extract the first 'v' parameter
+                video_id = query_params["v"][0]
                 encoded_video_id = urlencode({"youtube": video_id})
                 redirect_url = f"/?{encoded_video_id}"
                 return RedirectResponse(url=redirect_url)

+ 31 - 1
backend/open_webui/utils/plugin.py

@@ -157,7 +157,8 @@ def load_function_module_by_id(function_id, content=None):
             raise Exception("No Function class found in the module")
     except Exception as e:
         log.error(f"Error loading module: {function_id}: {e}")
-        del sys.modules[module_name]  # Cleanup by removing the module in case of error
+        # Cleanup by removing the module in case of error
+        del sys.modules[module_name]
 
         Functions.update_function_by_id(function_id, {"is_active": False})
         raise e
@@ -182,3 +183,32 @@ def install_frontmatter_requirements(requirements: str):
 
     else:
         log.info("No requirements found in frontmatter.")
+
+
+def install_admin_tool_and_function_dependencies():
+    """
+    Install all dependencies for all admin tools and active functions.
+
+    By first collecting all dependencies from the frontmatter of each tool and function,
+    and then installing them using pip. Duplicates or similar version specifications are
+    handled by pip as much as possible.
+    """
+    function_list = Functions.get_functions(active_only=True)
+    tool_list = Tools.get_tools()
+
+    all_dependencies = ""
+    try:
+        for function in function_list:
+            frontmatter = extract_frontmatter(replace_imports(function.content))
+            if dependencies := frontmatter.get("requirements"):
+                all_dependencies += f"{dependencies}, "
+        for tool in tool_list:
+            # Only install requirements for admin tools
+            if tool.user.role == "admin":
+                frontmatter = extract_frontmatter(replace_imports(tool.content))
+                if dependencies := frontmatter.get("requirements"):
+                    all_dependencies += f"{dependencies}, "
+
+        install_frontmatter_requirements(all_dependencies.strip(", "))
+    except Exception as e:
+        log.error(f"Error installing requirements: {e}")