ai-speedometer 1.3.4 → 1.3.7

package/cli.js

@@ -34,7 +34,10 @@ function parseCliArgs() {
   const parsed = {
     debug: false,
     bench: null,
+    benchCustom: null,
     apiKey: null,
+    baseUrl: null,
+    endpointFormat: null,
     useAiSdk: false,
     formatted: false,
     help: false
@@ -47,8 +50,14 @@ function parseCliArgs() {
       parsed.debug = true;
     } else if (arg === '--bench') {
       parsed.bench = args[++i];
+    } else if (arg === '--bench-custom') {
+      parsed.benchCustom = args[++i];
     } else if (arg === '--api-key') {
       parsed.apiKey = args[++i];
+    } else if (arg === '--base-url') {
+      parsed.baseUrl = args[++i];
+    } else if (arg === '--endpoint-format') {
+      parsed.endpointFormat = args[++i];
     } else if (arg === '--ai-sdk') {
       parsed.useAiSdk = true;
     } else if (arg === '--formatted') {
@@ -58,7 +67,7 @@ function parseCliArgs() {
     }
   }
 
-  return parsed;
+return parsed;
 }
 
 function showHelp() {
@@ -67,23 +76,65 @@ function showHelp() {
   console.log(colorText('Usage:', 'yellow'));
   console.log('  ai-speedometer                                  ' + colorText('# Interactive mode', 'dim'));
   console.log('  ai-speedometer --bench <provider:model>         ' + colorText('# Headless benchmark', 'dim'));
+  console.log('  ai-speedometer --bench-custom <provider:model>  ' + colorText('# Custom provider benchmark', 'dim'));
   console.log('');
   console.log(colorText('Options:', 'yellow'));
-  console.log('  --bench <provider:model>    ' + colorText('Run benchmark in headless mode', 'dim'));
-  console.log('  --api-key <key>             ' + colorText('Override API key (optional)', 'dim'));
-  console.log('  --ai-sdk                    ' + colorText('Use AI SDK instead of REST API', 'dim'));
-  console.log('  --formatted                 ' + colorText('Format JSON output for human readability', 'dim'));
-  console.log('  --debug                     ' + colorText('Enable debug logging', 'dim'));
-  console.log('  --help, -h                  ' + colorText('Show this help message', 'dim'));
+  console.log('  --bench <provider:model>      ' + colorText('Run benchmark in headless mode', 'dim'));
+  console.log('  --bench-custom <provider:model> ' + colorText('Run custom provider benchmark', 'dim'));
+  console.log('  --base-url <url>              ' + colorText('Base URL for custom provider', 'dim'));
+  console.log('  --api-key <key>               ' + colorText('API key for custom provider', 'dim'));
+  console.log('  --endpoint-format <format>    ' + colorText('Endpoint format (default: chat/completions)', 'dim'));
+  console.log('  --ai-sdk                      ' + colorText('Use AI SDK instead of REST API', 'dim'));
+  console.log('  --formatted                   ' + colorText('Format JSON output for human readability', 'dim'));
+  console.log('  --debug                       ' + colorText('Enable debug logging', 'dim'));
+  console.log('  --help, -h                    ' + colorText('Show this help message', 'dim'));
   console.log('');
   console.log(colorText('Examples:', 'yellow'));
   console.log('  ai-speedometer --bench openai:gpt-4');
   console.log('  ai-speedometer --bench anthropic:claude-3-opus --api-key "sk-..."');
-  console.log('  ai-speedometer --bench openai:gpt-4 --ai-sdk');
-  console.log('  ai-speedometer --bench openai:gpt-4 --formatted');
+  console.log('  ai-speedometer --bench-custom openai:gpt-4 --base-url "https://api.openai.com/v1" --api-key "sk-..."');
+  console.log('  ai-speedometer --bench-custom anthropic:claude --base-url "https://api.anthropic.com/v1" --api-key "sk-..." --endpoint-format "messages"');
   console.log('');
 }
 
+// Parse provider:model format, handling colons in model IDs
+function parseProviderModel(arg) {
+  const firstColonIndex = arg.indexOf(':');
+  if (firstColonIndex === -1) {
+    throw new Error(`Invalid format. Use provider:model (e.g., openai:gpt-4)`);
+  }
+  
+  const provider = arg.substring(0, firstColonIndex);
+  const model = arg.substring(firstColonIndex + 1);
+  
+  return { provider, model };
+}
+
+// Create temporary custom provider from CLI args
+function createCustomProviderFromCli(cliArgs) {
+  const { provider, model } = parseProviderModel(cliArgs.benchCustom);
+  
+  // Validate required arguments
+  if (!cliArgs.baseUrl) {
+    throw new Error('--base-url is required for custom provider benchmarking');
+  }
+  if (!cliArgs.apiKey) {
+    throw new Error('--api-key is required for custom provider benchmarking');
+  }
+  
+  const endpointFormat = cliArgs.endpointFormat || 'chat/completions';
+  
+  return {
+    id: provider,
+    name: provider,
+    type: 'openai-compatible', // Default to OpenAI compatible for custom providers
+    baseUrl: cliArgs.baseUrl,
+    apiKey: cliArgs.apiKey,
+    endpointFormat: endpointFormat,
+    models: [{ name: model, id: model }]
+  };
+}
+
 const cliArgs = parseCliArgs();
 const debugMode = cliArgs.debug;
 let logFile = null;
@@ -1350,7 +1401,7 @@ async function addCustomProviderCLI() {
       const modelName = await question(colorText('Model name: ', 'cyan'));
       if (!modelName.trim()) break;
       
-      const modelId = modelName.trim().toLowerCase().replace(/[^a-z0-9-]/g, '-') + '_' + Date.now();
+      const modelId = modelName.trim().toLowerCase().replace(/[^a-z0-9-]/g, '-');
       models.push({
         name: modelName.trim(),
         id: modelId
@@ -1360,7 +1411,7 @@ async function addCustomProviderCLI() {
     // Single model mode
     const modelName = await question(colorText('Enter model name (e.g., gpt-4): ', 'cyan'));
     if (modelName.trim()) {
-      const modelId = modelName.trim().toLowerCase().replace(/[^a-z0-9-]/g, '-') + '_' + Date.now();
+      const modelId = modelName.trim().toLowerCase().replace(/[^a-z0-9-]/g, '-');
       models.push({
         name: modelName.trim(),
         id: modelId
@@ -1676,7 +1727,7 @@ async function addModelsToExistingProvider() {
       const modelName = await question(colorText('Model name: ', 'cyan'));
       if (!modelName.trim()) break;
       
-      const modelId = modelName.trim().toLowerCase().replace(/[^a-z0-9-]/g, '-') + '_' + Date.now();
+      const modelId = modelName.trim().toLowerCase().replace(/[^a-z0-9-]/g, '-');
       const modelData = {
         name: modelName.trim(),
         id: modelId
@@ -1694,7 +1745,7 @@ async function addModelsToExistingProvider() {
     // Single model mode
     const modelName = await question(colorText('Enter model name: ', 'cyan'));
     if (modelName.trim()) {
-      const modelId = modelName.trim().toLowerCase().replace(/[^a-z0-9-]/g, '-') + '_' + Date.now();
+      const modelId = modelName.trim().toLowerCase().replace(/[^a-z0-9-]/g, '-');
       const modelData = {
         name: modelName.trim(),
         id: modelId
@@ -1746,9 +1797,12 @@ async function benchmarkSingleModelRest(model) {
       let streamedText = '';
       let tokenCount = 0;
       
-      // Use correct endpoint based on provider type
+      // Use correct endpoint based on provider type or custom format
       let endpoint;
-      if (model.providerType === 'anthropic') {
+      if (model.providerConfig.endpointFormat) {
+        // Use custom endpoint format
+        endpoint = '/' + model.providerConfig.endpointFormat;
+      } else if (model.providerType === 'anthropic') {
         endpoint = '/messages';
       } else if (model.providerType === 'google') {
         endpoint = '/models/' + actualModelId + ':streamGenerateContent';
@@ -1834,7 +1888,7 @@ async function benchmarkSingleModelRest(model) {
           isFirstChunk = false;
           // Show live TTFT result (only in interactive mode, not headless)
           const ttftSeconds = ((firstTokenTime - startTime) / 1000).toFixed(2);
-          if (!cliArgs.bench) {
+          if (!cliArgs.bench && !cliArgs.benchCustom) {
             console.log(colorText(`TTFT received at ${ttftSeconds}s for ${model.name}`, 'green'));
           }
         }
@@ -2172,10 +2226,71 @@ process.on('SIGINT', () => {
 });
 
 // Headless benchmark mode
-async function runHeadlessBenchmark(benchSpec, apiKey, useAiSdk) {
+async function runHeadlessBenchmark(benchSpec, apiKey, useAiSdk, cliArgs = null) {
   try {
-    // Parse provider:model format
-    const [providerSpec, modelName] = benchSpec.split(':');
+    // Check if this is a custom provider benchmark
+    if (cliArgs && cliArgs.benchCustom) {
+      // Handle custom provider
+      const customProvider = createCustomProviderFromCli(cliArgs);
+      
+      // Create model object for benchmarking
+      const modelConfig = {
+        ...customProvider.models[0],
+        providerName: customProvider.name,
+        providerType: customProvider.type,
+        providerId: customProvider.id,
+        providerConfig: {
+          baseUrl: customProvider.baseUrl,
+          apiKey: customProvider.apiKey,
+          endpointFormat: customProvider.endpointFormat
+        },
+        selected: true
+      };
+
+      // Run benchmark silently and get results
+      const result = await benchmarkSingleModelRest(modelConfig);
+
+      // Output JSON to stdout (same format as regular benchmarks)
+      const jsonOutput = {
+        provider: customProvider.name,
+        providerId: customProvider.id,
+        model: customProvider.models[0].name,
+        modelId: customProvider.models[0].id,
+        method: 'rest-api',
+        success: result.success,
+        totalTime: result.totalTime,
+        totalTimeSeconds: result.totalTime / 1000,
+        timeToFirstToken: result.timeToFirstToken,
+        timeToFirstTokenSeconds: result.timeToFirstToken / 1000,
+        tokensPerSecond: result.tokensPerSecond,
+        outputTokens: result.tokenCount,
+        promptTokens: result.promptTokens,
+        totalTokens: result.totalTokens,
+        is_estimated: !!(result.usedEstimateForOutput || result.usedEstimateForInput),
+        error: result.error || null
+      };
+
+      console.log(JSON.stringify(jsonOutput, null, cliArgs.formatted ? 2 : 0));
+      process.exit(result.success ? 0 : 1);
+    }
+
+    // Handle regular provider:model format
+    let providerSpec, modelName;
+    const colonIndex = benchSpec.indexOf(':');
+    if (colonIndex === -1) {
+      console.error(colorText('Error: Invalid --bench format. Use: provider:model', 'red'));
+      console.error(colorText('Example: --bench zai-code-anth:glm-4.6', 'yellow'));
+      process.exit(1);
+    }
+    
+    providerSpec = benchSpec.substring(0, colonIndex);
+    modelName = benchSpec.substring(colonIndex + 1);
+    
+    // Remove quotes from model name if present
+    if ((modelName.startsWith('"') && modelName.endsWith('"')) || 
+        (modelName.startsWith("'") && modelName.endsWith("'"))) {
+      modelName = modelName.slice(1, -1);
+    }
     
     if (!providerSpec || !modelName) {
       console.error(colorText('Error: Invalid --bench format. Use: provider:model', 'red'));
@@ -2202,19 +2317,22 @@ async function runHeadlessBenchmark(benchSpec, apiKey, useAiSdk) {
     }
 
     // Find the model
-    // Model IDs are prefixed with provider name (e.g., "zai-code-anth_glm-4.6")
-    // So we need to check:
-    // 1. Full ID match: "zai-code-anth_glm-4.6"
-    // 2. ID without provider prefix: "glm-4.6"
-    // 3. Name match: "GLM-4.6-anth"
+    // First try exact match with the provided model ID
+    // Then fall back to legacy matching for compatibility
     const model = provider.models.find(m => {
       const modelIdLower = m.id?.toLowerCase() || '';
       const modelNameLower = m.name?.toLowerCase() || '';
       const searchLower = modelName.toLowerCase();
       
-      // Check full ID match
+      // Exact ID match first (for quoted model IDs like "hf:moonshotai/Kimi-K2-Instruct-0905")
       if (modelIdLower === searchLower) return true;
       
+      // Legacy matching for backward compatibility:
+      // Model IDs are prefixed with provider name (e.g., "zai-code-anth_glm-4.6")
+      // So we need to check:
+      // 1. ID without provider prefix: "glm-4.6"
+      // 2. Name match: "GLM-4.6-anth"
+      
       // Check ID without provider prefix (strip "provider_" prefix)
       const idWithoutPrefix = modelIdLower.includes('_') 
         ? modelIdLower.split('_').slice(1).join('_') 
@@ -2305,16 +2423,20 @@ async function runHeadlessBenchmark(benchSpec, apiKey, useAiSdk) {
 }
 
 // Start the CLI
-if (require.main === module) {
+if (typeof require !== 'undefined' && require.main === module) {
   // Check if help flag
   if (cliArgs.help) {
     showHelp();
     process.exit(0);
   }
   
+  // Check if custom provider benchmark mode
+  if (cliArgs.benchCustom) {
+    runHeadlessBenchmark(cliArgs.benchCustom, cliArgs.apiKey, cliArgs.useAiSdk, cliArgs);
+  }
   // Check if headless benchmark mode
-  if (cliArgs.bench) {
-    runHeadlessBenchmark(cliArgs.bench, cliArgs.apiKey, cliArgs.useAiSdk);
+  else if (cliArgs.bench) {
+    runHeadlessBenchmark(cliArgs.bench, cliArgs.apiKey, cliArgs.useAiSdk, null);
   } else {
     // Interactive mode
     cleanupRecentModelsFromConfig().then(() => {