diff --git a/api/app/clients/tools/structured/GoogleSearch.js b/api/app/clients/tools/structured/GoogleSearch.js index bae1a458e0..c2cf72b9d6 100644 --- a/api/app/clients/tools/structured/GoogleSearch.js +++ b/api/app/clients/tools/structured/GoogleSearch.js @@ -12,9 +12,15 @@ class GoogleSearchResults extends Tool { this.envVarApiKey = 'GOOGLE_SEARCH_API_KEY'; this.envVarSearchEngineId = 'GOOGLE_CSE_ID'; this.override = fields.override ?? false; - this.apiKey = fields.apiKey ?? getEnvironmentVariable(this.envVarApiKey); + this.apiKey = fields[this.envVarApiKey] ?? getEnvironmentVariable(this.envVarApiKey); this.searchEngineId = - fields.searchEngineId ?? getEnvironmentVariable(this.envVarSearchEngineId); + fields[this.envVarSearchEngineId] ?? getEnvironmentVariable(this.envVarSearchEngineId); + + if (!this.override && (!this.apiKey || !this.searchEngineId)) { + throw new Error( + `Missing ${this.envVarApiKey} or ${this.envVarSearchEngineId} environment variable.`, + ); + } this.kwargs = fields?.kwargs ?? {}; this.name = 'google'; diff --git a/api/app/clients/tools/structured/TavilySearchResults.js b/api/app/clients/tools/structured/TavilySearchResults.js index b1691a63db..9a62053ff0 100644 --- a/api/app/clients/tools/structured/TavilySearchResults.js +++ b/api/app/clients/tools/structured/TavilySearchResults.js @@ -12,7 +12,7 @@ class TavilySearchResults extends Tool { this.envVar = 'TAVILY_API_KEY'; /* Used to initialize the Tool without necessary variables. */ this.override = fields.override ?? false; - this.apiKey = fields.apiKey ?? this.getApiKey(); + this.apiKey = fields[this.envVar] ?? this.getApiKey(); this.kwargs = fields?.kwargs ?? {}; this.name = 'tavily_search_results_json'; diff --git a/api/app/clients/tools/structured/specs/GoogleSearch.spec.js b/api/app/clients/tools/structured/specs/GoogleSearch.spec.js new file mode 100644 index 0000000000..ff11265301 --- /dev/null +++ b/api/app/clients/tools/structured/specs/GoogleSearch.spec.js @@ -0,0 +1,50 @@ +const GoogleSearch = require('../GoogleSearch'); + +jest.mock('node-fetch'); +jest.mock('@langchain/core/utils/env'); + +describe('GoogleSearch', () => { + let originalEnv; + const mockApiKey = 'mock_api'; + const mockSearchEngineId = 'mock_search_engine_id'; + + beforeAll(() => { + originalEnv = { ...process.env }; + }); + + beforeEach(() => { + jest.resetModules(); + process.env = { + ...originalEnv, + GOOGLE_SEARCH_API_KEY: mockApiKey, + GOOGLE_CSE_ID: mockSearchEngineId, + }; + }); + + afterEach(() => { + jest.clearAllMocks(); + process.env = originalEnv; + }); + + it('should use mockApiKey and mockSearchEngineId when environment variables are not set', () => { + const instance = new GoogleSearch({ + GOOGLE_SEARCH_API_KEY: mockApiKey, + GOOGLE_CSE_ID: mockSearchEngineId, + }); + expect(instance.apiKey).toBe(mockApiKey); + expect(instance.searchEngineId).toBe(mockSearchEngineId); + }); + + it('should throw an error if GOOGLE_SEARCH_API_KEY or GOOGLE_CSE_ID is missing', () => { + delete process.env.GOOGLE_SEARCH_API_KEY; + expect(() => new GoogleSearch()).toThrow( + 'Missing GOOGLE_SEARCH_API_KEY or GOOGLE_CSE_ID environment variable.', + ); + + process.env.GOOGLE_SEARCH_API_KEY = mockApiKey; + delete process.env.GOOGLE_CSE_ID; + expect(() => new GoogleSearch()).toThrow( + 'Missing GOOGLE_SEARCH_API_KEY or GOOGLE_CSE_ID environment variable.', + ); + }); +}); diff --git a/api/app/clients/tools/structured/specs/TavilySearchResults.spec.js b/api/app/clients/tools/structured/specs/TavilySearchResults.spec.js new file mode 100644 index 0000000000..5ea00140c7 --- /dev/null +++ b/api/app/clients/tools/structured/specs/TavilySearchResults.spec.js @@ -0,0 +1,38 @@ +const TavilySearchResults = require('../TavilySearchResults'); + +jest.mock('node-fetch'); +jest.mock('@langchain/core/utils/env'); + +describe('TavilySearchResults', () => { + let originalEnv; + const mockApiKey = 'mock_api_key'; + + beforeAll(() => { + originalEnv = { ...process.env }; + }); + + beforeEach(() => { + jest.resetModules(); + process.env = { + ...originalEnv, + TAVILY_API_KEY: mockApiKey, + }; + }); + + afterEach(() => { + jest.clearAllMocks(); + process.env = originalEnv; + }); + + it('should throw an error if TAVILY_API_KEY is missing', () => { + delete process.env.TAVILY_API_KEY; + expect(() => new TavilySearchResults()).toThrow('Missing TAVILY_API_KEY environment variable.'); + }); + + it('should use mockApiKey when TAVILY_API_KEY is not set in the environment', () => { + const instance = new TavilySearchResults({ + TAVILY_API_KEY: mockApiKey, + }); + expect(instance.apiKey).toBe(mockApiKey); + }); +});