Enhance DigitalOcean droplet function to support power on/off actions and update schemas accordingly

Author: jkpe

genai-examples/function-instructions.md

@@ -1,3 +1,3 @@
 ref: https://docs.digitalocean.com/products/genai-platform/concepts/function-instructions/
 
-Call this function when the user asks about their DigitalOcean droplets, virtual machines, instances, or servers. Use this function to retrieve information about one or more droplets in a DigitalOcean account.
+Call this function when the user asks about their DigitalOcean droplets, virtual machines, instances, or servers. Use this function to retrieve information about one or more droplets in a DigitalOcean account. You can also use this function to power on or power off a specific droplet when the user asks to start, turn on, boot, stop, turn off, or shut down a droplet.

genai-examples/input_schema.json

@@ -7,11 +7,16 @@
     "droplet_id": {
       "type": "string",
       "required": false,
-      "description": "Specific droplet ID to retrieve information for. If not provided, returns a list of droplets"
+      "description": "Specific droplet ID to retrieve information for or to perform an action on. If not provided, returns a list of droplets"
     },
     "tag": {
       "type": "string",
       "required": false,
       "description": "Filter droplets by tag"
+    },
+    "action": {
+      "type": "string",
+      "required": false,
+      "description": "Action to perform on a droplet. Supported values: 'power_on', 'power_off'. Requires droplet_id to be specified"
     }
   }

genai-examples/output_schema.json

@@ -10,10 +10,22 @@
     },
     "droplets": {
       "type": "string",
-      "description": "JSON string containing list of droplet information"
+      "description": "JSON string containing list of droplet information",
+      "required": false
     },
     "count": {
       "type": "number",
-      "description": "Total number of droplets returned"
+      "description": "Total number of droplets returned",
+      "required": false
+    },
+    "action": {
+      "type": "object",
+      "description": "Response from the droplet action API (when power_on or power_off is performed)",
+      "required": false
+    },
+    "message": {
+      "type": "string",
+      "description": "Status message for action operations",
+      "required": false
     }
   }

packages/do-api/api/__main__.py

@@ -4,23 +4,25 @@
 
 def main(args):
     """
-    Retrieve droplet information from the DigitalOcean API.
+    Retrieve or control droplet information from the DigitalOcean API.
     
     Args:
         args (dict): Input parameters
             - api_token: DigitalOcean API token
-            - droplet_id: Optional specific droplet ID to retrieve
+            - droplet_id: Optional specific droplet ID to retrieve or control
             - tag: Optional tag to filter droplets
             - limit: Optional limit for the number of droplets to return (default: 10)
+            - action: Optional action to perform on droplet (power_on, power_off)
     
     Returns:
-        dict: API response containing droplet information in the 'body' key
+        dict: API response containing droplet information or action status in the 'body' key
     """
     # Extract parameters
     api_token = os.environ.get('DO_API_TOKEN')  # Get token from environment variable
     droplet_id = args.get('droplet_id')
     tag = args.get('tag')
     limit = args.get('limit', 10)
+    action = args.get('action')
 
     # Validate environment variable
     if not api_token:
@@ -35,8 +37,41 @@ def main(args):
         # Initialize DigitalOcean client
         client = Client(token=api_token)
 
+        # Check if we need to perform an action on a droplet
+        if droplet_id and action:
+            droplet_id = int(droplet_id)  # Convert to integer as the API expects
+            
+            if action == 'power_on':
+                # Call the post method with the action type
+                response = client.droplet_actions.post(droplet_id=droplet_id, body={"type": "power_on"})
+                return {
+                    'body': {
+                        'action': response,
+                        'status': 'success',
+                        'message': f'Power on initiated for droplet {droplet_id}'
+                    }
+                }
+            elif action == 'power_off':
+                # Call the post method with the action type
+                response = client.droplet_actions.post(droplet_id=droplet_id, body={"type": "power_off"})
+                return {
+                    'body': {
+                        'action': response,
+                        'status': 'success',
+                        'message': f'Power off initiated for droplet {droplet_id}'
+                    }
+                }
+            else:
+                return {
+                    'body': {
+                        'error': f'Unsupported action: {action}. Supported actions: power_on, power_off',
+                        'status': 'error'
+                    }
+                }
+        
         # Get specific droplet or list of droplets
         if droplet_id:
+            droplet_id = int(droplet_id)  # Convert to integer as the API expects
             response = client.droplets.get(droplet_id=droplet_id)
             droplets = [response['droplet']] if 'droplet' in response else []
         else:
@@ -90,9 +125,30 @@ def main(args):
 
 # For local testing
 if __name__ == "__main__":
-    test_args = {
-        'api_token': 'your_test_token',
+    import sys
+    
+    # Example 1: List droplets (default behavior)
+    test_args_list = {
         'limit': 5
     }
+    
+    # Example 2: Power on a specific droplet
+    test_args_power_on = {
+        'droplet_id': 12345,  # Replace with actual droplet ID
+        'action': 'power_on'
+    }
+    
+    # Example 3: Power off a specific droplet
+    test_args_power_off = {
+        'droplet_id': 12345,  # Replace with actual droplet ID
+        'action': 'power_off'
+    }
+    
+    # Choose which example to run
+    # Uncomment the test_args you want to use
+    test_args = test_args_list
+    # test_args = test_args_power_on
+    # test_args = test_args_power_off
+    
     result = main(test_args)
     print(json.dumps(result, indent=2))