fix: Use tool secrets in toolkits (#271)

~~Note: Don't merge until the correct secrets have been added to Arcade
Cloud.~~

Ready to merge, the feature is already on its way to prod.

---------

Co-authored-by: Eric Gustin <eric@arcade.dev>

toolkits/code_sandbox/arcade_code_sandbox/tools/e2b.py

@@ -1,16 +1,16 @@
 from typing import Annotated
 
-from arcade.sdk import tool
+from arcade.sdk import ToolContext, tool
 from e2b_code_interpreter import Sandbox
 
 from arcade_code_sandbox.tools.models import E2BSupportedLanguage
-from arcade_code_sandbox.tools.utils import get_secret
 
 # See https://e2b.dev/docs to learn more about E2B
 
 
-@tool
+@tool(requires_secrets=["E2B_API_KEY"])
 def run_code(
+    context: ToolContext,
     code: Annotated[str, "The code to run"],
     language: Annotated[
         E2BSupportedLanguage, "The language of the code"
@@ -19,7 +19,7 @@ def run_code(
     """
     Run code in a sandbox and return the output.
     """
-    api_key = get_secret("E2B_API_KEY")
+    api_key = context.get_secret("E2B_API_KEY")
 
     with Sandbox(api_key=api_key) as sbx:
         execution = sbx.run_code(code=code, language=language)
@@ -29,15 +29,16 @@ def run_code(
 
 # Note: Not recommended to use tool_choice='generate' with this tool
 #       since it contains base64 encoded image.
-@tool
+@tool(requires_secrets=["E2B_API_KEY"])
 def create_static_matplotlib_chart(
+    context: ToolContext,
     code: Annotated[str, "The Python code to run"],
 ) -> Annotated[dict, "A dictionary with the following keys: base64_image, logs, error"]:
     """
     Run the provided Python code to generate a static matplotlib chart.
     The resulting chart is returned as a base64 encoded image.
     """
-    api_key = get_secret("E2B_API_KEY")
+    api_key = context.get_secret("E2B_API_KEY")
 
     with Sandbox(api_key=api_key) as sbx:
         execution = sbx.run_code(code=code)

toolkits/code_sandbox/arcade_code_sandbox/tools/utils.py

@@ -1,9 +0,0 @@
-import os
-from typing import Any, Optional
-
-
-def get_secret(name: str, default: Optional[Any] = None) -> Any:
-    secret = os.getenv(name)
-    if secret is None and default is not None:
-        return default
-    return secret

toolkits/code_sandbox/pyproject.toml

@@ -1,12 +1,12 @@
 [tool.poetry]
 name = "arcade_code_sandbox"
-version = "0.1.9"
+version = "1.0.0"
 description = "LLM tools for running code in a sandbox"
 authors = ["Arcade <dev@arcade.dev>"]
 
 [tool.poetry.dependencies]
 python = "^3.10"
-arcade-ai = ">=0.1,<2.0"
+arcade-ai = ">=1.0.5,<2.0"
 e2b-code-interpreter = "^1.0.1"
 
 [tool.poetry.dev-dependencies]

toolkits/code_sandbox/tests/test_e2b.py

@@ -1,6 +1,8 @@
 from unittest.mock import MagicMock, patch
 
 import pytest
+from arcade.core.schema import ToolSecretItem
+from arcade.sdk import ToolContext
 from arcade.sdk.errors import ToolExecutionError
 
 from arcade_code_sandbox.tools.e2b import create_static_matplotlib_chart, run_code
@@ -13,32 +15,37 @@ def mock_sandbox():
         yield mock.return_value.__enter__.return_value
 
 
-def test_run_code_success(mock_sandbox):
+@pytest.fixture
+def mock_context():
+    return ToolContext(secrets=[ToolSecretItem(key="e2b_api_key", value="fake_api_key")])
+
+
+def test_run_code_success(mock_sandbox, mock_context):
     mock_execution = MagicMock()
     mock_execution.to_json.return_value = '{"result": "success"}'
     mock_sandbox.run_code.return_value = mock_execution
 
-    result = run_code("print('Hello, World!')", E2BSupportedLanguage.PYTHON)
+    result = run_code(mock_context, "print('Hello, World!')", E2BSupportedLanguage.PYTHON)
     assert result == '{"result": "success"}'
 
 
-def test_run_code_error(mock_sandbox):
+def test_run_code_error(mock_sandbox, mock_context):
     mock_execution = MagicMock()
     mock_execution.to_json.side_effect = ToolExecutionError("Execution failed")
     mock_sandbox.run_code.return_value = mock_execution
 
     with pytest.raises(ToolExecutionError, match="Execution failed"):
-        run_code("print('Hello, World!')", E2BSupportedLanguage.PYTHON)
+        run_code(mock_context, "print('Hello, World!')", E2BSupportedLanguage.PYTHON)
 
 
-def test_create_static_matplotlib_chart_success(mock_sandbox):
+def test_create_static_matplotlib_chart_success(mock_sandbox, mock_context):
     mock_execution = MagicMock()
     mock_execution.results = [MagicMock(png="base64encodedimage")]
     mock_execution.logs.to_json.return_value = '{"logs": "log data"}'
     mock_execution.error = None
     mock_sandbox.run_code.return_value = mock_execution
 
-    result = create_static_matplotlib_chart("import matplotlib.pyplot as plt")
+    result = create_static_matplotlib_chart(mock_context, "import matplotlib.pyplot as plt")
     assert result == {
         "base64_image": "base64encodedimage",
         "logs": '{"logs": "log data"}',
@@ -46,14 +53,14 @@ def test_create_static_matplotlib_chart_success(mock_sandbox):
     }
 
 
-def test_create_static_matplotlib_chart_error(mock_sandbox):
+def test_create_static_matplotlib_chart_error(mock_sandbox, mock_context):
     mock_execution = MagicMock()
     mock_execution.results = []
     mock_execution.logs.to_json.return_value = '{"logs": "log data"}'
     mock_execution.error.to_json.return_value = '{"error": "some error"}'
     mock_sandbox.run_code.return_value = mock_execution
 
-    result = create_static_matplotlib_chart("import matplotlib.pyplot as plt")
+    result = create_static_matplotlib_chart(mock_context, "import matplotlib.pyplot as plt")
     assert result == {
         "base64_image": None,
         "logs": '{"logs": "log data"}',

toolkits/search/arcade_search/tools/google.py

@@ -2,21 +2,18 @@ import json
 from typing import Annotated
 
 import serpapi
-from arcade.sdk import tool
+from arcade.sdk import ToolContext, tool
-
-from arcade_search.tools.utils import get_secret
 
 
-@tool
+@tool(requires_secrets=["SERP_API_KEY"])
 async def search_google(
+    context: ToolContext,
     query: Annotated[str, "Search query"],
     n_results: Annotated[int, "Number of results to retrieve"] = 5,
 ) -> str:
     """Search Google using SerpAPI and return organic search results."""
 
-    api_key = get_secret("SERP_API_KEY")
+    api_key = context.get_secret("SERP_API_KEY")
-    if not api_key:
-        raise ValueError("SERP_API_KEY is not set")
 
     client = serpapi.Client(api_key=api_key)
     params = {"engine": "google", "q": query}

toolkits/search/arcade_search/tools/utils.py

@@ -1,11 +0,0 @@
-import os
-from typing import Any, Optional
-
-
-def get_secret(name: str, default: Optional[Any] = None) -> Any:
-    secret = os.getenv(name)
-    if secret is None:
-        if default is not None:
-            return default
-        raise ValueError(f"Secret {name} is not set.")
-    return secret

toolkits/search/pyproject.toml

@@ -1,12 +1,12 @@
 [tool.poetry]
 name = "arcade_search"
-version = "0.1.10"
+version = "1.0.0"
 description = "Tools for searching the web"
 authors = ["Arcade <dev@arcade.dev>"]
 
 [tool.poetry.dependencies]
 python = "^3.10"
-arcade-ai = ">=0.1,<2.0"
+arcade-ai = ">=1.0.5,<2.0"
 serpapi = "^0.1.5"
 
 [tool.poetry.dev-dependencies]

toolkits/search/tests/test_google.py

@@ -2,16 +2,20 @@ import json
 from unittest.mock import patch
 
 import pytest
+from arcade.core.schema import ToolSecretItem
+from arcade.sdk import ToolContext
 
 from arcade_search.tools.google import search_google
 
-GET_SECRET_PATCH_TARGET = "arcade_search.tools.google.get_secret"  # noqa: S105
+
+@pytest.fixture
+def mock_context():
+    return ToolContext(secrets=[ToolSecretItem(key="serp_api_key", value="fake_api_key")])
 
 
 @pytest.mark.asyncio
-async def test_search_google_success():
+async def test_search_google_success(mock_context):
     with (
-        patch(GET_SECRET_PATCH_TARGET, return_value="fake_api_key"),
         patch("serpapi.Client") as MockClient,
     ):
         mock_client_instance = MockClient.return_value
@@ -23,7 +27,7 @@ async def test_search_google_success():
             ]
         }
 
-        result = await search_google("test query", 2)
+        result = await search_google(mock_context, "test query", 2)
 
         expected_result = json.dumps([
             {"title": "Result 1", "link": "http://example.com/1"},
@@ -33,15 +37,14 @@ async def test_search_google_success():
 
 
 @pytest.mark.asyncio
-async def test_search_google_no_results():
+async def test_search_google_no_results(mock_context):
     with (
-        patch(GET_SECRET_PATCH_TARGET, return_value="fake_api_key"),
         patch("serpapi.Client") as MockClient,
     ):
         mock_client_instance = MockClient.return_value
         mock_client_instance.search.return_value.as_dict.return_value = {"organic_results": []}
 
-        result = await search_google("test query", 2)
+        result = await search_google(mock_context, "test query", 2)
 
         expected_result = json.dumps([])
         assert result == expected_result

toolkits/web/arcade_web/tools/firecrawl.py

@@ -1,17 +1,17 @@
 from typing import Annotated, Any, Optional
 
-from arcade.sdk import tool
+from arcade.sdk import ToolContext, tool
 from firecrawl import FirecrawlApp
 
 from arcade_web.tools.models import Formats
-from arcade_web.tools.utils import get_secret
 
 
 # TODO: Support actions. This would enable clicking, scrolling, screenshotting, etc.
 # TODO: Support extract.
 # TODO: Support headers param?
-@tool
+@tool(requires_secrets=["FIRECRAWL_API_KEY"])
 async def scrape_url(
+    context: ToolContext,
     url: Annotated[str, "URL to scrape"],
     formats: Annotated[
         Optional[list[Formats]], "Formats to retrieve. Defaults to ['markdown']."
@@ -31,7 +31,7 @@ async def scrape_url(
 ) -> Annotated[dict[str, Any], "Scraped data in specified formats"]:
     """Scrape a URL using Firecrawl and return the data in specified formats."""
 
-    api_key = get_secret("FIRECRAWL_API_KEY")
+    api_key = context.get_secret("FIRECRAWL_API_KEY")
 
     formats = formats or [Formats.MARKDOWN]
 
@@ -50,8 +50,9 @@ async def scrape_url(
 
 
 # TODO: Support scrapeOptions.
-@tool
+@tool(requires_secrets=["FIRECRAWL_API_KEY"])
 async def crawl_website(
+    context: ToolContext,
     url: Annotated[str, "URL to crawl"],
     exclude_paths: Annotated[list[str] | None, "URL patterns to exclude from the crawl"] = None,
     include_paths: Annotated[list[str] | None, "URL patterns to include in the crawl"] = None,
@@ -75,7 +76,7 @@ async def crawl_website(
     If the crawl is synchronous, then returns the crawl data.
     """
 
-    api_key = get_secret("FIRECRAWL_API_KEY")
+    api_key = context.get_secret("FIRECRAWL_API_KEY")
 
     app = FirecrawlApp(api_key=api_key)
     params = {
@@ -102,15 +103,16 @@ async def crawl_website(
     return dict(response)
 
 
-@tool
+@tool(requires_secrets=["FIRECRAWL_API_KEY"])
 async def get_crawl_status(
+    context: ToolContext,
     crawl_id: Annotated[str, "The ID of the crawl job"],
 ) -> Annotated[dict[str, Any], "Crawl status information"]:
     """
     Get the status of a Firecrawl 'crawl' that is either in progress or recently completed.
     """
 
-    api_key = get_secret("FIRECRAWL_API_KEY")
+    api_key = context.get_secret("FIRECRAWL_API_KEY")
 
     app = FirecrawlApp(api_key=api_key)
     crawl_status = app.check_crawl_status(crawl_id)
@@ -125,13 +127,14 @@ async def get_crawl_status(
 #       then the Firecrawl API response will have a next_url field.
 @tool
 async def get_crawl_data(
+    context: ToolContext,
     crawl_id: Annotated[str, "The ID of the crawl job"],
 ) -> Annotated[dict[str, Any], "Crawl data information"]:
     """
     Get the data of a Firecrawl 'crawl' that is either in progress or recently completed.
     """
 
-    api_key = get_secret("FIRECRAWL_API_KEY")
+    api_key = context.get_secret("FIRECRAWL_API_KEY")
 
     app = FirecrawlApp(api_key=api_key)
     crawl_data = app.check_crawl_status(crawl_id)
@@ -139,15 +142,16 @@ async def get_crawl_data(
     return dict(crawl_data)
 
 
-@tool
+@tool(requires_secrets=["FIRECRAWL_API_KEY"])
 async def cancel_crawl(
+    context: ToolContext,
     crawl_id: Annotated[str, "The ID of the asynchronous crawl job to cancel"],
 ) -> Annotated[dict[str, Any], "Cancellation status information"]:
     """
     Cancel an asynchronous crawl job that is in progress using the Firecrawl API.
     """
 
-    api_key = get_secret("FIRECRAWL_API_KEY")
+    api_key = context.get_secret("FIRECRAWL_API_KEY")
 
     app = FirecrawlApp(api_key=api_key)
     cancellation_status = app.cancel_crawl(crawl_id)
@@ -155,8 +159,9 @@ async def cancel_crawl(
     return dict(cancellation_status)
 
 
-@tool
+@tool(requires_secrets=["FIRECRAWL_API_KEY"])
 async def map_website(
+    context: ToolContext,
     url: Annotated[str, "The base URL to start crawling from"],
     search: Annotated[Optional[str], "Search query to use for mapping"] = None,
     ignore_sitemap: Annotated[bool, "Ignore the website sitemap when crawling"] = True,
@@ -167,7 +172,7 @@ async def map_website(
     Map a website from a single URL to a map of the entire website.
     """
 
-    api_key = get_secret("FIRECRAWL_API_KEY")
+    api_key = context.get_secret("FIRECRAWL_API_KEY")
 
     app = FirecrawlApp(api_key=api_key)
     params: dict[str, Any] = {

toolkits/web/arcade_web/tools/utils.py

@@ -1,9 +0,0 @@
-import os
-from typing import Any, Optional
-
-
-def get_secret(name: str, default: Optional[Any] = None) -> Any:
-    secret = os.getenv(name)
-    if secret is None and default is not None:
-        return default
-    return secret

toolkits/web/pyproject.toml

@@ -1,12 +1,12 @@
 [tool.poetry]
 name = "arcade_web"
-version = "0.1.9"
+version = "1.0.0"
 description = "LLM tools for web-related tasks"
 authors = ["Arcade <dev@arcade.dev>"]
 
 [tool.poetry.dependencies]
 python = "^3.10"
-arcade-ai = ">=0.1,<2.0"
+arcade-ai = ">=1.0.5,<2.0"
 firecrawl-py = "^1.3.1"
 
 [tool.poetry.dev-dependencies]

toolkits/web/tests/test_firecrawl.py

@@ -1,6 +1,8 @@
-from unittest.mock import AsyncMock, patch
+from unittest.mock import patch
 
 import pytest
+from arcade.core.schema import ToolSecretItem
+from arcade.sdk import ToolContext
 from arcade.sdk.errors import ToolExecutionError
 
 from arcade_web.tools.firecrawl import (
@@ -15,9 +17,7 @@ from arcade_web.tools.firecrawl import (
 
 @pytest.fixture
 def mock_context():
-    context = AsyncMock()
+    return ToolContext(secrets=[ToolSecretItem(key="firecrawl_api_key", value="fake_api_key")])
-    context.authorization.token = "mock_token"  # noqa: S105
-    return context
 
 
 @pytest.fixture
@@ -27,50 +27,50 @@ def mock_firecrawl_app():
 
 
 @pytest.mark.asyncio
-async def test_scrape_url_success(mock_firecrawl_app):
+async def test_scrape_url_success(mock_firecrawl_app, mock_context):
     mock_firecrawl_app.scrape_url.return_value = {"data": "scraped content"}
 
-    result = await scrape_url("http://example.com")
+    result = await scrape_url(mock_context, "http://example.com")
     assert result == {"data": "scraped content"}
 
 
 @pytest.mark.asyncio
-async def test_crawl_website_success(mock_firecrawl_app):
+async def test_crawl_website_success(mock_firecrawl_app, mock_context):
     mock_firecrawl_app.async_crawl_url.return_value = {"crawl_id": "12345"}
 
-    result = await crawl_website("http://example.com")
+    result = await crawl_website(mock_context, "http://example.com")
     assert result == {"crawl_id": "12345"}
 
 
 @pytest.mark.asyncio
-async def test_get_crawl_status_success(mock_firecrawl_app):
+async def test_get_crawl_status_success(mock_firecrawl_app, mock_context):
     mock_firecrawl_app.check_crawl_status.return_value = {"status": "completed"}
 
-    result = await get_crawl_status("12345")
+    result = await get_crawl_status(mock_context, "12345")
     assert result == {"status": "completed"}
 
 
 @pytest.mark.asyncio
-async def test_get_crawl_data_success(mock_firecrawl_app):
+async def test_get_crawl_data_success(mock_firecrawl_app, mock_context):
     mock_firecrawl_app.check_crawl_status.return_value = {"data": "crawl data"}
 
-    result = await get_crawl_data("12345")
+    result = await get_crawl_data(mock_context, "12345")
     assert result == {"data": "crawl data"}
 
 
 @pytest.mark.asyncio
-async def test_cancel_crawl_success(mock_firecrawl_app):
+async def test_cancel_crawl_success(mock_firecrawl_app, mock_context):
     mock_firecrawl_app.cancel_crawl.return_value = {"status": "cancelled"}
 
-    result = await cancel_crawl("12345")
+    result = await cancel_crawl(mock_context, "12345")
     assert result == {"status": "cancelled"}
 
 
 @pytest.mark.asyncio
-async def test_map_website_success(mock_firecrawl_app):
+async def test_map_website_success(mock_firecrawl_app, mock_context):
     mock_firecrawl_app.map_url.return_value = {"map": "website map"}
 
-    result = await map_website("http://example.com")
+    result = await map_website(mock_context, "http://example.com")
     assert result == {"map": "website map"}
 
 
@@ -86,7 +86,7 @@ async def test_map_website_success(mock_firecrawl_app):
         (map_website, ("http://example.com",), "Error mapping website"),
     ],
 )
-async def test_firecrawl_error(mock_firecrawl_app, method, params, error_message):
+async def test_firecrawl_error(mock_firecrawl_app, mock_context, method, params, error_message):
     mock_firecrawl_app.scrape_url.side_effect = Exception(error_message)
     mock_firecrawl_app.async_crawl_url.side_effect = Exception(error_message)
     mock_firecrawl_app.check_crawl_status.side_effect = Exception(error_message)
@@ -94,4 +94,4 @@ async def test_firecrawl_error(mock_firecrawl_app, method, params, error_message
     mock_firecrawl_app.map_url.side_effect = Exception(error_message)
 
     with pytest.raises(ToolExecutionError):
-        await method(*params)
+        await method(mock_context, *params)