Update app.py

app.py

@@ -6,54 +6,42 @@ import time
 
 app = Flask(__name__)
 
-# Global variables to store workspace and bot IDs
-GLOBAL_WORKSPACE_ID = None
-GLOBAL_BOT_ID = None
-
-
-# Authorization value used in requests (should be updated with a valid Authorization),
+# Authorization value used in requests (update with a valid Authorization)
 TOKEN = "Bearer bp_pat_vTuxol25N0ymBpYaWqtWpFfGPKt260IfT784"
+
 # -------------------------------------------------------------------
-# Helper functions for random bot/workspace names
+# Helper functions
 # -------------------------------------------------------------------
 def generate_random_name(length=5):
-    """Generate a random name for workspace or bot"""
     return ''.join(random.choices(string.ascii_letters, k=length))
 
 # -------------------------------------------------------------------
-# Functions to create/delete workspaces and bots
+# Workspace and Bot management
 # -------------------------------------------------------------------
 def create_workspace():
-    """Create a new workspace and return its ID"""
-    ws_url = "https://api.botpress.cloud/v1/admin/workspaces"
+    url = "https://api.botpress.cloud/v1/admin/workspaces"
     headers = {
         "User-Agent": "Mozilla/5.0",
         "Authorization": TOKEN
     }
     payload = {"name": generate_random_name()}
-
     try:
-        response = requests.post(ws_url, headers=headers, json=payload)
-        if response.status_code == 200:
-            response_json = response.json()
-            workspace_id = response_json.get('id')
-            print(f"Successfully created workspace: {workspace_id}")
-            return workspace_id
-        else:
-            print(f"Workspace creation failed with: {response.status_code}, {response.text}")
-            return None
+        r = requests.post(url, headers=headers, json=payload, timeout=60)
+        if r.status_code == 200:
+            ws_id = r.json().get("id")
+            if ws_id:
+                print(f"Workspace created: {ws_id}")
+                return ws_id
+        print(f"Workspace creation failed: {r.status_code} {r.text}")
     except Exception as e:
-        print(f"Error creating workspace: {str(e)}")
-        return None
-
+        print(f"Error creating workspace: {e}")
+    return None
 
 def create_bot(workspace_id):
-    """Create a new bot in the specified workspace and return its ID"""
     if not workspace_id:
-        print("Cannot create bot: No workspace ID provided")
+        print("No workspace_id provided for bot creation")
         return None
-
-    bot_url = "https://api.botpress.cloud/v1/admin/bots"
+    url = "https://api.botpress.cloud/v1/admin/bots"
     headers = {
         "User-Agent": "Mozilla/5.0",
         "x-workspace-id": workspace_id,
@@ -61,40 +49,28 @@ def create_bot(workspace_id):
         "Content-Type": "application/json"
     }
     payload = {"name": generate_random_name()}
-
     try:
-        response = requests.post(bot_url, headers=headers, json=payload)
-        if response.status_code == 200:
-            response_json = response.json()
-            bot_id = response_json.get("bot", {}).get("id")
-            if not bot_id:
-                print("Bot ID not found in the response.")
-                return None
-
-            print(f"Successfully created bot: {bot_id} in workspace: {workspace_id}")
-
-            # Install integration for the new bot
-            integration_success = install_bot_integration(bot_id, workspace_id)
-            if integration_success:
-                print(f"Successfully installed integration for bot {bot_id}")
+        r = requests.post(url, headers=headers, json=payload, timeout=60)
+        if r.status_code == 200:
+            bot_id = r.json().get("bot", {}).get("id")
+            if bot_id:
+                print(f"Bot created: {bot_id} in workspace: {workspace_id}")
+                # Optionally install integration (same as original)
+                installed = install_bot_integration(bot_id, workspace_id)
+                if installed:
+                    print(f"Integration installed for bot {bot_id}")
+                else:
+                    print(f"Integration install failed for bot {bot_id} (continuing)")
                 return bot_id
-            else:
-                print(f"Failed to install integration for bot {bot_id}")
-                return bot_id  # Still return the bot ID even if integration fails
-        else:
-            print(f"Bot creation failed with: {response.status_code}, {response.text}")
-            return None
+        print(f"Bot creation failed: {r.status_code} {r.text}")
     except Exception as e:
-        print(f"Error creating bot: {str(e)}")
-        return None
-
+        print(f"Error creating bot: {e}")
+    return None
 
 def install_bot_integration(bot_id, workspace_id):
-    """Install required integration for the bot to function properly"""
     if not bot_id or not workspace_id:
