sentienceapi 0.92.2__py3-none-any.whl → 0.98.0__py3-none-any.whl

sentience/backends/snapshot.py

@@ -0,0 +1,483 @@
+"""
+Backend-agnostic snapshot for browser-use integration.
+
+Takes Sentience snapshots using BrowserBackend protocol,
+enabling element grounding with browser-use or other frameworks.
+
+Usage with browser-use:
+    from sentience.backends import BrowserUseAdapter, snapshot, CachedSnapshot
+
+    adapter = BrowserUseAdapter(session)
+    backend = await adapter.create_backend()
+
+    # Take snapshot
+    snap = await snapshot(backend)
+    print(f"Found {len(snap.elements)} elements")
+
+    # With caching (reuse if fresh)
+    cache = CachedSnapshot(backend, max_age_ms=2000)
+    snap1 = await cache.get()  # Fresh snapshot
+    snap2 = await cache.get()  # Returns cached if < 2s old
+    cache.invalidate()  # Force refresh on next get()
+"""
+
+import asyncio
+import time
+from typing import TYPE_CHECKING, Any
+
+from ..constants import SENTIENCE_API_URL
+from ..models import Snapshot, SnapshotOptions
+from ..snapshot import (
+    _build_snapshot_payload,
+    _merge_api_result_with_local,
+    _post_snapshot_to_gateway_async,
+)
+from .exceptions import ExtensionDiagnostics, ExtensionNotLoadedError, SnapshotError
+
+if TYPE_CHECKING:
+    from .protocol import BrowserBackend
+
+
+def _is_execution_context_destroyed_error(e: Exception) -> bool:
+    """
+    Playwright (and other browser backends) can throw while a navigation is in-flight.
+
+    Common symptoms:
+    - "Execution context was destroyed, most likely because of a navigation"
+    - "Cannot find context with specified id"
+    """
+    msg = str(e).lower()
+    return (
+        "execution context was destroyed" in msg
+        or "most likely because of a navigation" in msg
+        or "cannot find context with specified id" in msg
+    )
+
+
+async def _eval_with_navigation_retry(
+    backend: "BrowserBackend",
+    expression: str,
+    *,
+    retries: int = 10,
+    settle_state: str = "interactive",
+    settle_timeout_ms: int = 10000,
+) -> Any:
+    """
+    Evaluate JS, retrying once/ twice if the page is mid-navigation.
+
+    This makes snapshots resilient to cases like:
+    - press Enter (navigation) → snapshot immediately → context destroyed
+    """
+    last_err: Exception | None = None
+    for attempt in range(retries + 1):
+        try:
+            return await backend.eval(expression)
+        except Exception as e:
+            last_err = e
+            if not _is_execution_context_destroyed_error(e) or attempt >= retries:
+                raise
+            # Navigation is in-flight; wait for new document context then retry.
+            try:
+                await backend.wait_ready_state(state=settle_state, timeout_ms=settle_timeout_ms)  # type: ignore[arg-type]
+            except Exception:
+                # If readyState polling also fails mid-nav, still retry after a short backoff.
+                pass
+            # Exponential-ish backoff (caps quickly), tuned for real navigations.
+            await asyncio.sleep(min(0.25 * (attempt + 1), 1.5))
+
+    # Unreachable in practice, but keeps type-checkers happy.
+    raise last_err if last_err else RuntimeError("eval failed")
+
+
+class CachedSnapshot:
+    """
+    Snapshot cache with staleness detection.
+
+    Caches snapshots and returns cached version if still fresh.
+    Useful for reducing redundant snapshot calls in action loops.
+
+    Usage:
+        cache = CachedSnapshot(backend, max_age_ms=2000)
+
+        # First call takes fresh snapshot
+        snap1 = await cache.get()
+
+        # Second call returns cached if < 2s old
+        snap2 = await cache.get()
+
+        # Invalidate after actions that change DOM
+        await click(backend, element.bbox)
+        cache.invalidate()
+
+        # Next get() will take fresh snapshot
+        snap3 = await cache.get()
+    """
+
+    def __init__(
+        self,
+        backend: "BrowserBackend",
+        max_age_ms: int = 2000,
+        options: SnapshotOptions | None = None,
+    ) -> None:
+        """
+        Initialize cached snapshot.
+
+        Args:
+            backend: BrowserBackend implementation
+            max_age_ms: Maximum cache age in milliseconds (default: 2000)
+            options: Default snapshot options
+        """
+        self._backend = backend
+        self._max_age_ms = max_age_ms
+        self._options = options
+        self._cached: Snapshot | None = None
+        self._cached_at: float = 0  # timestamp in seconds
+        self._cached_url: str | None = None
+
+    async def get(
+        self,
+        options: SnapshotOptions | None = None,
+        force_refresh: bool = False,
+    ) -> Snapshot:
+        """
+        Get snapshot, using cache if fresh.
+
+        Args:
+            options: Override default options for this call
+            force_refresh: If True, always take fresh snapshot
+
+        Returns:
+            Snapshot (cached or fresh)
+        """
+        # Check if we need to refresh
+        if force_refresh or self._is_stale():
+            self._cached = await snapshot(
+                self._backend,
+                options or self._options,
+            )
+            self._cached_at = time.time()
+            self._cached_url = self._cached.url
+
+        assert self._cached is not None
+        return self._cached
+
+    def invalidate(self) -> None:
+        """
+        Invalidate cache, forcing refresh on next get().
+
+        Call this after actions that modify the DOM.
+        """
+        self._cached = None
+        self._cached_at = 0
+        self._cached_url = None
+
+    def _is_stale(self) -> bool:
+        """Check if cache is stale and needs refresh."""
+        if self._cached is None:
+            return True
+
+        # Check age
+        age_ms = (time.time() - self._cached_at) * 1000
+        if age_ms > self._max_age_ms:
+            return True
+
+        return False
+
+    @property
+    def is_cached(self) -> bool:
+        """Check if a cached snapshot exists."""
+        return self._cached is not None
+
+    @property
+    def age_ms(self) -> float:
+        """Get age of cached snapshot in milliseconds."""
+        if self._cached is None:
+            return float("inf")
+        return (time.time() - self._cached_at) * 1000
+
+
+async def snapshot(
+    backend: "BrowserBackend",
+    options: SnapshotOptions | None = None,
+) -> Snapshot:
+    """
+    Take a Sentience snapshot using the backend protocol.
+
+    This function respects the `use_api` option and can call either:
+    - Server-side API (Pro/Enterprise tier) when `use_api=True` and API key is provided
+    - Local extension (Free tier) when `use_api=False` or no API key
+
+    Requires:
+        - Sentience extension loaded in browser (via --load-extension)
+        - Extension injected window.sentience API
+
+    Args:
+        backend: BrowserBackend implementation (CDPBackendV0, PlaywrightBackend, etc.)
+        options: Snapshot options (limit, filter, screenshot, use_api, sentience_api_key, etc.)
+
+    Returns:
+        Snapshot with elements, viewport, and optional screenshot
+
+    Example:
+        from sentience.backends import BrowserUseAdapter
+        from sentience.backends.snapshot import snapshot
+        from sentience.models import SnapshotOptions
+
+        adapter = BrowserUseAdapter(session)
+        backend = await adapter.create_backend()
+
+        # Basic snapshot (uses local extension)
+        snap = await snapshot(backend)
+
+        # With server-side API (Pro/Enterprise tier)
+        snap = await snapshot(backend, SnapshotOptions(
+            use_api=True,
+            sentience_api_key="sk_pro_xxxxx",
+            limit=100,
+            screenshot=True
+        ))
+
+        # Force local extension (Free tier)
+        snap = await snapshot(backend, SnapshotOptions(
+            use_api=False
+        ))
+    """
+    if options is None:
+        options = SnapshotOptions()
+
+    # Determine if we should use server-side API
+    # Same logic as main snapshot() function in sentience/snapshot.py
+    should_use_api = (
+        options.use_api if options.use_api is not None else (options.sentience_api_key is not None)
+    )
+
+    if should_use_api and options.sentience_api_key:
+        # Use server-side API (Pro/Enterprise tier)
+        return await _snapshot_via_api(backend, options)
+    else:
+        # Use local extension (Free tier)
+        return await _snapshot_via_extension(backend, options)
+
+
+async def _wait_for_extension(
+    backend: "BrowserBackend",
+    timeout_ms: int = 5000,
+) -> None:
+    """
+    Wait for Sentience extension to inject window.sentience API.
+
+    Args:
+        backend: BrowserBackend implementation
+        timeout_ms: Maximum wait time
+
+    Raises:
+        RuntimeError: If extension not injected within timeout
+    """
+    import asyncio
+    import logging
+
+    logger = logging.getLogger("sentience.backends.snapshot")
+
+    start = time.monotonic()
+    timeout_sec = timeout_ms / 1000.0
+    poll_count = 0
+
+    logger.debug(f"Waiting for extension injection (timeout={timeout_ms}ms)...")
+
+    while True:
+        elapsed = time.monotonic() - start
+        poll_count += 1
+
+        if poll_count % 10 == 0:  # Log every 10 polls (~1 second)
+            logger.debug(f"Extension poll #{poll_count}, elapsed={elapsed*1000:.0f}ms")
+
+        if elapsed >= timeout_sec:
+            # Gather diagnostics
+            try:
+                diag_dict = await backend.eval(
+                    """
+                    (() => ({
+                        sentience_defined: typeof window.sentience !== 'undefined',
+                        sentience_snapshot: typeof window.sentience?.snapshot === 'function',
+                        url: window.location.href,
+                        extension_id: document.documentElement.dataset.sentienceExtensionId || null,
+                        has_content_script: !!document.documentElement.dataset.sentienceExtensionId
+                    }))()
+                """
+                )
+                diagnostics = ExtensionDiagnostics.from_dict(diag_dict)
+                logger.debug(f"Extension diagnostics: {diag_dict}")
+            except Exception as e:
+                diagnostics = ExtensionDiagnostics(error=f"Could not gather diagnostics: {e}")
+
+            raise ExtensionNotLoadedError.from_timeout(
+                timeout_ms=timeout_ms,
+                diagnostics=diagnostics,
+            )
+
+        # Check if extension is ready
+        try:
+            ready = await backend.eval(
+                "typeof window.sentience !== 'undefined' && "
+                "typeof window.sentience.snapshot === 'function'"
+            )
+            if ready:
+                return
+        except Exception:
+            pass  # Keep polling
+
+        await asyncio.sleep(0.1)
+
+
+async def _snapshot_via_extension(
+    backend: "BrowserBackend",
+    options: SnapshotOptions,
+) -> Snapshot:
+    """Take snapshot using local extension (Free tier)"""
+    # Wait for extension injection
+    await _wait_for_extension(backend, timeout_ms=5000)
+
+    # Build options dict for extension API
+    ext_options = _build_extension_options(options)
+
+    # Call extension's snapshot function
+    result = await _eval_with_navigation_retry(
+        backend,
+        f"""
+        (() => {{
+            const options = {_json_serialize(ext_options)};
+            return window.sentience.snapshot(options);
+        }})()
+    """,
+    )
+
+    if result is None:
+        # Try to get URL for better error message
+        try:
+            url = await backend.eval("window.location.href")
+        except Exception:
+            url = None
+        raise SnapshotError.from_null_result(url=url)
+
+    # Show overlay if requested
+    if options.show_overlay:
+        raw_elements = result.get("raw_elements", [])
+        if raw_elements:
+            await _eval_with_navigation_retry(
+                backend,
+                f"""
+                (() => {{
+                    if (window.sentience && window.sentience.showOverlay) {{
+                        window.sentience.showOverlay({_json_serialize(raw_elements)}, null);
+                    }}
+                }})()
+            """,
+            )
+
+    # Build and return Snapshot
+    return Snapshot(**result)
+
+
+async def _snapshot_via_api(
+    backend: "BrowserBackend",
+    options: SnapshotOptions,
+) -> Snapshot:
+    """Take snapshot using server-side API (Pro/Enterprise tier)"""
+    # Default API URL (same as main snapshot function)
+    api_url = SENTIENCE_API_URL
+
+    # Wait for extension injection (needed even for API mode to collect raw data)
+    await _wait_for_extension(backend, timeout_ms=5000)
+
+    # Step 1: Get raw data from local extension (always happens locally)
+    raw_options: dict[str, Any] = {}
+    if options.screenshot is not False:
+        raw_options["screenshot"] = options.screenshot
+
+    # Call extension to get raw elements
+    raw_result = await _eval_with_navigation_retry(
+        backend,
+        f"""
+        (() => {{
+            const options = {_json_serialize(raw_options)};
+            return window.sentience.snapshot(options);
+        }})()
+    """,
+    )
+
+    if raw_result is None:
+        try:
+            url = await backend.eval("window.location.href")
+        except Exception:
+            url = None
+        raise SnapshotError.from_null_result(url=url)
+
+    # Step 2: Send to server for smart ranking/filtering
+    payload = _build_snapshot_payload(raw_result, options)
+
+    try:
+        api_result = await _post_snapshot_to_gateway_async(
+            payload, options.sentience_api_key, api_url
+        )
+
+        # Merge API result with local data (screenshot, etc.)
+        snapshot_data = _merge_api_result_with_local(api_result, raw_result)
+
+        # Show visual overlay if requested (use API-ranked elements)
+        if options.show_overlay:
+            elements = api_result.get("elements", [])
+            if elements:
+                await _eval_with_navigation_retry(
+                    backend,
+                    f"""
+                    (() => {{
+                        if (window.sentience && window.sentience.showOverlay) {{
+                            window.sentience.showOverlay({_json_serialize(elements)}, null);
+                        }}
+                    }})()
+                """,
+                )
+
+        return Snapshot(**snapshot_data)
+    except (RuntimeError, ValueError):
+        # Re-raise validation errors as-is
+        raise
+    except Exception as e:
+        # Fallback to local extension on API error
+        # This matches the behavior of the main snapshot function
+        raise RuntimeError(
+            f"Server-side snapshot API failed: {e}. "
+            "Try using use_api=False to use local extension instead."
+        ) from e
+
+
+def _build_extension_options(options: SnapshotOptions) -> dict[str, Any]:
+    """Build options dict for extension API call."""
+    ext_options: dict[str, Any] = {}
+
+    # Screenshot config
+    if options.screenshot is not False:
+        if hasattr(options.screenshot, "model_dump"):
+            ext_options["screenshot"] = options.screenshot.model_dump()
+        else:
+            ext_options["screenshot"] = options.screenshot
+
+    # Limit (only if not default)
+    if options.limit != 50:
+        ext_options["limit"] = options.limit
+
+    # Filter
+    if options.filter is not None:
+        if hasattr(options.filter, "model_dump"):
+            ext_options["filter"] = options.filter.model_dump()
+        else:
+            ext_options["filter"] = options.filter
+
+    return ext_options
+
+
+def _json_serialize(obj: Any) -> str:
+    """Serialize object to JSON string for embedding in JS."""
+    import json
+
+    return json.dumps(obj)