🔧 fix: API Key Handling for GoogleSearch and TavilySearch Tools (#3541)

* fix: get user-provided API key using environment variable names as keys

* feat: Add error handling for missing API key and search engine ID

* feat: Add GoogleSearch and TavilySearchResults specs for environment variables handling

---------

Co-authored-by: Dongwoo Jeong <dongwoo.jeong@lge.com>
This commit is contained in:
Dongwoo Jeong 2024-08-06 08:29:16 +09:00 committed by GitHub
parent 80773d0bce
commit 6879de0bf1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 97 additions and 3 deletions

View file

@ -12,9 +12,15 @@ class GoogleSearchResults extends Tool {
this.envVarApiKey = 'GOOGLE_SEARCH_API_KEY'; this.envVarApiKey = 'GOOGLE_SEARCH_API_KEY';
this.envVarSearchEngineId = 'GOOGLE_CSE_ID'; this.envVarSearchEngineId = 'GOOGLE_CSE_ID';
this.override = fields.override ?? false; this.override = fields.override ?? false;
this.apiKey = fields.apiKey ?? getEnvironmentVariable(this.envVarApiKey); this.apiKey = fields[this.envVarApiKey] ?? getEnvironmentVariable(this.envVarApiKey);
this.searchEngineId = 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.kwargs = fields?.kwargs ?? {};
this.name = 'google'; this.name = 'google';

View file

@ -12,7 +12,7 @@ class TavilySearchResults extends Tool {
this.envVar = 'TAVILY_API_KEY'; this.envVar = 'TAVILY_API_KEY';
/* Used to initialize the Tool without necessary variables. */ /* Used to initialize the Tool without necessary variables. */
this.override = fields.override ?? false; this.override = fields.override ?? false;
this.apiKey = fields.apiKey ?? this.getApiKey(); this.apiKey = fields[this.envVar] ?? this.getApiKey();
this.kwargs = fields?.kwargs ?? {}; this.kwargs = fields?.kwargs ?? {};
this.name = 'tavily_search_results_json'; this.name = 'tavily_search_results_json';

View file

@ -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.',
);
});
});

View file

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