-        print("Cannot install integration: Missing bot ID or workspace ID")
+        print("Missing bot_id or workspace_id for integration")
         return False
-
     url = f"https://api.botpress.cloud/v1/admin/bots/{bot_id}"
     headers = {
         "User-Agent": "Mozilla/5.0",
@@ -103,7 +79,6 @@ def install_bot_integration(bot_id, workspace_id):
         "x-bot-id": bot_id,
         "x-workspace-id": workspace_id
     }
-    # Integration payload
     payload = {
         "integrations": {
             "intver_01JSCCBZGK79MXK2FSHKHTRKFD": {
@@ -111,90 +86,62 @@ def install_bot_integration(bot_id, workspace_id):
             }
         }
     }
-
     try:
-        response = requests.put(url, headers=headers, json=payload)
-        if response.status_code == 200:
-            print(f"Successfully installed integration for bot {bot_id}")
+        r = requests.put(url, headers=headers, json=payload, timeout=60)
+        if r.status_code == 200:
             return True
-        else:
-            print(f"Failed to install integration: {response.status_code}, {response.text}")
-            return False
+        print(f"Integration install failed: {r.status_code} {r.text}")
     except Exception as e:
-        print(f"Error installing integration: {str(e)}")
-        return False
-
+        print(f"Error installing integration: {e}")
+    return False
 
 def delete_bot(bot_id, workspace_id):
-    """Delete a bot from the specified workspace"""
     if not bot_id or not workspace_id:
-        print("Cannot delete bot: Missing bot ID or workspace ID")
+        print("Missing bot_id or workspace_id for bot deletion")
         return False
-
     url = f"https://api.botpress.cloud/v1/admin/bots/{bot_id}"
     headers = {
         "User-Agent": "Mozilla/5.0",
         "x-workspace-id": workspace_id,
         "Authorization": TOKEN
     }
-
     try:
-        response = requests.delete(url, headers=headers)
-        if response.status_code in [200, 204]:
-            print(f"Successfully deleted bot: {bot_id}")
+        r = requests.delete(url, headers=headers, timeout=60)
+        if r.status_code in [200, 204]:
+            print(f"Bot deleted: {bot_id}")
             return True
-        else:
-            print(f"Failed to delete bot: {response.status_code}, {response.text}")
-            return False
+        print(f"Bot deletion failed: {r.status_code} {r.text}")
     except Exception as e:
-        print(f"Error deleting bot: {str(e)}")
-        return False
-
+        print(f"Error deleting bot: {e}")
+    return False
 
 def delete_workspace(workspace_id):
-    """Delete a workspace"""
     if not workspace_id:
-        print("Cannot delete workspace: No workspace ID provided")
+        print("Missing workspace_id for workspace deletion")
         return False
-
     url = f"https://api.botpress.cloud/v1/admin/workspaces/{workspace_id}"
     headers = {
         "User-Agent": "Mozilla/5.0",
         "Authorization": TOKEN
     }
-
     try:
-        response = requests.delete(url, headers=headers)
-        if response.status_code in [200, 204]:
-            print(f"Successfully deleted workspace: {workspace_id}")
+        r = requests.delete(url, headers=headers, timeout=60)
+        if r.status_code in [200, 204]:
+            print(f"Workspace deleted: {workspace_id}")
             return True
-        else:
-            print(f"Failed to delete workspace: {response.status_code}, {response.text}")
-            return False
+        print(f"Workspace deletion failed: {r.status_code} {r.text}")
     except Exception as e:
-        print(f"Error deleting workspace: {str(e)}")
-        return False
-
+        print(f"Error deleting workspace: {e}")
+    return False
 
 # -------------------------------------------------------------------
-# Browse PAGE Function
+# Capture Screenshot (Botpress chat action)
 # -------------------------------------------------------------------
-def browse_pages(urls, wait_for=350, bot_id=None, workspace_id=None):
-    """
-    Browses specific web pages using the Botpress API
-    Args:
-        urls (list): List of URLs to browse
-        wait_for (int): Wait time in milliseconds
-        bot_id (str): The bot ID to use for the request
-        workspace_id (str): The workspace ID to use for the request
-    Returns:
-        tuple: (browse_results, bot_id, workspace_id)
-    """
+def capture_screenshot(url, width, height, full_page, bot_id, workspace_id, timeout=120):
     if not bot_id or not workspace_id:
