custom-providers.ts

/**
 * @fileoverview Custom provider handlers for model setup
 */

import { CUSTOM_PROVIDERS } from '@tm/core';
import chalk from 'chalk';
import inquirer from 'inquirer';
import { getAzureBaseURL } from '../../lib/model-management.js';
import { validateOllamaModel, validateOpenRouterModel } from './fetchers.js';
import { CUSTOM_PROVIDER_IDS } from './types.js';
import type {
	CustomProviderConfig,
	CustomProviderId,
	ModelRole
} from './types.js';

/**
 * Configuration for all custom providers
 */
export const customProviderConfigs: Record<
	keyof typeof CUSTOM_PROVIDER_IDS,
	CustomProviderConfig
> = {
	OPENROUTER: {
		id: '__CUSTOM_OPENROUTER__',
		name: '* Custom OpenRouter model',
		provider: CUSTOM_PROVIDERS.OPENROUTER,
		promptMessage: (role) =>
			`Enter the custom OpenRouter Model ID for the ${role} role:`,
		validate: async (modelId) => {
			const isValid = await validateOpenRouterModel(modelId);
			if (!isValid) {
				console.error(
					chalk.red(
						`Error: Model ID "${modelId}" not found in the live OpenRouter model list. Please check the ID.`
					)
				);
			}
			return isValid;
		}
	},
	OLLAMA: {
		id: '__CUSTOM_OLLAMA__',
		name: '* Custom Ollama model',
		provider: CUSTOM_PROVIDERS.OLLAMA,
		requiresBaseURL: true,
		defaultBaseURL: 'http://localhost:11434/api',
		promptMessage: (role) =>
			`Enter the custom Ollama Model ID for the ${role} role:`,
		validate: async (modelId, baseURL) => {
			const urlToCheck = baseURL || 'http://localhost:11434/api';
			const isValid = await validateOllamaModel(modelId, urlToCheck);
			if (!isValid) {
				console.error(
					chalk.red(
						`Error: Model ID "${modelId}" not found in the Ollama instance. Please verify the model is pulled and available.`
					)
				);
				console.log(
					chalk.yellow(
						`You can check available models with: curl ${urlToCheck}/tags`
					)
				);
			}
			return isValid;
		}
	},
	BEDROCK: {
		id: '__CUSTOM_BEDROCK__',
		name: '* Custom Bedrock model',
		provider: CUSTOM_PROVIDERS.BEDROCK,
		promptMessage: (role) =>
			`Enter the custom Bedrock Model ID for the ${role} role (e.g., anthropic.claude-3-sonnet-20240229-v1:0):`,
		checkEnvVars: () => {
			if (
				!process.env.AWS_ACCESS_KEY_ID ||
				!process.env.AWS_SECRET_ACCESS_KEY
			) {
				console.warn(
					chalk.yellow(
						'Warning: AWS_ACCESS_KEY_ID and/or AWS_SECRET_ACCESS_KEY environment variables are missing. Will fallback to system configuration (ex: aws config files or ec2 instance profiles).'
					)
				);
			}
			return true;
		}
	},
	AZURE: {
		id: '__CUSTOM_AZURE__',
		name: '* Custom Azure OpenAI model',
		provider: CUSTOM_PROVIDERS.AZURE,
		requiresBaseURL: true,
		promptMessage: (role) =>
			`Enter the Azure deployment name for the ${role} role (e.g., gpt-4o):`,
		checkEnvVars: () => {
			if (!process.env.AZURE_OPENAI_API_KEY) {
				console.error(
					chalk.red(
						'Error: AZURE_OPENAI_API_KEY environment variable is missing. Please set it before using Azure models.'
					)
				);
				return false;
			}
			return true;
		}
	},
	VERTEX: {
		id: '__CUSTOM_VERTEX__',
		name: '* Custom Vertex model',
		provider: CUSTOM_PROVIDERS.VERTEX,
		promptMessage: (role) =>
			`Enter the custom Vertex AI Model ID for the ${role} role (e.g., gemini-1.5-pro-002):`,
		checkEnvVars: () => {
			if (
				!process.env.GOOGLE_API_KEY &&
				!process.env.GOOGLE_APPLICATION_CREDENTIALS
			) {
				console.error(
					chalk.red(
						'Error: Either GOOGLE_API_KEY or GOOGLE_APPLICATION_CREDENTIALS environment variable is required. Please set one before using custom Vertex models.'
					)
				);
				return false;
			}
			return true;
		}
	},
	VERTEX_ANTHROPIC: {
		id: '__CUSTOM_VERTEX_ANTHROPIC__',
		name: '* Custom Vertex Anthropic model',
		provider: CUSTOM_PROVIDERS.VERTEX_ANTHROPIC,
		promptMessage: (role) =>
			`Enter the custom Vertex AI Anthropic Model ID for the ${role} role (e.g., claude-sonnet-4-6):`,
		checkEnvVars: () => {
			if (
				!process.env.GOOGLE_API_KEY &&
				!process.env.GOOGLE_APPLICATION_CREDENTIALS
			) {
				console.error(
					chalk.red(
						'Error: Either GOOGLE_API_KEY or GOOGLE_APPLICATION_CREDENTIALS environment variable is required. Please set one before using Vertex Anthropic models.'
					)
				);
				return false;
			}
			return true;
		}
	},
	LMSTUDIO: {
		id: '__CUSTOM_LMSTUDIO__',
		name: '* Custom LMStudio model',
		provider: CUSTOM_PROVIDERS.LMSTUDIO,
		requiresBaseURL: true,
		defaultBaseURL: 'http://localhost:1234/v1',
		promptMessage: (role) =>
			`Enter the custom LM Studio Model ID for the ${role} role:`,
		checkEnvVars: () => {
			console.log(
				chalk.blue(
					'Note: LM Studio runs locally. Make sure the LM Studio server is running.'
				)
			);
			return true;
		}
	},
	OPENAI_COMPATIBLE: {
		id: '__CUSTOM_OPENAI_COMPATIBLE__',
		name: '* Custom OpenAI-compatible model',
		provider: CUSTOM_PROVIDERS.OPENAI_COMPATIBLE,
		promptMessage: (role) =>
			`Enter the custom OpenAI-compatible Model ID for the ${role} role:`,
		requiresBaseURL: true,
		checkEnvVars: () => {
			console.log(
				chalk.blue(
					'Note: This will configure a generic OpenAI-compatible provider. Make sure your API endpoint is accessible.'
				)
			);
			return true;
		}
	}
};

