r2 uploads

Author: CyrusNuevoDia

lib/asyncio.py

@@ -0,0 +1,6 @@
+from functools import partial
+
+import asyncer
+
+syncify = partial(asyncer.syncify, raise_sync_error=False)
+asyncify = partial(asyncer.asyncify, abandon_on_cancel=True)

lib/browser.py

@@ -1,7 +1,8 @@
 import logging
+from pathlib import Path
 from os import environ as env
 
-from browser_use import BrowserProfile
+from browser_use import Agent, BrowserProfile
 from kernel import AsyncKernel, KernelContext
 from kernel.types import BrowserCreateParams
 
@@ -14,6 +15,12 @@
 logger = logging.getLogger(__name__)
 
 
+def downloaded_files(agent: Agent) -> list[Path]:
+    if downloads_path := agent.browser_profile.downloads_path:
+        return list(Path(downloads_path).glob("*"))
+    return []
+
+
 async def create_browser(ctx: KernelContext, request: BrowserAgentRequest):
     invocation_id = ctx.invocation_id
     headless = request.headless

lib/models.py

@@ -1,5 +1,6 @@
-from pathlib import Path
-from browser_use import Agent, AgentHistoryList
+import typing as t
+
+from browser_use.agent.views import AgentHistoryList
 from pydantic import BaseModel
 
 from lib.ai import ModelProvider
@@ -24,21 +25,20 @@ class BrowserAgentRequest(BaseModel):
 
 class BrowserAgentResponse(BaseModel):
     session: str
+    success: bool
     duration: float
     result: str
-    files: list[str]
-    errors: list[str]
+    downloads: dict[str, str]
 
     @classmethod
-    def build(cls, session: str, agent: Agent, history: AgentHistoryList):
-        files: list[str] = []
-        if downloads_path := agent.browser_profile.downloads_path:
-            files = [str(path) for path in Path(downloads_path).glob("*")]
-
+    def from_run(
+        cls,
+        run: AgentHistoryList,
+        **kwargs,
+    ) -> t.Self:
         return cls(
-            session=session,
-            duration=history.total_duration_seconds(),
-            result=history.final_result(),
-            files=files,
-            errors=list(filter(bool, history.errors())),
+            success=run.is_successful(),
+            duration=run.total_duration_seconds(),
+            result=run.final_result(),
+            **kwargs,
         )

lib/storage.py

@@ -0,0 +1,59 @@
+import asyncio
+import typing as t
+from os import environ as env
+from pathlib import Path
+
+import boto3
+import orjson
+from botocore.client import Config
+
+from lib.asyncio import asyncify
+
+BUCKET = env.get("R2_S3_BUCKET", "browser-agent")
+PRESIGNED_URL_EXPIRES_IN = 24 * 60 * 60  # 12 hours
+
+client = boto3.client(
+    service_name="s3",
+    endpoint_url=env["R2_S3_ENDPOINT_URL"],
+    aws_access_key_id=env["R2_S3_ACCESS_KEY_ID"],
+    aws_secret_access_key=env["R2_S3_SECRET_ACCESS_KEY"],
+    region_name="auto",
+    config=Config(signature_version="s3v4"),
+)
+
+
+@asyncify
+def upload_file(file: Path | str, key: str):
+    client.upload_file(
+        Bucket=BUCKET,
+        File=str(file),
+        Key=key,
+    )
+    return client.generate_presigned_url(
+        "get_object",
+        Params={"Bucket": BUCKET, "Key": key},
+        ExpiresIn=PRESIGNED_URL_EXPIRES_IN,
+    )
+
+
+@asyncify
+def upload_json(data: t.Any, key: str) -> str:
+    client.put_object(
+        Bucket=BUCKET,
+        Key=key,
+        Body=orjson.dumps(data, option=orjson.OPT_INDENT_2),
+    )
+    return client.generate_presigned_url(
+        "get_object",
+        Params={"Bucket": BUCKET, "Key": key},
+        ExpiresIn=PRESIGNED_URL_EXPIRES_IN,
+    )
+
+
+async def upload_files(
+    dir: str,
+    files: list[Path | str],
+) -> dict[str, str]:
+    keys = [f"{dir}/{Path(f).name}" for f in files]
+    urls = await asyncio.gather(*[upload_file(f, k) for f, k in zip(files, keys)])
+    return dict(zip(keys, urls))

main.py

@@ -1,12 +1,14 @@
+import asyncio
 import logging
 
 from browser_use import Agent
 from kernel import App, KernelContext
 from zenbase_llml import llml
 
 from lib.ai import AGENT_INSTRUCTIONS, ChatFactory
-from lib.browser import create_browser
+from lib.browser import create_browser, downloaded_files
 from lib.models import BrowserAgentRequest, BrowserAgentResponse
+from lib.storage import upload_files, upload_json
 
 logger = logging.getLogger(__name__)
 
@@ -31,7 +33,7 @@ async def perform(ctx: KernelContext, params: dict):
         "input": request.input,
         "notes": """
         Your browser will automatically:
-        1. Download the PDF file upon viewing it. Just wait for it.
+        1. Download the PDF file upon viewing it. Just wait for it. You do not need to read the PDF.
         2. Solve CAPTCHAs or similar tests. Just wait for it.
         """,
     }
@@ -44,7 +46,16 @@ async def perform(ctx: KernelContext, params: dict):
         flash_mode=request.flash,
     )
 
-    history = await agent.run(max_steps=request.max_steps)
+    trajectory = await agent.run(max_steps=request.max_steps)
 
-    response = BrowserAgentResponse.build(session, agent, history)
+    (downloads,) = await asyncio.gather(
+        upload_files(dir=session, files=downloaded_files(agent)),
+        upload_json(trajectory.model_dump(), key=f"{session}/trajectory.json"),
+    )
+
+    response = BrowserAgentResponse.from_run(
+        trajectory,
+        session=session,
+        downloads=downloads,
+    )
     return response.model_dump()

pyproject.toml

@@ -6,8 +6,11 @@ readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
     "anyio>=4.10.0",
+    "asyncer>=0.0.8",
+    "boto3>=1.40.25",
     "browser-use>=0.7.2",
     "kernel>=0.11.0",
+    "orjson>=3.11.3",
     "pydantic>=2.10.6",
     "sorcery>=0.2.2",
     "zenbase-llml>=0.4.0",