-        print("Cannot perform browse: Missing bot ID or workspace ID")
-        return {"error": "Missing bot ID or workspace ID"}, bot_id, workspace_id
+        return {"error": "Missing bot_id or workspace_id"}
 
-    # Prepare the headers
+    api_url = "https://api.botpress.cloud/v1/chat/actions"
     headers = {
         "User-Agent": "Mozilla/5.0",
         "x-bot-id": bot_id,
@@ -202,212 +149,189 @@ def browse_pages(urls, wait_for=350, bot_id=None, workspace_id=None):
         "Content-Type": "application/json",
         "Authorization": TOKEN
     }
-
-    # Prepare the payload for the API
     payload = {
-        "type": "browser:browsePages",
+        "type": "browser:captureScreenshot",
         "input": {
-            "urls": urls,
-            "waitFor": wait_for
+            "url": url,
+            "width": width,
+            "height": height,
+            "fullPage": bool(full_page)
         }
     }
 
-    botpress_url = "https://api.botpress.cloud/v1/chat/actions"
     max_retries = 3
-    timeout = 60
-
-    # Flag to track if we need to create new IDs due to quota exceeded
-    quota_exceeded = False
-
-    # Attempt to send the request
     for attempt in range(max_retries):
         try:
-            print(f"Attempt {attempt+1}: Sending browse request to Botpress API with bot_id={bot_id}, workspace_id={workspace_id}")
-            response = requests.post(botpress_url, json=payload, headers=headers, timeout=timeout)
-
-            # If successful (200)
-            if response.status_code == 200:
-                data = response.json()
-                print(f"Successfully received browse results from Botpress API")
-                return data, bot_id, workspace_id
-
-            # Check for quota exceeded error specifically
-            elif response.status_code == 403:
-                error_data = response.json()
-                error_message = error_data.get('message', '')
-
-                # Check if this is the specific quota exceeded error
-                if "has reached its usage limit for ai spend" in error_message:
-                    print(f"Quota exceeded error detected: {error_message}")
-                    quota_exceeded = True
-                    break
-                else:
-                    print(f"Received 403 error but not quota exceeded: {error_message}")
-                    if attempt < max_retries - 1:
-                        time.sleep(2)
-                        continue
-                    else:
-                        return {"error": f"Unable to get browse results (Error 403)."}, bot_id, workspace_id
-
-            # Handle network errors or timeouts (just retry)
-            elif response.status_code in [404, 408, 502, 503, 504]:
-                print(f"Received error {response.status_code}. Retrying...")
-                time.sleep(3)  # Wait before retrying
+            print(f"Capture attempt {attempt+1} for {url}")
+            r = requests.post(api_url, headers=headers, json=payload, timeout=timeout)
+            if r.status_code == 200:
+                data = r.json()
+                # Expecting:
+                # { "output": {"imageUrl": "...", "htmlUrl": "..."}, "meta": {"cached": false}}
+                return data
+            elif r.status_code in [404, 408, 502, 503, 504]:
+                print(f"Transient error {r.status_code}: retrying...")
+                time.sleep(2)
                 continue
-
-            # Any other error status code
+            elif r.status_code == 403:
+                print(f"Forbidden (403): {r.text}")
+                # Not rotating bot/workspace here; per-request lifecycle handles cleanup
+                return {"error": "Forbidden", "status": 403, "details": safe_text(r)}
             else:
-                print(f"Received unexpected error: {response.status_code}, {response.text}")
-                if attempt < max_retries - 1:
-                    time.sleep(2)
-                    continue
-                else:
-                    return {"error": f"Unable to get browse results (Error {response.status_code})."}, bot_id, workspace_id
-
+                print(f"Unexpected {r.status_code}: {r.text}")
+                return {"error": "Unexpected error", "status": r.status_code, "details": safe_text(r)}
         except requests.exceptions.Timeout:
-            print(f"Request timed out. Retrying...")
+            print("Request timed out, retrying...")
             if attempt < max_retries - 1:
-                time.sleep(2)
                 continue
-            else:
-                return {"error": "The browse request timed out. Please try again."}, bot_id, workspace_id
-
+            return {"error": "Timeout"}
         except Exception as e:
-            print(f"Error during request: {str(e)}")
+            print(f"Exception: {e}")
             if attempt < max_retries - 1:
-                time.sleep(2)
+                time.sleep(1)
                 continue