/**
 * Handle custom provider selection
 */
export async function handleCustomProvider(
	providerId: CustomProviderId,
	role: ModelRole,
	currentModel: {
		modelId?: string | null;
		provider?: string | null;
		baseURL?: string | null;
	} | null = null,
	projectRoot?: string
): Promise<{
	modelId: string | null;
	provider: string | null;
	baseURL?: string | null;
	success: boolean;
}> {
	// Find the matching config
	const configEntry = Object.entries(customProviderConfigs).find(
		([_, config]) => config.id === providerId
	);

	if (!configEntry) {
		console.error(chalk.red(`Unknown custom provider: ${providerId}`));
		return { modelId: null, provider: null, success: false };
	}

	const config = configEntry[1];

	// Check environment variables if needed
	if (config.checkEnvVars && !config.checkEnvVars()) {
		return { modelId: null, provider: null, success: false };
	}

	// Prompt for baseURL if required
	let baseURL: string | null = null;
	if (config.requiresBaseURL) {
		// Determine the appropriate default baseURL
		let defaultBaseURL: string;
		if (currentModel?.provider === config.provider && currentModel?.baseURL) {
			// Already using this provider - preserve existing baseURL
			defaultBaseURL = currentModel.baseURL;
		} else if (config.provider === CUSTOM_PROVIDERS.AZURE && projectRoot) {
			// For Azure, try to use the global azureBaseURL from config
			defaultBaseURL = getAzureBaseURL(projectRoot) || '';
		} else {
			// Switching providers or no existing baseURL - use fallback default
			defaultBaseURL = config.defaultBaseURL || '';
		}

		const baseURLAnswer = await inquirer.prompt([
			{
				type: 'input',
				name: 'baseURL',
				message: `Enter the base URL for the ${role} role:`,
				default: defaultBaseURL,
				validate: (input: string) => {
					if (!input || input.trim() === '') {
						return `Base URL is required for ${config.provider} providers`;
					}
					try {
						new URL(input);
						return true;
					} catch {
						return 'Please enter a valid URL';
					}
				}
			}
		]);
		baseURL = baseURLAnswer.baseURL;
	}

	// Prompt for custom ID
	const { customId } = await inquirer.prompt([
		{
			type: 'input',
			name: 'customId',
			message: config.promptMessage(role)
		}
	]);

	if (!customId) {
		console.log(chalk.yellow('No custom ID entered. Skipping role.'));
		return { modelId: null, provider: null, success: true };
	}

	// Validate if validation function exists
	if (config.validate) {
		const isValid = await config.validate(customId, baseURL || undefined);
		if (!isValid) {
			return { modelId: null, provider: null, success: false };
		}
	} else {
		console.log(
			chalk.blue(
				`Custom ${config.provider} model "${customId}" will be used. No validation performed.`
			)
		);
	}

	return {
		modelId: customId,
		provider: config.provider,
		baseURL: baseURL,
		success: true
	};
}

/**
 * Get all custom provider options for display
 */
export function getCustomProviderOptions(): Array<{
	name: string;
	value: CustomProviderId;
	short: string;
}> {
	return Object.values(customProviderConfigs).map((config) => ({
		name: config.name,
		value: config.id,
		short: config.name
	}));
}