fix(ai-services): add logic for API key checking in fallback sequence
This commit is contained in:
Eyal Toledano
@@ -10,6 +10,7 @@ const mockGetFallbackModelId = jest.fn();
|
||||
const mockGetParametersForRole = jest.fn();
|
||||
const mockGetUserId = jest.fn();
|
||||
const mockGetDebugFlag = jest.fn();
|
||||
const mockIsApiKeySet = jest.fn();
|
||||
|
||||
// --- Mock MODEL_MAP Data ---
|
||||
// Provide a simplified structure sufficient for cost calculation tests
|
||||
@@ -29,6 +30,12 @@ const mockModelMap = {
|
||||
id: 'test-research-model',
|
||||
cost_per_1m_tokens: { input: 1, output: 1, currency: 'USD' }
|
||||
}
|
||||
],
|
||||
openai: [
|
||||
{
|
||||
id: 'test-openai-model',
|
||||
cost_per_1m_tokens: { input: 2, output: 6, currency: 'USD' }
|
||||
}
|
||||
]
|
||||
// Add other providers/models if needed for specific tests
|
||||
};
|
||||
@@ -45,7 +52,8 @@ jest.unstable_mockModule('../../scripts/modules/config-manager.js', () => ({
|
||||
getUserId: mockGetUserId,
|
||||
getDebugFlag: mockGetDebugFlag,
|
||||
MODEL_MAP: mockModelMap,
|
||||
getBaseUrlForRole: mockGetBaseUrlForRole
|
||||
getBaseUrlForRole: mockGetBaseUrlForRole,
|
||||
isApiKeySet: mockIsApiKeySet
|
||||
}));
|
||||
|
||||
// Mock AI Provider Modules
|
||||
@@ -67,7 +75,24 @@ jest.unstable_mockModule('../../src/ai-providers/perplexity.js', () => ({
|
||||
generatePerplexityObject: mockGeneratePerplexityObject
|
||||
}));
|
||||
|
||||
// ... Mock other providers (google, openai, etc.) similarly ...
|
||||
const mockGenerateOpenAIText = jest.fn();
|
||||
const mockStreamOpenAIText = jest.fn();
|
||||
const mockGenerateOpenAIObject = jest.fn();
|
||||
jest.unstable_mockModule('../../src/ai-providers/openai.js', () => ({
|
||||
generateOpenAIText: mockGenerateOpenAIText,
|
||||
streamOpenAIText: mockStreamOpenAIText,
|
||||
generateOpenAIObject: mockGenerateOpenAIObject
|
||||
}));
|
||||
|
||||
// Mock ollama provider (for special case testing - API key is optional)
|
||||
const mockGenerateOllamaText = jest.fn();
|
||||
const mockStreamOllamaText = jest.fn();
|
||||
const mockGenerateOllamaObject = jest.fn();
|
||||
jest.unstable_mockModule('../../src/ai-providers/ollama.js', () => ({
|
||||
generateOllamaText: mockGenerateOllamaText,
|
||||
streamOllamaText: mockStreamOllamaText,
|
||||
generateOllamaObject: mockGenerateOllamaObject
|
||||
}));
|
||||
|
||||
// Mock utils logger, API key resolver, AND findProjectRoot
|
||||
const mockLog = jest.fn();
|
||||
@@ -112,6 +137,8 @@ describe('Unified AI Services', () => {
|
||||
mockResolveEnvVariable.mockImplementation((key) => {
|
||||
if (key === 'ANTHROPIC_API_KEY') return 'mock-anthropic-key';
|
||||
if (key === 'PERPLEXITY_API_KEY') return 'mock-perplexity-key';
|
||||
if (key === 'OPENAI_API_KEY') return 'mock-openai-key';
|
||||
if (key === 'OLLAMA_API_KEY') return 'mock-ollama-key';
|
||||
return null;
|
||||
});
|
||||
|
||||
@@ -119,6 +146,7 @@ describe('Unified AI Services', () => {
|
||||
mockFindProjectRoot.mockReturnValue(fakeProjectRoot);
|
||||
mockGetDebugFlag.mockReturnValue(false);
|
||||
mockGetUserId.mockReturnValue('test-user-id'); // Add default mock for getUserId
|
||||
mockIsApiKeySet.mockReturnValue(true); // Default to true for most tests
|
||||
});
|
||||
|
||||
describe('generateTextService', () => {
|
||||
@@ -333,13 +361,218 @@ describe('Unified AI Services', () => {
|
||||
expect(mockGenerateAnthropicText).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
// Add more tests for edge cases:
|
||||
// - Missing API keys (should throw from _resolveApiKey)
|
||||
// - Unsupported provider configured (should skip and log)
|
||||
// - Missing provider/model config for a role (should skip and log)
|
||||
// - Missing prompt
|
||||
// - Different initial roles (research, fallback)
|
||||
// - generateObjectService (mock schema, check object result)
|
||||
// - streamTextService (more complex to test, might need stream helpers)
|
||||
// New tests for API key checking and fallback sequence
|
||||
// These tests verify that:
|
||||
// 1. The system checks if API keys are set before trying to use a provider
|
||||
// 2. If a provider's API key is missing, it skips to the next provider in the fallback sequence
|
||||
// 3. The system throws an appropriate error if all providers' API keys are missing
|
||||
// 4. Ollama is a special case where API key is optional and not checked
|
||||
// 5. Session context is correctly used for API key checks
|
||||
|
||||
test('should skip provider with missing API key and try next in fallback sequence', async () => {
|
||||
// Setup isApiKeySet to return false for anthropic but true for perplexity
|
||||
mockIsApiKeySet
|
||||
.mockImplementation((provider, session, root) => {
|
||||
if (provider === 'anthropic') return false; // Main provider has no key
|
||||
return true; // Other providers have keys
|
||||
});
|
||||
|
||||
// Mock perplexity text response (since we'll skip anthropic)
|
||||
mockGeneratePerplexityText.mockResolvedValue({
|
||||
text: 'Perplexity response (skipped to research)',
|
||||
usage: { inputTokens: 20, outputTokens: 30, totalTokens: 50 }
|
||||
});
|
||||
|
||||
const params = {
|
||||
role: 'main',
|
||||
prompt: 'Skip main provider test',
|
||||
session: { env: {} }
|
||||
};
|
||||
|
||||
const result = await generateTextService(params);
|
||||
|
||||
// Should have gotten the perplexity response
|
||||
expect(result.mainResult).toBe('Perplexity response (skipped to research)');
|
||||
|
||||
// Should check API keys
|
||||
expect(mockIsApiKeySet).toHaveBeenCalledWith('anthropic', params.session, fakeProjectRoot);
|
||||
expect(mockIsApiKeySet).toHaveBeenCalledWith('perplexity', params.session, fakeProjectRoot);
|
||||
|
||||
// Should log a warning
|
||||
expect(mockLog).toHaveBeenCalledWith(
|
||||
'warn',
|
||||
expect.stringContaining(`Skipping role 'main' (Provider: anthropic): API key not set or invalid.`)
|
||||
);
|
||||
|
||||
// Should NOT call anthropic provider
|
||||
expect(mockGenerateAnthropicText).not.toHaveBeenCalled();
|
||||
|
||||
// Should call perplexity provider
|
||||
expect(mockGeneratePerplexityText).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
test('should skip multiple providers with missing API keys and use first available', async () => {
|
||||
// Setup: Main and fallback providers have no keys, only research has a key
|
||||
mockIsApiKeySet
|
||||
.mockImplementation((provider, session, root) => {
|
||||
if (provider === 'anthropic') return false; // Main and fallback are both anthropic
|
||||
if (provider === 'perplexity') return true; // Research has a key
|
||||
return false;
|
||||
});
|
||||
|
||||
// Define different providers for testing multiple skips
|
||||
mockGetFallbackProvider.mockReturnValue('openai'); // Different from main
|
||||
mockGetFallbackModelId.mockReturnValue('test-openai-model');
|
||||
|
||||
// Mock isApiKeySet to return false for both main and fallback
|
||||
mockIsApiKeySet
|
||||
.mockImplementation((provider, session, root) => {
|
||||
if (provider === 'anthropic') return false; // Main provider has no key
|
||||
if (provider === 'openai') return false; // Fallback provider has no key
|
||||
return true; // Research provider has a key
|
||||
});
|
||||
|
||||
// Mock perplexity text response (since we'll skip to research)
|
||||
mockGeneratePerplexityText.mockResolvedValue({
|
||||
text: 'Research response after skipping main and fallback',
|
||||
usage: { inputTokens: 20, outputTokens: 30, totalTokens: 50 }
|
||||
});
|
||||
|
||||
const params = {
|
||||
role: 'main',
|
||||
prompt: 'Skip multiple providers test',
|
||||
session: { env: {} }
|
||||
};
|
||||
|
||||
const result = await generateTextService(params);
|
||||
|
||||
// Should have gotten the perplexity (research) response
|
||||
expect(result.mainResult).toBe('Research response after skipping main and fallback');
|
||||
|
||||
// Should check API keys for all three roles
|
||||
expect(mockIsApiKeySet).toHaveBeenCalledWith('anthropic', params.session, fakeProjectRoot);
|
||||
expect(mockIsApiKeySet).toHaveBeenCalledWith('openai', params.session, fakeProjectRoot);
|
||||
expect(mockIsApiKeySet).toHaveBeenCalledWith('perplexity', params.session, fakeProjectRoot);
|
||||
|
||||
// Should log warnings for both skipped providers
|
||||
expect(mockLog).toHaveBeenCalledWith(
|
||||
'warn',
|
||||
expect.stringContaining(`Skipping role 'main' (Provider: anthropic): API key not set or invalid.`)
|
||||
);
|
||||
expect(mockLog).toHaveBeenCalledWith(
|
||||
'warn',
|
||||
expect.stringContaining(`Skipping role 'fallback' (Provider: openai): API key not set or invalid.`)
|
||||
);
|
||||
|
||||
// Should NOT call skipped providers
|
||||
expect(mockGenerateAnthropicText).not.toHaveBeenCalled();
|
||||
expect(mockGenerateOpenAIText).not.toHaveBeenCalled();
|
||||
|
||||
// Should call perplexity provider
|
||||
expect(mockGeneratePerplexityText).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
test('should throw error if all providers in sequence have missing API keys', async () => {
|
||||
// Mock all providers to have missing API keys
|
||||
mockIsApiKeySet.mockReturnValue(false);
|
||||
|
||||
const params = {
|
||||
role: 'main',
|
||||
prompt: 'All API keys missing test',
|
||||
session: { env: {} }
|
||||
};
|
||||
|
||||
// Should throw error since all providers would be skipped
|
||||
await expect(generateTextService(params)).rejects.toThrow(
|
||||
'AI service call failed for all configured roles'
|
||||
);
|
||||
|
||||
// Should log warnings for all skipped providers
|
||||
expect(mockLog).toHaveBeenCalledWith(
|
||||
'warn',
|
||||
expect.stringContaining(`Skipping role 'main' (Provider: anthropic): API key not set or invalid.`)
|
||||
);
|
||||
expect(mockLog).toHaveBeenCalledWith(
|
||||
'warn',
|
||||
expect.stringContaining(`Skipping role 'fallback' (Provider: anthropic): API key not set or invalid.`)
|
||||
);
|
||||
expect(mockLog).toHaveBeenCalledWith(
|
||||
'warn',
|
||||
expect.stringContaining(`Skipping role 'research' (Provider: perplexity): API key not set or invalid.`)
|
||||
);
|
||||
|
||||
// Should log final error
|
||||
expect(mockLog).toHaveBeenCalledWith(
|
||||
'error',
|
||||
expect.stringContaining('All roles in the sequence [main, fallback, research] failed.')
|
||||
);
|
||||
|
||||
// Should NOT call any providers
|
||||
expect(mockGenerateAnthropicText).not.toHaveBeenCalled();
|
||||
expect(mockGeneratePerplexityText).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test('should not check API key for Ollama provider and try to use it', async () => {
|
||||
// Setup: Set main provider to ollama
|
||||
mockGetMainProvider.mockReturnValue('ollama');
|
||||
mockGetMainModelId.mockReturnValue('llama3');
|
||||
|
||||
// Mock Ollama text generation to succeed
|
||||
mockGenerateOllamaText.mockResolvedValue({
|
||||
text: 'Ollama response (no API key required)',
|
||||
usage: { inputTokens: 10, outputTokens: 10, totalTokens: 20 }
|
||||
});
|
||||
|
||||
const params = {
|
||||
role: 'main',
|
||||
prompt: 'Ollama special case test',
|
||||
session: { env: {} }
|
||||
};
|
||||
|
||||
const result = await generateTextService(params);
|
||||
|
||||
// Should have gotten the Ollama response
|
||||
expect(result.mainResult).toBe('Ollama response (no API key required)');
|
||||
|
||||
// isApiKeySet shouldn't be called for Ollama
|
||||
// Note: This is indirect - the code just doesn't check isApiKeySet for ollama
|
||||
// so we're verifying ollama provider was called despite isApiKeySet being mocked to false
|
||||
mockIsApiKeySet.mockReturnValue(false); // Should be ignored for Ollama
|
||||
|
||||
// Should call Ollama provider
|
||||
expect(mockGenerateOllamaText).toHaveBeenCalledTimes(1);
|
||||
});
|
||||
|
||||
test('should correctly use the provided session for API key check', async () => {
|
||||
// Mock custom session object with env vars
|
||||
const customSession = { env: { ANTHROPIC_API_KEY: 'session-api-key' } };
|
||||
|
||||
// Setup API key check to verify the session is passed correctly
|
||||
mockIsApiKeySet
|
||||
.mockImplementation((provider, session, root) => {
|
||||
// Only return true if the correct session was provided
|
||||
return session === customSession;
|
||||
});
|
||||
|
||||
// Mock the anthropic response
|
||||
mockGenerateAnthropicText.mockResolvedValue({
|
||||
text: 'Anthropic response with session key',
|
||||
usage: { inputTokens: 10, outputTokens: 10, totalTokens: 20 }
|
||||
});
|
||||
|
||||
const params = {
|
||||
role: 'main',
|
||||
prompt: 'Session API key test',
|
||||
session: customSession
|
||||
};
|
||||
|
||||
const result = await generateTextService(params);
|
||||
|
||||
// Should check API key with the custom session
|
||||
expect(mockIsApiKeySet).toHaveBeenCalledWith('anthropic', customSession, fakeProjectRoot);
|
||||
|
||||
// Should have gotten the anthropic response
|
||||
expect(result.mainResult).toBe('Anthropic response with session key');
|
||||
});
|
||||
});
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user