-            else:
-                return {"error": f"Unable to get browse results: {str(e)}"}, bot_id, workspace_id
-
-    # If quota exceeded, we need to create new resources
-    if quota_exceeded:
-        print("Quota exceeded. Creating new workspace and bot...")
-
-        # First delete the bot, then the workspace (in that order)
-        if bot_id and workspace_id:
-            print(f"Deleting bot {bot_id} first...")
-            delete_success = delete_bot(bot_id, workspace_id)
-            if delete_success:
-                print(f"Successfully deleted bot {bot_id}")
-            else:
-                print(f"Failed to delete bot {bot_id}")
+            return {"error": str(e)}
 
-            # Wait 3 seconds before deleting the workspace
-            print("Waiting 3 seconds before deleting workspace...")
-            time.sleep(3)
-
-            print(f"Now deleting workspace {workspace_id}...")
-            ws_delete_success = delete_workspace(workspace_id)
-            if ws_delete_success:
-                print(f"Successfully deleted workspace {workspace_id}")
-            else:
-                print(f"Failed to delete workspace {workspace_id}")
-
-        # Create new workspace
-        new_workspace_id = create_workspace()
-        if not new_workspace_id:
-            return {"error": "Failed to create a new workspace after quota exceeded. Please try again later."}, bot_id, workspace_id
-
-        # Create new bot in the new workspace
-        new_bot_id = create_bot(new_workspace_id)
-        if not new_bot_id:
-            return {"error": "Failed to create a new bot after quota exceeded. Please try again later."}, new_workspace_id, workspace_id
-
-        # Update headers with new bot ID
-        headers["x-bot-id"] = new_bot_id
-
-        # Try one more time with the new IDs
-        try:
-            print(f"Retrying with new bot_id={new_bot_id}, workspace_id={new_workspace_id}")
-            retry_response = requests.post(botpress_url, json=payload, headers=headers, timeout=timeout)
-
-            if retry_response.status_code == 200:
-                data = retry_response.json()
-                print(f"Successfully received browse results with new IDs")
-                return data, new_bot_id, new_workspace_id
-            else:
-                print(f"Failed with new IDs: {retry_response.status_code}, {retry_response.text}")
-                return {"error": f"Unable to get browse results with new credentials."}, new_bot_id, new_workspace_id
-
-        except Exception as e:
-            print(f"Error with new IDs: {str(e)}")
-            return {"error": f"Unable to get browse results with new credentials: {str(e)}"}, new_bot_id, new_workspace_id
-
-    # Should not reach here due to the handling in the loop
-    return {"error": "Unable to get browse results."}, bot_id, workspace_id
+    return {"error": "Failed after retries"}
 
+def safe_text(response):
+    try:
+        return response.text
+    except Exception:
+        return "<unavailable>"
 
 # -------------------------------------------------------------------
-# Flask Endpoints
+# Flask Endpoint
 # -------------------------------------------------------------------
-@app.route("/browse", methods=["POST"])
-def browse_endpoint():
+@app.route("/capture", methods=["POST"])
+def capture_endpoint():
     """
-    Endpoint for browsing specific web pages.
-    Expects JSON with:
+    JSON body example:
+    {
+      "urls": ["https://x.com/USTreasury", "https://example.com"],   // required, array of one or more
+      "width": 10850,                                                // optional, default 1080
+      "height": 1920,                                                // optional, default 1920
+      "fullPage": true                                               // optional, default true
+    }
+    Response:
     {
-      "urls": ["https://example.com", "https://example.org"],
-      "wait_for": 350                // Optional, defaults to 350ms
+      "results": [
+        {
+          "url": "...",
+          "output": {
+            "imageUrl": "...png",
+            "htmlUrl": "...html"
+          },
+          "meta": { "cached": false }
+        },
+        { ... }
+      ],
+      "errors": [
+        { "url": "...", "error": "...", "status": 403, "details": "..." }
+      ]
     }
-    Returns the raw browse results JSON from the Botpress API
+    Behavior:
+    - Create workspace, then bot.
+    - Run capture for all URLs with the same bot/workspace.
+    - Delete bot first, then workspace, at the end of the request.
     """
-    global GLOBAL_WORKSPACE_ID, GLOBAL_BOT_ID
-
-    # Parse JSON from request
-    data = request.get_json(force=True)
-    urls = data.get("urls", [])
-    wait_for = data.get("wait_for", 350)
-
-    # Validate URLs
-    if not urls or not isinstance(urls, list):
-        return jsonify({"error": "URLs parameter is required and must be a list"}), 400
-
-    # Validate wait_for
+    # Parse and validate body
     try:
