arena/browserbase_tools.py

"""Browserbase Fetch and Search tools for the DeepAgent swarm."""

from __future__ import annotations

import json
import os
from typing import Any

from browserbase import Browserbase


def _client() -> Browserbase:
    api_key = os.environ.get("BROWSERBASE_API_KEY", "")
    if not api_key:
        raise RuntimeError("BROWSERBASE_API_KEY is required for fetch/search tools")
    return Browserbase(api_key=api_key)


def browserbase_fetch(url: str, allow_redirects: bool = False) -> str:
    """Fetch and return the content of a web page without a browser.

    Use this to retrieve HTML, JSON, or text content from a URL.
    Best for static pages or API endpoints that don't require JavaScript.

    Args:
        url: The URL to fetch.
        allow_redirects: Whether to follow HTTP redirects (default False).

    Returns:
        The page content as a string, or an error message if the fetch fails.
    """
    try:
        response = _client().fetch_api.create(url=url, allow_redirects=allow_redirects)
        if hasattr(response, "content"):
            content = response.content
            if len(content) > 10000:
                content = content[:10000] + "\n...[truncated]"
            return str(content)
        return str(response)
    except Exception as exc:
        return f"Fetch error: {exc}"


def browserbase_search(query: str, num_results: int = 5) -> str:
    """Search the web and return structured results without a browser.

    Use this to find URLs, documentation, or reference material.
    Returns title, URL, and description for each result.

    Args:
        query: The search query string.
        num_results: Number of results to return (1-25, default 5).

    Returns:
        Formatted JSON string with search results, or an error message.
    """
    try:
        response = _client().search_api.create(query=query, num_results=num_results)
        if hasattr(response, "results"):
            formatted = []
            for r in response.results:
                entry = {}
                if hasattr(r, "title"):
                    entry["title"] = r.title
                if hasattr(r, "url"):
                    entry["url"] = r.url
                if hasattr(r, "description"):
                    entry["description"] = r.description
                formatted.append(entry)
            return json.dumps(formatted, indent=2, ensure_ascii=False)
        return str(response)
    except Exception as exc:
        return f"Search error: {exc}"


def _fetch_tool_schema() -> dict[str, Any]:
    return {
        "name": "browserbase_fetch",
        "description": "Fetch and return the content of a web page without a browser. Use for static pages or APIs.",
        "parameters": {
            "type": "object",
            "properties": {
                "url": {
                    "type": "string",
                    "description": "The URL to fetch.",
                },
                "allow_redirects": {
                    "type": "boolean",
                    "description": "Whether to follow HTTP redirects (default False).",
                },
            },
            "required": ["url"],
        },
    }


def _search_tool_schema() -> dict[str, Any]:
    return {
        "name": "browserbase_search",
        "description": "Search the web for documentation, references, or examples. Returns title, URL, and description.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "The search query string.",
                },
                "num_results": {
                    "type": "integer",
                    "description": "Number of results to return (1-25, default 5).",
                },
            },
            "required": ["query"],
        },
    }


SWARM_TOOL_SCHEMAS: list[dict[str, Any]] = [_fetch_tool_schema(), _search_tool_schema()]

SWARM_TOOLS: dict[str, Any] = {
    "browserbase_fetch": browserbase_fetch,
    "browserbase_search": browserbase_search,
}