From e9ee3bba40e293b3930c71e7c29e359c569053f6 Mon Sep 17 00:00:00 2001 From: Nate Barbettini Date: Tue, 4 Mar 2025 13:35:36 -0800 Subject: [PATCH] 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 --- .../arcade_code_sandbox/tools/e2b.py | 13 +++---- .../arcade_code_sandbox/tools/utils.py | 9 ----- toolkits/code_sandbox/pyproject.toml | 4 +-- toolkits/code_sandbox/tests/test_e2b.py | 23 +++++++----- toolkits/search/arcade_search/tools/google.py | 11 +++--- toolkits/search/arcade_search/tools/utils.py | 11 ------ toolkits/search/pyproject.toml | 4 +-- toolkits/search/tests/test_google.py | 17 +++++---- toolkits/web/arcade_web/tools/firecrawl.py | 31 +++++++++------- toolkits/web/arcade_web/tools/utils.py | 9 ----- toolkits/web/pyproject.toml | 4 +-- toolkits/web/tests/test_firecrawl.py | 36 +++++++++---------- 12 files changed, 78 insertions(+), 94 deletions(-) delete mode 100644 toolkits/code_sandbox/arcade_code_sandbox/tools/utils.py delete mode 100644 toolkits/search/arcade_search/tools/utils.py delete mode 100644 toolkits/web/arcade_web/tools/utils.py diff --git a/toolkits/code_sandbox/arcade_code_sandbox/tools/e2b.py b/toolkits/code_sandbox/arcade_code_sandbox/tools/e2b.py index 51ba6fe2..26432b7d 100644 --- a/toolkits/code_sandbox/arcade_code_sandbox/tools/e2b.py +++ b/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) diff --git a/toolkits/code_sandbox/arcade_code_sandbox/tools/utils.py b/toolkits/code_sandbox/arcade_code_sandbox/tools/utils.py deleted file mode 100644 index c03b10da..00000000 --- a/toolkits/code_sandbox/arcade_code_sandbox/tools/utils.py +++ /dev/null @@ -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 diff --git a/toolkits/code_sandbox/pyproject.toml b/toolkits/code_sandbox/pyproject.toml index ade5c7ae..1c12516c 100644 --- a/toolkits/code_sandbox/pyproject.toml +++ b/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 "] [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] diff --git a/toolkits/code_sandbox/tests/test_e2b.py b/toolkits/code_sandbox/tests/test_e2b.py index 2f5c2a76..8af6d49c 100644 --- a/toolkits/code_sandbox/tests/test_e2b.py +++ b/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"}', diff --git a/toolkits/search/arcade_search/tools/google.py b/toolkits/search/arcade_search/tools/google.py index bf4821d2..6ff35207 100644 --- a/toolkits/search/arcade_search/tools/google.py +++ b/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_search.tools.utils import get_secret +from arcade.sdk import ToolContext, tool -@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") - if not api_key: - raise ValueError("SERP_API_KEY is not set") + api_key = context.get_secret("SERP_API_KEY") client = serpapi.Client(api_key=api_key) params = {"engine": "google", "q": query} diff --git a/toolkits/search/arcade_search/tools/utils.py b/toolkits/search/arcade_search/tools/utils.py deleted file mode 100644 index 94394baf..00000000 --- a/toolkits/search/arcade_search/tools/utils.py +++ /dev/null @@ -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 diff --git a/toolkits/search/pyproject.toml b/toolkits/search/pyproject.toml index b678317e..0503c213 100644 --- a/toolkits/search/pyproject.toml +++ b/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 "] [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] diff --git a/toolkits/search/tests/test_google.py b/toolkits/search/tests/test_google.py index 1734a1cb..619be7dd 100644 --- a/toolkits/search/tests/test_google.py +++ b/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 diff --git a/toolkits/web/arcade_web/tools/firecrawl.py b/toolkits/web/arcade_web/tools/firecrawl.py index fc245bd6..b7efa08b 100644 --- a/toolkits/web/arcade_web/tools/firecrawl.py +++ b/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] = { diff --git a/toolkits/web/arcade_web/tools/utils.py b/toolkits/web/arcade_web/tools/utils.py deleted file mode 100644 index c03b10da..00000000 --- a/toolkits/web/arcade_web/tools/utils.py +++ /dev/null @@ -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 diff --git a/toolkits/web/pyproject.toml b/toolkits/web/pyproject.toml index 1c85e8b7..a69c172f 100644 --- a/toolkits/web/pyproject.toml +++ b/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 "] [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] diff --git a/toolkits/web/tests/test_firecrawl.py b/toolkits/web/tests/test_firecrawl.py index e31d7302..1d5e24b0 100644 --- a/toolkits/web/tests/test_firecrawl.py +++ b/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() - context.authorization.token = "mock_token" # noqa: S105 - return context + return ToolContext(secrets=[ToolSecretItem(key="firecrawl_api_key", value="fake_api_key")]) @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)