Update app.py

try generalizing again

app.py

@@ -40,50 +40,100 @@ def my_custom_tool(arg1:str, arg2:int)-> str: #it's import to specify the return
 
 
 def load_image_sources():
-    with open("image_sources.json", "r") as file:
+    with open("expanded.json", "r") as file:
         return json.load(file)
 
 image_sources = load_image_sources()
 
+
 def get_image_url(image_type: str):
     """Finds the best match for the given image type using fuzzy matching."""
     choices = list(image_sources.keys())  # Get all available keys
     best_match, score, *rest = process.extractOne(image_type, choices)  
-    
+
     if score > 90:  # Set a threshold to ensure a reasonable match
-        print(best_match)
-        return image_sources[best_match]
+        print(f"Best match found: {best_match}")
+        return {
+            "url": image_sources[best_match]["url"],
+            "width": image_sources[best_match]["width"],
+            "height": image_sources[best_match]["height"]
+        }
     else:
-       #return None  # No good match found
-        return "https://editor.p5js.org/kfahn/full/2XD5Y8MiV"
+        # Return default values if no good match is found
+        return {
+            "url": "https://editor.p5js.org/kfahn/full/2XD5Y8MiV",
+            "width": 800,
+            "height": 800
+        }
+
+
+# def get_image_url(image_type: str):
+#     """Finds the best match for the given image type using fuzzy matching."""
+#     choices = list(image_sources.keys())  # Get all available keys
+#     best_match, score, *rest = process.extractOne(image_type, choices)  
+    
+#     if score > 90:  # Set a threshold to ensure a reasonable match
+#         print(best_match)
+#         return image_sources[best_match]
+#     else:
+#        #return None  # No good match found
+#         return "https://editor.p5js.org/kfahn/full/2XD5Y8MiV"
+
+# async def capture_screenshot(image_type: str):
+#     """Launches Playwright and uses user input,  if any, to captures a screenshot of an image from p5.js."""
+#     print("Launching Playwright...")
+#     async with async_playwright() as p:
+#         browser = await p.chromium.launch(headless=True)
+#         page = await browser.new_page()
+
+#         #url = "https://openprocessing.org/sketch/2539973"
+#         url = "https://editor.p5js.org/kfahn/full/2XD5Y8MiV"
+       
+#         if image_type: 
+#             image_url = get_image_url(image_type)
+#         else:
+#             image_url = url
+        
+#         print(f"Opening image from p5 sketch: {image_url}")
+#         await page.goto(image_url, timeout=120000)  # Wait for the image page to load
+
+#         print("Waiting for image element...")
+#        # await page.wait_for_selector("img", timeout=120000)  # Wait for the <img> to be visible
+#         await page.wait_for_timeout(5000)  # Allow sketch to fully render
+#         print("Capturing screenshot...")
+#         await page.set_viewport_size({"width": 800, "height": 800})
+#         await page.locator("iframe").screenshot(path="img.png")
+#         await browser.close()
+#         print("Screenshot saved!")
 
 async def capture_screenshot(image_type: str):
-    """Launches Playwright and uses user input,  if any, to captures a screenshot of an image from p5.js."""
+    """Launches Playwright and captures a screenshot of an image from p5.js."""
     print("Launching Playwright...")
     async with async_playwright() as p:
         browser = await p.chromium.launch(headless=True)
         page = await browser.new_page()
 
-        #url = "https://openprocessing.org/sketch/2539973"
-        url = "https://editor.p5js.org/kfahn/full/2XD5Y8MiV"
-       
-        if image_type: 
-            image_url = get_image_url(image_type)
-        else:
-            image_url = url
-        
+        image_data = get_image_url(image_type)
+        image_url = image_data["url"]
+        width = image_data["width"]
+        height = image_data["height"]
+
         print(f"Opening image from p5 sketch: {image_url}")
         await page.goto(image_url, timeout=120000)  # Wait for the image page to load
 
-        print("Waiting for image element...")
-       # await page.wait_for_selector("img", timeout=120000)  # Wait for the <img> to be visible
+        print("Waiting for sketch to render...")
         await page.wait_for_timeout(5000)  # Allow sketch to fully render
-        print("Capturing screenshot...")
-        await page.set_viewport_size({"width": 800, "height": 800})
+
+        print(f"Setting viewport to {width}x{height}...")
+        await page.set_viewport_size({"width": width, "height": height})
+
+        print("Capturing screenshot of canvas...")
         await page.locator("iframe").screenshot(path="img.png")
+
         await browser.close()
         print("Screenshot saved!")
-   
+
+
      
 @tool
 def grab_image(image_type: str) -> Image: