admin-valentin/arcade-mcp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

fix: Use tool secrets in toolkits (#271)

Browse source 
~~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>
This commit is contained in:
Nate Barbettini 2025-03-04 13:35:36 -08:00 committed by GitHub
parent 75da4bf8b0
commit e9ee3bba40
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
12 changed files with 78 additions and 94 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
toolkits/code_sandbox/arcade_code_sandbox/tools/e2b.py
View file

@ -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)

9
toolkits/code_sandbox/arcade_code_sandbox/tools/utils.py
View file

@ -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

4
toolkits/code_sandbox/pyproject.toml
View file

@ -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]

23
toolkits/code_sandbox/tests/test_e2b.py
View file

@ -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"}',

11
toolkits/search/arcade_search/tools/google.py
View file

@ -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}

11
toolkits/search/arcade_search/tools/utils.py
View file

@ -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

4
toolkits/search/pyproject.toml
View file

@ -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]

17
toolkits/search/tests/test_google.py
View file

@ -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

31
toolkits/web/arcade_web/tools/firecrawl.py
View file

@ -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] = {

9
toolkits/web/arcade_web/tools/utils.py
View file

@ -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

4
toolkits/web/pyproject.toml
View file

@ -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]

36
toolkits/web/tests/test_firecrawl.py
View file

@ -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)