-        wait_for = int(wait_for)
-        if wait_for <= 0:
-            wait_for = 350
-            print(f"Invalid wait_for value. Using default: {wait_for}")
-    except (ValueError, TypeError):
-        wait_for = 350
-        print(f"Invalid wait_for format. Using default: {wait_for}")
-
-    # If we don't yet have a workspace or bot, create them
-    if not GLOBAL_WORKSPACE_ID or not GLOBAL_BOT_ID:
-        print("No existing IDs found. Creating new workspace and bot...")
-        GLOBAL_WORKSPACE_ID = create_workspace()
-        if GLOBAL_WORKSPACE_ID:
-            GLOBAL_BOT_ID = create_bot(GLOBAL_WORKSPACE_ID)
+        body = request.get_json(force=True) or {}
+    except Exception:
+        return jsonify({"error": "Invalid JSON"}), 400
 
-        # If creation failed
-        if not GLOBAL_WORKSPACE_ID or not GLOBAL_BOT_ID:
-            return jsonify({"error": "Failed to initialize browse capability. Please try again later."}), 500
-
-    # Call our browse function
-    print(f"Sending browse request with urls='{urls}', wait_for={wait_for}")
-    print(f"Using existing bot_id={GLOBAL_BOT_ID}, workspace_id={GLOBAL_WORKSPACE_ID}")
-
-    browse_results, updated_bot_id, updated_workspace_id = browse_pages(
-        urls,
-        wait_for,
-        GLOBAL_BOT_ID,
-        GLOBAL_WORKSPACE_ID
-    )
-
-    # Update global IDs if they changed
-    if updated_bot_id != GLOBAL_BOT_ID or updated_workspace_id != GLOBAL_WORKSPACE_ID:
-        print(f"Updating global IDs: bot_id={updated_bot_id}, workspace_id={updated_workspace_id}")
-        GLOBAL_BOT_ID = updated_bot_id
-        GLOBAL_WORKSPACE_ID = updated_workspace_id
+    urls = body.get("urls")
+    if not urls or not isinstance(urls, list):
+        return jsonify({"error": "urls is required and must be a list"}), 400
 
-    # Return the raw API response
-    return jsonify(browse_results)
+    # Optional params
+    width = body.get("width", 1080)
+    height = body.get("height", 1920)
+    full_page = body.get("fullPage", True)
 
+    # Normalize numeric params
+    try:
+        width = int(width)
+        if width <= 0:
+            width = 1080
+    except Exception:
+        width = 1080
+    try:
+        height = int(height)
+        if height <= 0:
+            height = 1920
+    except Exception:
+        height = 1920
+    full_page = bool(full_page)
+
+    workspace_id = None
+    bot_id = None
+
+    # Create (Workspace then Bot)
+    workspace_id = create_workspace()
+    if not workspace_id:
+        return jsonify({"error": "Failed to create workspace"}), 500
+
+    bot_id = create_bot(workspace_id)
+    if not bot_id:
+        # Cleanup: if bot creation fails, delete workspace
+        delete_workspace(workspace_id)
+        return jsonify({"error": "Failed to create bot"}), 500
+
+    results = []
+    errors = []
+
+    # Process multiple URLs using the same bot/workspace
+    for url in urls:
+        if not isinstance(url, str) or not url.strip():
+            errors.append({"url": url, "error": "Invalid URL"})
+            continue
+
+        capture = capture_screenshot(
+            url=url.strip(),
+            width=width,
+            height=height,
+            full_page=full_page,
+            bot_id=bot_id,
+            workspace_id=workspace_id
+        )
+
+        if isinstance(capture, dict) and capture.get("error"):
+            errors.append({"url": url, **capture})
+        else:
+            # Ensure consistent output structure
+            output = capture.get("output", {})
+            meta = capture.get("meta", {})
+            results.append({
+                "url": url,
+                "output": {
+                    "imageUrl": output.get("imageUrl"),
+                    "htmlUrl": output.get("htmlUrl")
+                },
+                "meta": meta
+            })
+
+    # Delete (Bot first, then Workspace) regardless of success
+    try:
+        if bot_id:
+            delete_bot(bot_id, workspace_id)
+    finally:
+        if workspace_id:
+            delete_workspace(workspace_id)
+
+    response_body = {"results": results}
+    if errors:
+        response_body["errors"] = errors
+
+    # If all failed, return 502; else 200
+    if results:
+        return jsonify(response_body), 200
+    else:
+        return jsonify(response_body), 502
 
 # -------------------------------------------------------------------
 # Run the Flask app
 # -------------------------------------------------------------------
-
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860, debug=True)