refac

backend/open_webui/tools/builtin.py

@@ -3255,10 +3255,22 @@ async def update_task(
 # =============================================================================
 
 
+async def _validate_owned_automation_folder(user_id: str, folder_id: Optional[str]) -> Optional[str]:
+    if not folder_id:
+        return None
+    from open_webui.models.folders import Folders
+
+    folder = await Folders.get_folder_by_id_and_user_id(folder_id, user_id)
+    if not folder:
+        raise ValueError('Folder not found')
+    return folder.id
+
+
 async def create_automation(
     name: str,
     prompt: str,
     rrule: str,
+    folder_id: Optional[str] = None,
     __request__: Request = None,
     __user__: dict = None,
     __metadata__: dict = None,
@@ -3281,6 +3293,7 @@ async def create_automation(
     :param name: A short descriptive name for the automation
     :param prompt: The prompt/instructions to execute on each run
     :param rrule: An iCalendar RRULE string defining the schedule
+    :param folder_id: Optional owner-owned folder ID for generated chats
     :return: JSON with the created automation details including id, next scheduled runs
     """
     if __request__ is None:
@@ -3308,6 +3321,11 @@ async def create_automation(
         if not model_id:
             return json.dumps({'error': 'Could not detect current model'})
 
+        try:
+            folder_id = await _validate_owned_automation_folder(user_id, folder_id)
+        except ValueError as e:
+            return json.dumps({'error': str(e)})
+
         # Validate the RRULE
         try:
             validate_rrule(rrule, tz=user.timezone)
@@ -3322,6 +3340,7 @@ async def create_automation(
         tz = user.timezone
         form = AutomationForm(
             name=name,
+            folder_id=folder_id,
             data=AutomationData(
                 prompt=prompt,
                 model_id=model_id,
@@ -3337,6 +3356,7 @@ async def create_automation(
                 'status': 'success',
                 'id': automation.id,
                 'name': automation.name,
+                'folder_id': automation.folder_id,
                 'model_id': model_id,
                 'is_active': automation.is_active,
                 'next_runs': next_n_runs_ns(rrule, tz=tz),
@@ -3354,6 +3374,7 @@ async def update_automation(
     prompt: Optional[str] = None,
     rrule: Optional[str] = None,
     model_id: Optional[str] = None,
+    folder_id: Optional[str] = None,
     __request__: Request = None,
     __user__: dict = None,
 ) -> str:
@@ -3365,6 +3386,7 @@ async def update_automation(
     :param prompt: New prompt/instructions (optional)
     :param rrule: New iCalendar RRULE schedule string (optional). See create_automation for format examples.
     :param model_id: New model ID to use (optional)
+    :param folder_id: New owner-owned folder ID (optional); pass an empty string to clear
     :return: JSON with the updated automation details
     """
     if __request__ is None:
@@ -3395,6 +3417,13 @@ async def update_automation(
         new_prompt = prompt if prompt is not None else automation.data.get('prompt', '')
         new_model_id = model_id if model_id is not None else automation.data.get('model_id', '')
         new_rrule = rrule if rrule is not None else automation.data.get('rrule', '')
+        if folder_id is None:
+            new_folder_id = automation.folder_id
+        else:
+            try:
+                new_folder_id = await _validate_owned_automation_folder(user_id, folder_id)
+            except ValueError as e:
+                return json.dumps({'error': str(e)})
 
         # Validate RRULE if changed
         if rrule is not None:
@@ -3411,6 +3440,7 @@ async def update_automation(
         tz = user.timezone
         form = AutomationForm(
             name=new_name,
+            folder_id=new_folder_id,
             data=AutomationData(
                 prompt=new_prompt,
                 model_id=new_model_id,
@@ -3426,6 +3456,7 @@ async def update_automation(
                 'status': 'success',
                 'id': updated.id,
                 'name': updated.name,
+                'folder_id': updated.folder_id,
                 'model_id': new_model_id,
                 'is_active': updated.is_active,
                 'next_runs': next_n_runs_ns(new_rrule, tz=tz),
@@ -3439,6 +3470,7 @@ async def update_automation(
 
 async def list_automations(
     status: Optional[str] = None,
+    folder_id: Optional[str] = None,
     count: int = 10,
     __request__: Request = None,
     __user__: dict = None,
@@ -3447,6 +3479,7 @@ async def list_automations(
     List the user's scheduled automations.
 
     :param status: Filter by status: "active", "paused", or omit for all
+    :param folder_id: Optional owner-owned folder ID filter; pass an empty string to clear the folder filter
     :param count: Maximum number of automations to return (default: 10)
     :return: JSON list of automations with id, name, prompt snippet, schedule, status, and next runs
     """
@@ -3463,10 +3496,16 @@ async def list_automations(
 
         user_id = __user__.get('id')
         user = await Users.get_user_by_id(user_id)
+        if folder_id:
+            try:
+                folder_id = await _validate_owned_automation_folder(user_id, folder_id)
+            except ValueError as e:
+                return json.dumps({'error': str(e)})
 
         result = await Automations.search_automations(
             user_id=user_id,
             status=status,
+            folder_id=folder_id,
             skip=0,
             limit=count,
         )
@@ -3481,6 +3520,7 @@ async def list_automations(
                 {
                     'id': item.id,
                     'name': item.name,
+                    'folder_id': item.folder_id,
                     'prompt_snippet': snippet,
                     'model_id': item.data.get('model_id', ''),
                     'rrule': rrule,