Merge pull request #397 from noriellecruz/main

fix: screenshot not showing properly if the website is slow

bot/plugins/webshot.py

@@ -1,3 +1,4 @@
+import os, requests, random, string
 from typing import Dict
 from .plugin import Plugin
 
@@ -11,21 +12,49 @@ def get_source_name(self) -> str:
     def get_spec(self) -> [Dict]:
         return [{
             "name": "screenshot_website",
-            "description": "Show screenshot/image of a website from a given url or domain name",
+            "description": "Show screenshot/image of a website from a given url or domain name.",
             "parameters": {
                 "type": "object",
                 "properties": {
-                    "url": {"type": "string", "description": "Website url or domain name"}
+                    "url": {"type": "string", "description": "Website url or domain name. Correctly formatted url is required. Example: https://www.google.com"}
                 },
                 "required": ["url"],
             },
         }]
+
+    def generate_random_string(self, length):
+        characters = string.ascii_letters + string.digits
+        return ''.join(random.choice(characters) for _ in range(length))
 
     async def execute(self, function_name, **kwargs) -> Dict:
-        return {
-            'direct_result': {
-                'kind': 'photo',
-                'format': 'url',
-                'value': f'https://image.thum.io/get/maxAge/12/width/720/{kwargs["url"]}'
-            }
-        }
+        try:
+            image_url = f'https://image.thum.io/get/maxAge/12/width/720/{kwargs["url"]}'
+
+            # preload url first
+            requests.get(image_url)
+
+            # download the actual image
+            response = requests.get(image_url, timeout=30)
+
+            if response.status_code == 200:
+                if not os.path.exists("uploads/webshot"):
+                    os.makedirs("uploads/webshot")
+
+                image_file_path = os.path.join("uploads/webshot", f"{self.generate_random_string(15)}.png")
+                with open(image_file_path, "wb") as f:
+                    f.write(response.content)
+
+                return {
+                    'direct_result': {
+                        'kind': 'photo',
+                        'format': 'path',
+                        'value': image_file_path
+                    }
+                }
+            else:
+                return {'result': 'Unable to screenshot website'}
+        except:
+            if 'image_file_path' in locals():
+                os.remove(image_file_path)
+
+            return {'result': 'Unable to screenshot website'}