浏览代码

refac/enh: ability to export/sync function valves

Timothy Jaeryang Baek 1 月之前
父节点
当前提交
d0f338bb99
共有 2 个文件被更改,包括 36 次插入12 次删除
  1. 30 7
      backend/open_webui/models/functions.py
  2. 6 5
      backend/open_webui/routers/functions.py

+ 30 - 7
backend/open_webui/models/functions.py

@@ -54,6 +54,22 @@ class FunctionModel(BaseModel):
     model_config = ConfigDict(from_attributes=True)
     model_config = ConfigDict(from_attributes=True)
 
 
 
 
+class FunctionWithValvesModel(BaseModel):
+    id: str
+    user_id: str
+    name: str
+    type: str
+    content: str
+    meta: FunctionMeta
+    valves: Optional[dict] = None
+    is_active: bool = False
+    is_global: bool = False
+    updated_at: int  # timestamp in epoch
+    created_at: int  # timestamp in epoch
+
+    model_config = ConfigDict(from_attributes=True)
+
+
 ####################
 ####################
 # Forms
 # Forms
 ####################
 ####################
@@ -111,8 +127,8 @@ class FunctionsTable:
             return None
             return None
 
 
     def sync_functions(
     def sync_functions(
-        self, user_id: str, functions: list[FunctionModel]
-    ) -> list[FunctionModel]:
+        self, user_id: str, functions: list[FunctionWithValvesModel]
+    ) -> list[FunctionWithValvesModel]:
         # Synchronize functions for a user by updating existing ones, inserting new ones, and removing those that are no longer present.
         # Synchronize functions for a user by updating existing ones, inserting new ones, and removing those that are no longer present.
         try:
         try:
             with get_db() as db:
             with get_db() as db:
@@ -166,17 +182,24 @@ class FunctionsTable:
         except Exception:
         except Exception:
             return None
             return None
 
 
-    def get_functions(self, active_only=False) -> list[FunctionModel]:
+    def get_functions(
+        self, active_only=False, include_valves=False
+    ) -> list[FunctionModel | FunctionWithValvesModel]:
         with get_db() as db:
         with get_db() as db:
             if active_only:
             if active_only:
+                functions = db.query(Function).filter_by(is_active=True).all()
+
+            else:
+                functions = db.query(Function).all()
+
+            if include_valves:
                 return [
                 return [
-                    FunctionModel.model_validate(function)
-                    for function in db.query(Function).filter_by(is_active=True).all()
+                    FunctionWithValvesModel.model_validate(function)
+                    for function in functions
                 ]
                 ]
             else:
             else:
                 return [
                 return [
-                    FunctionModel.model_validate(function)
-                    for function in db.query(Function).all()
+                    FunctionModel.model_validate(function) for function in functions
                 ]
                 ]
 
 
     def get_functions_by_type(
     def get_functions_by_type(

+ 6 - 5
backend/open_webui/routers/functions.py

@@ -10,6 +10,7 @@ from open_webui.models.functions import (
     FunctionForm,
     FunctionForm,
     FunctionModel,
     FunctionModel,
     FunctionResponse,
     FunctionResponse,
+    FunctionWithValvesModel,
     Functions,
     Functions,
 )
 )
 from open_webui.utils.plugin import (
 from open_webui.utils.plugin import (
@@ -46,9 +47,9 @@ async def get_functions(user=Depends(get_verified_user)):
 ############################
 ############################
 
 
 
 
-@router.get("/export", response_model=list[FunctionModel])
-async def get_functions(user=Depends(get_admin_user)):
-    return Functions.get_functions()
+@router.get("/export", response_model=list[FunctionModel | FunctionWithValvesModel])
+async def get_functions(include_valves: bool = False, user=Depends(get_admin_user)):
+    return Functions.get_functions(include_valves=include_valves)
 
 
 
 
 ############################
 ############################
@@ -132,10 +133,10 @@ async def load_function_from_url(
 
 
 
 
 class SyncFunctionsForm(BaseModel):
 class SyncFunctionsForm(BaseModel):
-    functions: list[FunctionModel] = []
+    functions: list[FunctionWithValvesModel] = []
 
 
 
 
-@router.post("/sync", response_model=list[FunctionModel])
+@router.post("/sync", response_model=list[FunctionWithValvesModel])
 async def sync_functions(
 async def sync_functions(
     request: Request, form_data: SyncFunctionsForm, user=Depends(get_admin_user)
     request: Request, form_data: SyncFunctionsForm, user=Depends(get_admin_user)
 ):
 ):