diff --git a/src/crewai_tools/tools/exa_tools/exa_search_tool.py b/src/crewai_tools/tools/exa_tools/exa_search_tool.py index 332576039..5dbd1df0b 100644 --- a/src/crewai_tools/tools/exa_tools/exa_search_tool.py +++ b/src/crewai_tools/tools/exa_tools/exa_search_tool.py @@ -43,10 +43,18 @@ class EXASearchTool(BaseTool): description="API key for Exa services", json_schema_extra={"required": False}, ) + base_url: Optional[str] = Field( + default_factory=lambda: os.getenv("EXA_BASE_URL"), + description="API server url", + json_schema_extra={"required": False}, + ) env_vars: List[EnvVar] = [ EnvVar( name="EXA_API_KEY", description="API key for Exa services", required=False ), + EnvVar( + name="EXA_BASE_URL", description="API url for the Exa services", required=False + ), ] def __init__( @@ -73,7 +81,10 @@ class EXASearchTool(BaseTool): raise ImportError( "You are missing the 'exa_py' package. Would you like to install it?" ) - self.client = Exa(api_key=self.api_key) + client_kwargs = {"api_key": self.api_key} + if self.base_url: + client_kwargs["base_url"] = self.base_url + self.client = Exa(**client_kwargs) self.content = content self.summary = summary self.type = type diff --git a/tests/tools/exa_search_tool_test.py b/tests/tools/exa_search_tool_test.py index 17c92e2f4..02e9b4cc2 100644 --- a/tests/tools/exa_search_tool_test.py +++ b/tests/tools/exa_search_tool_test.py @@ -15,18 +15,58 @@ def mock_exa_api_key(): yield def test_exa_search_tool_initialization(): - with patch("crewai_tools.tools.exa_tools.exa_search_tool.Exa") as mock_exa_class: - api_key = "test_api_key" - tool = EXASearchTool(api_key=api_key) + with patch.dict(os.environ, {}, clear=True): + with patch("crewai_tools.tools.exa_tools.exa_search_tool.Exa") as mock_exa_class: + api_key = "test_api_key" + tool = EXASearchTool(api_key=api_key) - assert tool.api_key == api_key - assert tool.content is False - assert tool.summary is False - assert tool.type == "auto" - mock_exa_class.assert_called_once_with(api_key=api_key) + assert tool.api_key == api_key + assert tool.content is False + assert tool.summary is False + assert tool.type == "auto" + mock_exa_class.assert_called_once_with(api_key=api_key) def test_exa_search_tool_initialization_with_env(mock_exa_api_key): + with patch.dict(os.environ, {"EXA_API_KEY": "test_key_from_env"}, clear=True): + with patch("crewai_tools.tools.exa_tools.exa_search_tool.Exa") as mock_exa_class: + EXASearchTool() + mock_exa_class.assert_called_once_with(api_key="test_key_from_env") + + +def test_exa_search_tool_initialization_with_base_url(): + with patch.dict(os.environ, {}, clear=True): + with patch("crewai_tools.tools.exa_tools.exa_search_tool.Exa") as mock_exa_class: + api_key = "test_api_key" + base_url = "https://custom.exa.api.com" + tool = EXASearchTool(api_key=api_key, base_url=base_url) + + assert tool.api_key == api_key + assert tool.base_url == base_url + assert tool.content is False + assert tool.summary is False + assert tool.type == "auto" + mock_exa_class.assert_called_once_with(api_key=api_key, base_url=base_url) + + +@pytest.fixture +def mock_exa_base_url(): + with patch.dict(os.environ, {"EXA_BASE_URL": "https://env.exa.api.com"}): + yield + + +def test_exa_search_tool_initialization_with_env_base_url(mock_exa_api_key, mock_exa_base_url): with patch("crewai_tools.tools.exa_tools.exa_search_tool.Exa") as mock_exa_class: EXASearchTool() - mock_exa_class.assert_called_once_with(api_key="test_key_from_env") + mock_exa_class.assert_called_once_with(api_key="test_key_from_env", base_url="https://env.exa.api.com") + + +def test_exa_search_tool_initialization_without_base_url(): + with patch.dict(os.environ, {}, clear=True): + with patch("crewai_tools.tools.exa_tools.exa_search_tool.Exa") as mock_exa_class: + api_key = "test_api_key" + tool = EXASearchTool(api_key=api_key) + + assert tool.api_key == api_key + assert tool.base_url is None + mock_exa_class.assert_called_once_with(api_key=api_key)