feat(config): Add Fallback Model and Expanded Provider Support
Introduces a configurable fallback model and adds support for additional AI provider API keys in the environment setup.
- **Add Fallback Model Configuration (.taskmasterconfig):**
- Implemented a new section in .
- Configured as the default fallback model, enhancing resilience if the primary model fails.
- **Update Default Model Configuration (.taskmasterconfig):**
- Changed the default model to .
- Changed the default model to .
- **Add API Key Examples (assets/env.example):**
- Added example environment variables for:
- (for OpenAI/OpenRouter)
- (for Google Gemini)
- (for XAI Grok)
- Included format comments for clarity.
This commit is contained in:
Eyal Toledano
16
tests/fixtures/.taskmasterconfig
vendored
Normal file
16
tests/fixtures/.taskmasterconfig
vendored
Normal file
@@ -0,0 +1,16 @@
|
||||
{
|
||||
"models": {
|
||||
"main": {
|
||||
"provider": "openai",
|
||||
"modelId": "gpt-4o"
|
||||
},
|
||||
"research": {
|
||||
"provider": "perplexity",
|
||||
"modelId": "sonar-pro"
|
||||
},
|
||||
"fallback": {
|
||||
"provider": "anthropic",
|
||||
"modelId": "claude-3-haiku-20240307"
|
||||
}
|
||||
}
|
||||
}
|
||||
350
tests/integration/cli/commands.test.js
Normal file
350
tests/integration/cli/commands.test.js
Normal file
@@ -0,0 +1,350 @@
|
||||
import { jest } from '@jest/globals';
|
||||
|
||||
// --- Define mock functions ---
|
||||
const mockGetMainModelId = jest.fn().mockReturnValue('claude-3-opus');
|
||||
const mockGetResearchModelId = jest.fn().mockReturnValue('gpt-4-turbo');
|
||||
const mockGetFallbackModelId = jest.fn().mockReturnValue('claude-3-haiku');
|
||||
const mockSetMainModel = jest.fn().mockResolvedValue(true);
|
||||
const mockSetResearchModel = jest.fn().mockResolvedValue(true);
|
||||
const mockSetFallbackModel = jest.fn().mockResolvedValue(true);
|
||||
const mockGetAvailableModels = jest.fn().mockReturnValue([
|
||||
{ id: 'claude-3-opus', name: 'Claude 3 Opus', provider: 'anthropic' },
|
||||
{ id: 'gpt-4-turbo', name: 'GPT-4 Turbo', provider: 'openai' },
|
||||
{ id: 'claude-3-haiku', name: 'Claude 3 Haiku', provider: 'anthropic' },
|
||||
{ id: 'claude-3-sonnet', name: 'Claude 3 Sonnet', provider: 'anthropic' }
|
||||
]);
|
||||
|
||||
// Mock UI related functions
|
||||
const mockDisplayHelp = jest.fn();
|
||||
const mockDisplayBanner = jest.fn();
|
||||
const mockLog = jest.fn();
|
||||
const mockStartLoadingIndicator = jest.fn(() => ({ stop: jest.fn() }));
|
||||
const mockStopLoadingIndicator = jest.fn();
|
||||
|
||||
// --- Setup mocks using unstable_mockModule (recommended for ES modules) ---
|
||||
jest.unstable_mockModule('../../../scripts/modules/config-manager.js', () => ({
|
||||
getMainModelId: mockGetMainModelId,
|
||||
getResearchModelId: mockGetResearchModelId,
|
||||
getFallbackModelId: mockGetFallbackModelId,
|
||||
setMainModel: mockSetMainModel,
|
||||
setResearchModel: mockSetResearchModel,
|
||||
setFallbackModel: mockSetFallbackModel,
|
||||
getAvailableModels: mockGetAvailableModels,
|
||||
VALID_PROVIDERS: ['anthropic', 'openai']
|
||||
}));
|
||||
|
||||
jest.unstable_mockModule('../../../scripts/modules/ui.js', () => ({
|
||||
displayHelp: mockDisplayHelp,
|
||||
displayBanner: mockDisplayBanner,
|
||||
log: mockLog,
|
||||
startLoadingIndicator: mockStartLoadingIndicator,
|
||||
stopLoadingIndicator: mockStopLoadingIndicator
|
||||
}));
|
||||
|
||||
// --- Mock chalk for consistent output formatting ---
|
||||
const mockChalk = {
|
||||
red: jest.fn((text) => text),
|
||||
yellow: jest.fn((text) => text),
|
||||
blue: jest.fn((text) => text),
|
||||
green: jest.fn((text) => text),
|
||||
gray: jest.fn((text) => text),
|
||||
dim: jest.fn((text) => text),
|
||||
bold: {
|
||||
cyan: jest.fn((text) => text),
|
||||
white: jest.fn((text) => text),
|
||||
red: jest.fn((text) => text)
|
||||
},
|
||||
cyan: {
|
||||
bold: jest.fn((text) => text)
|
||||
},
|
||||
white: {
|
||||
bold: jest.fn((text) => text)
|
||||
}
|
||||
};
|
||||
// Default function for chalk itself
|
||||
mockChalk.default = jest.fn((text) => text);
|
||||
// Add the methods to the function itself for dual usage
|
||||
Object.keys(mockChalk).forEach((key) => {
|
||||
if (key !== 'default') mockChalk.default[key] = mockChalk[key];
|
||||
});
|
||||
|
||||
jest.unstable_mockModule('chalk', () => ({
|
||||
default: mockChalk.default
|
||||
}));
|
||||
|
||||
// --- Import modules (AFTER mock setup) ---
|
||||
let configManager, ui, chalk;
|
||||
|
||||
describe('CLI Models Command (Action Handler Test)', () => {
|
||||
// Setup dynamic imports before tests run
|
||||
beforeAll(async () => {
|
||||
configManager = await import('../../../scripts/modules/config-manager.js');
|
||||
ui = await import('../../../scripts/modules/ui.js');
|
||||
chalk = (await import('chalk')).default;
|
||||
});
|
||||
|
||||
// --- Replicate the action handler logic from commands.js ---
|
||||
async function modelsAction(options) {
|
||||
options = options || {}; // Ensure options object exists
|
||||
const availableModels = configManager.getAvailableModels();
|
||||
|
||||
const findProvider = (modelId) => {
|
||||
const modelInfo = availableModels.find((m) => m.id === modelId);
|
||||
return modelInfo?.provider;
|
||||
};
|
||||
|
||||
let modelSetAction = false;
|
||||
|
||||
try {
|
||||
if (options.setMain) {
|
||||
const modelId = options.setMain;
|
||||
if (typeof modelId !== 'string' || modelId.trim() === '') {
|
||||
console.error(
|
||||
chalk.red('Error: --set-main flag requires a valid model ID.')
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
const provider = findProvider(modelId);
|
||||
if (!provider) {
|
||||
console.error(
|
||||
chalk.red(
|
||||
`Error: Model ID "${modelId}" not found in available models.`
|
||||
)
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
if (await configManager.setMainModel(provider, modelId)) {
|
||||
console.log(
|
||||
chalk.green(`Main model set to: ${modelId} (Provider: ${provider})`)
|
||||
);
|
||||
modelSetAction = true;
|
||||
} else {
|
||||
console.error(chalk.red(`Failed to set main model.`));
|
||||
process.exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
if (options.setResearch) {
|
||||
const modelId = options.setResearch;
|
||||
if (typeof modelId !== 'string' || modelId.trim() === '') {
|
||||
console.error(
|
||||
chalk.red('Error: --set-research flag requires a valid model ID.')
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
const provider = findProvider(modelId);
|
||||
if (!provider) {
|
||||
console.error(
|
||||
chalk.red(
|
||||
`Error: Model ID "${modelId}" not found in available models.`
|
||||
)
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
if (await configManager.setResearchModel(provider, modelId)) {
|
||||
console.log(
|
||||
chalk.green(
|
||||
`Research model set to: ${modelId} (Provider: ${provider})`
|
||||
)
|
||||
);
|
||||
modelSetAction = true;
|
||||
} else {
|
||||
console.error(chalk.red(`Failed to set research model.`));
|
||||
process.exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
if (options.setFallback) {
|
||||
const modelId = options.setFallback;
|
||||
if (typeof modelId !== 'string' || modelId.trim() === '') {
|
||||
console.error(
|
||||
chalk.red('Error: --set-fallback flag requires a valid model ID.')
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
const provider = findProvider(modelId);
|
||||
if (!provider) {
|
||||
console.error(
|
||||
chalk.red(
|
||||
`Error: Model ID "${modelId}" not found in available models.`
|
||||
)
|
||||
);
|
||||
process.exit(1);
|
||||
}
|
||||
if (await configManager.setFallbackModel(provider, modelId)) {
|
||||
console.log(
|
||||
chalk.green(
|
||||
`Fallback model set to: ${modelId} (Provider: ${provider})`
|
||||
)
|
||||
);
|
||||
modelSetAction = true;
|
||||
} else {
|
||||
console.error(chalk.red(`Failed to set fallback model.`));
|
||||
process.exit(1);
|
||||
}
|
||||
}
|
||||
|
||||
if (!modelSetAction) {
|
||||
const currentMain = configManager.getMainModelId();
|
||||
const currentResearch = configManager.getResearchModelId();
|
||||
const currentFallback = configManager.getFallbackModelId();
|
||||
|
||||
if (!availableModels || availableModels.length === 0) {
|
||||
console.log(chalk.yellow('No models defined in configuration.'));
|
||||
return;
|
||||
}
|
||||
|
||||
// Create a mock table for testing - avoid using Table constructor
|
||||
const mockTableData = [];
|
||||
availableModels.forEach((model) => {
|
||||
if (model.id.startsWith('[') && model.id.endsWith(']')) return;
|
||||
mockTableData.push([
|
||||
model.id,
|
||||
model.name || 'N/A',
|
||||
model.provider || 'N/A',
|
||||
model.id === currentMain ? chalk.green(' ✓') : '',
|
||||
model.id === currentResearch ? chalk.green(' ✓') : '',
|
||||
model.id === currentFallback ? chalk.green(' ✓') : ''
|
||||
]);
|
||||
});
|
||||
|
||||
// In a real implementation, we would use cli-table3, but for testing
|
||||
// we'll just log 'Mock Table Output'
|
||||
console.log('Mock Table Output');
|
||||
}
|
||||
} catch (error) {
|
||||
// Use ui.log mock if available, otherwise console.error
|
||||
(ui.log || console.error)(
|
||||
`Error processing models command: ${error.message}`,
|
||||
'error'
|
||||
);
|
||||
if (error.stack) {
|
||||
(ui.log || console.error)(error.stack, 'debug');
|
||||
}
|
||||
throw error; // Re-throw for test failure
|
||||
}
|
||||
}
|
||||
// --- End of Action Handler Logic ---
|
||||
|
||||
let originalConsoleLog;
|
||||
let originalConsoleError;
|
||||
let originalProcessExit;
|
||||
|
||||
beforeEach(() => {
|
||||
// Reset all mocks
|
||||
jest.clearAllMocks();
|
||||
|
||||
// Save original console methods
|
||||
originalConsoleLog = console.log;
|
||||
originalConsoleError = console.error;
|
||||
originalProcessExit = process.exit;
|
||||
|
||||
// Mock console and process.exit
|
||||
console.log = jest.fn();
|
||||
console.error = jest.fn();
|
||||
process.exit = jest.fn((code) => {
|
||||
throw new Error(`process.exit(${code}) called`);
|
||||
});
|
||||
});
|
||||
|
||||
afterEach(() => {
|
||||
// Restore original console methods
|
||||
console.log = originalConsoleLog;
|
||||
console.error = originalConsoleError;
|
||||
process.exit = originalProcessExit;
|
||||
});
|
||||
|
||||
// --- Test Cases (Calling modelsAction directly) ---
|
||||
|
||||
it('should call setMainModel with correct provider and ID', async () => {
|
||||
const modelId = 'claude-3-opus';
|
||||
const expectedProvider = 'anthropic';
|
||||
await modelsAction({ setMain: modelId });
|
||||
expect(mockSetMainModel).toHaveBeenCalledWith(expectedProvider, modelId);
|
||||
expect(console.log).toHaveBeenCalledWith(
|
||||
expect.stringContaining(`Main model set to: ${modelId}`)
|
||||
);
|
||||
expect(console.log).toHaveBeenCalledWith(
|
||||
expect.stringContaining(`(Provider: ${expectedProvider})`)
|
||||
);
|
||||
});
|
||||
|
||||
it('should show an error if --set-main model ID is not found', async () => {
|
||||
await expect(
|
||||
modelsAction({ setMain: 'non-existent-model' })
|
||||
).rejects.toThrow(/process.exit/); // Expect exit call
|
||||
expect(mockSetMainModel).not.toHaveBeenCalled();
|
||||
expect(console.error).toHaveBeenCalledWith(
|
||||
expect.stringContaining('Model ID "non-existent-model" not found')
|
||||
);
|
||||
});
|
||||
|
||||
it('should call setResearchModel with correct provider and ID', async () => {
|
||||
const modelId = 'gpt-4-turbo';
|
||||
const expectedProvider = 'openai';
|
||||
await modelsAction({ setResearch: modelId });
|
||||
expect(mockSetResearchModel).toHaveBeenCalledWith(
|
||||
expectedProvider,
|
||||
modelId
|
||||
);
|
||||
expect(console.log).toHaveBeenCalledWith(
|
||||
expect.stringContaining(`Research model set to: ${modelId}`)
|
||||
);
|
||||
expect(console.log).toHaveBeenCalledWith(
|
||||
expect.stringContaining(`(Provider: ${expectedProvider})`)
|
||||
);
|
||||
});
|
||||
|
||||
it('should call setFallbackModel with correct provider and ID', async () => {
|
||||
const modelId = 'claude-3-haiku';
|
||||
const expectedProvider = 'anthropic';
|
||||
await modelsAction({ setFallback: modelId });
|
||||
expect(mockSetFallbackModel).toHaveBeenCalledWith(
|
||||
expectedProvider,
|
||||
modelId
|
||||
);
|
||||
expect(console.log).toHaveBeenCalledWith(
|
||||
expect.stringContaining(`Fallback model set to: ${modelId}`)
|
||||
);
|
||||
expect(console.log).toHaveBeenCalledWith(
|
||||
expect.stringContaining(`(Provider: ${expectedProvider})`)
|
||||
);
|
||||
});
|
||||
|
||||
it('should call all set*Model functions when all flags are used', async () => {
|
||||
const mainModelId = 'claude-3-opus';
|
||||
const researchModelId = 'gpt-4-turbo';
|
||||
const fallbackModelId = 'claude-3-haiku';
|
||||
const mainProvider = 'anthropic';
|
||||
const researchProvider = 'openai';
|
||||
const fallbackProvider = 'anthropic';
|
||||
|
||||
await modelsAction({
|
||||
setMain: mainModelId,
|
||||
setResearch: researchModelId,
|
||||
setFallback: fallbackModelId
|
||||
});
|
||||
expect(mockSetMainModel).toHaveBeenCalledWith(mainProvider, mainModelId);
|
||||
expect(mockSetResearchModel).toHaveBeenCalledWith(
|
||||
researchProvider,
|
||||
researchModelId
|
||||
);
|
||||
expect(mockSetFallbackModel).toHaveBeenCalledWith(
|
||||
fallbackProvider,
|
||||
fallbackModelId
|
||||
);
|
||||
});
|
||||
|
||||
it('should call specific get*ModelId and getAvailableModels and log table when run without flags', async () => {
|
||||
await modelsAction({}); // Call with empty options
|
||||
|
||||
expect(mockGetMainModelId).toHaveBeenCalled();
|
||||
expect(mockGetResearchModelId).toHaveBeenCalled();
|
||||
expect(mockGetFallbackModelId).toHaveBeenCalled();
|
||||
expect(mockGetAvailableModels).toHaveBeenCalled();
|
||||
|
||||
expect(console.log).toHaveBeenCalled();
|
||||
// Check the mocked Table.toString() was used via console.log
|
||||
expect(console.log).toHaveBeenCalledWith('Mock Table Output');
|
||||
});
|
||||
});
|
||||
@@ -25,9 +25,9 @@ global.wait = (ms) => new Promise((resolve) => setTimeout(resolve, ms));
|
||||
if (process.env.SILENCE_CONSOLE === 'true') {
|
||||
global.console = {
|
||||
...console,
|
||||
log: jest.fn(),
|
||||
info: jest.fn(),
|
||||
warn: jest.fn(),
|
||||
error: jest.fn()
|
||||
log: () => {},
|
||||
info: () => {},
|
||||
warn: () => {},
|
||||
error: () => {}
|
||||
};
|
||||
}
|
||||
|
||||
@@ -59,7 +59,8 @@ const DEFAULT_CONFIG = {
|
||||
const VALID_CUSTOM_CONFIG = {
|
||||
models: {
|
||||
main: { provider: 'openai', modelId: 'gpt-4o' },
|
||||
research: { provider: 'google', modelId: 'gemini-1.5-pro-latest' }
|
||||
research: { provider: 'google', modelId: 'gemini-1.5-pro-latest' },
|
||||
fallback: { provider: undefined, modelId: undefined }
|
||||
}
|
||||
};
|
||||
|
||||
@@ -67,6 +68,7 @@ const PARTIAL_CONFIG = {
|
||||
models: {
|
||||
main: { provider: 'openai', modelId: 'gpt-4-turbo' }
|
||||
// research missing
|
||||
// fallback will be added by readConfig
|
||||
}
|
||||
};
|
||||
|
||||
@@ -90,9 +92,66 @@ const resetMocks = () => {
|
||||
mockWriteFileSync.mockReset();
|
||||
mockMkdirSync.mockReset();
|
||||
|
||||
// Default behaviors
|
||||
mockExistsSync.mockReturnValue(true);
|
||||
mockReadFileSync.mockReturnValue(JSON.stringify(DEFAULT_CONFIG));
|
||||
// Default behaviors - CRITICAL: Mock supported-models.json read
|
||||
mockReadFileSync.mockImplementation((filePath) => {
|
||||
if (filePath.endsWith('supported-models.json')) {
|
||||
// Return a mock structure including allowed_roles
|
||||
return JSON.stringify({
|
||||
openai: [
|
||||
{
|
||||
id: 'gpt-4o',
|
||||
swe_score: 0,
|
||||
cost_per_1m_tokens: null,
|
||||
allowed_roles: ['main', 'fallback']
|
||||
},
|
||||
{
|
||||
id: 'gpt-4',
|
||||
swe_score: 0,
|
||||
cost_per_1m_tokens: null,
|
||||
allowed_roles: ['main', 'fallback']
|
||||
}
|
||||
],
|
||||
google: [
|
||||
{
|
||||
id: 'gemini-1.5-pro-latest',
|
||||
swe_score: 0,
|
||||
cost_per_1m_tokens: null,
|
||||
allowed_roles: ['main', 'fallback']
|
||||
}
|
||||
],
|
||||
perplexity: [
|
||||
{
|
||||
id: 'sonar-pro',
|
||||
swe_score: 0,
|
||||
cost_per_1m_tokens: null,
|
||||
allowed_roles: ['main', 'fallback', 'research']
|
||||
}
|
||||
],
|
||||
anthropic: [
|
||||
{
|
||||
id: 'claude-3-opus-20240229',
|
||||
swe_score: 0,
|
||||
cost_per_1m_tokens: null,
|
||||
allowed_roles: ['main', 'fallback']
|
||||
},
|
||||
{
|
||||
id: 'claude-3.5-sonnet-20240620',
|
||||
swe_score: 0,
|
||||
cost_per_1m_tokens: null,
|
||||
allowed_roles: ['main', 'fallback']
|
||||
}
|
||||
]
|
||||
// Add other providers/models as needed for specific tests
|
||||
});
|
||||
} else if (filePath === MOCK_CONFIG_PATH) {
|
||||
// Default for .taskmasterconfig reads
|
||||
return JSON.stringify(DEFAULT_CONFIG);
|
||||
}
|
||||
// Handle other potential reads or throw an error for unexpected paths
|
||||
throw new Error(`Unexpected readFileSync call in test: ${filePath}`);
|
||||
});
|
||||
|
||||
mockExistsSync.mockReturnValue(true); // Default to file existing
|
||||
};
|
||||
|
||||
// Set up module before tests
|
||||
@@ -253,10 +312,9 @@ describe('readConfig', () => {
|
||||
// --- writeConfig Tests ---
|
||||
describe('writeConfig', () => {
|
||||
test('should write valid config to file', () => {
|
||||
mockExistsSync.mockReturnValue(true);
|
||||
const success = configManager.writeConfig(
|
||||
VALID_CUSTOM_CONFIG,
|
||||
MOCK_PROJECT_ROOT
|
||||
MOCK_CONFIG_PATH
|
||||
);
|
||||
expect(success).toBe(true);
|
||||
expect(mockExistsSync).toHaveBeenCalledWith(MOCK_CONFIG_PATH);
|
||||
@@ -265,34 +323,29 @@ describe('writeConfig', () => {
|
||||
JSON.stringify(VALID_CUSTOM_CONFIG, null, 2),
|
||||
'utf-8'
|
||||
);
|
||||
expect(console.error).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
test('should return false and log error if write fails', () => {
|
||||
mockExistsSync.mockReturnValue(true);
|
||||
const writeError = new Error('Disk full');
|
||||
mockWriteFileSync.mockImplementation(() => {
|
||||
throw writeError;
|
||||
throw new Error('Disk full');
|
||||
});
|
||||
|
||||
const success = configManager.writeConfig(
|
||||
VALID_CUSTOM_CONFIG,
|
||||
MOCK_PROJECT_ROOT
|
||||
MOCK_CONFIG_PATH
|
||||
);
|
||||
|
||||
expect(success).toBe(false);
|
||||
expect(console.error).toHaveBeenCalledWith(
|
||||
expect.stringContaining(
|
||||
'Error writing to /mock/project/.taskmasterconfig: Disk full.'
|
||||
`Error writing configuration to ${MOCK_CONFIG_PATH}: Disk full`
|
||||
)
|
||||
);
|
||||
});
|
||||
|
||||
test('should return false if config file does not exist', () => {
|
||||
mockExistsSync.mockReturnValue(false);
|
||||
const success = configManager.writeConfig(
|
||||
VALID_CUSTOM_CONFIG,
|
||||
MOCK_PROJECT_ROOT
|
||||
);
|
||||
const success = configManager.writeConfig(VALID_CUSTOM_CONFIG);
|
||||
|
||||
expect(success).toBe(false);
|
||||
expect(mockWriteFileSync).not.toHaveBeenCalled();
|
||||
|
||||
Reference in New Issue
Block a user