@artemiskit/cli 0.1.8 → 0.2.2

package/src/commands/init.ts

@@ -7,7 +7,13 @@ import { appendFile, mkdir, readFile, writeFile } from 'node:fs/promises';
 import { join } from 'node:path';
 import chalk from 'chalk';
 import { Command } from 'commander';
-import { createSpinner, icons } from '../ui/index.js';
+import {
+  type InitWizardResult,
+  createSpinner,
+  icons,
+  isInteractive,
+  runInitWizard,
+} from '../ui/index.js';
 import { checkForUpdateAndNotify, getCurrentVersion } from '../utils/update-checker.js';
 
 const DEFAULT_CONFIG = `# ArtemisKit Configuration
@@ -86,6 +92,67 @@ const ENV_KEYS = [
   'ANTHROPIC_API_KEY=',
 ];
 
+/**
+ * Generate config content from wizard results
+ */
+function generateConfigFromWizard(wizard: InitWizardResult): string {
+  const providerConfigs: Record<string, string> = {
+    openai: `  openai:
+    apiKey: \${OPENAI_API_KEY}
+    defaultModel: ${wizard.model}`,
+    'azure-openai': `  azure-openai:
+    apiKey: \${AZURE_OPENAI_API_KEY}
+    resourceName: \${AZURE_OPENAI_RESOURCE}
+    deploymentName: \${AZURE_OPENAI_DEPLOYMENT}
+    apiVersion: "2024-02-15-preview"`,
+    anthropic: `  anthropic:
+    apiKey: \${ANTHROPIC_API_KEY}
+    defaultModel: ${wizard.model}`,
+    google: `  google:
+    apiKey: \${GOOGLE_AI_API_KEY}
+    defaultModel: ${wizard.model}`,
+    mistral: `  mistral:
+    apiKey: \${MISTRAL_API_KEY}
+    defaultModel: ${wizard.model}`,
+    ollama: `  ollama:
+    baseUrl: http://localhost:11434
+    defaultModel: ${wizard.model}`,
+  };
+
+  const storageConfig =
+    wizard.storageType === 'supabase'
+      ? `storage:
+  type: supabase
+  supabaseUrl: \${SUPABASE_URL}
+  supabaseKey: \${SUPABASE_ANON_KEY}`
+      : `storage:
+  type: local
+  basePath: ./artemis-runs`;
+
+  return `# ArtemisKit Configuration
+project: ${wizard.projectName}
+
+# Default provider settings
+provider: ${wizard.provider}
+model: ${wizard.model}
+
+# Provider configurations
+providers:
+${providerConfigs[wizard.provider] || providerConfigs.openai}
+
+# Storage configuration
+${storageConfig}
+
+# Scenarios directory
+scenariosDir: ./scenarios
+
+# Output settings
+output:
+  format: json
+  dir: ./artemis-output
+`;
+}
+
 function renderWelcomeBanner(): string {
   // Brand color for "KIT" portion: #fb923c (orange)
   const brandColor = chalk.hex('#fb923c');
@@ -230,82 +297,119 @@ export function initCommand(): Command {
     .description('Initialize ArtemisKit in the current directory')
     .option('-f, --force', 'Overwrite existing configuration')
     .option('--skip-env', 'Skip adding environment variables to .env')
-    .action(async (options: { force?: boolean; skipEnv?: boolean }) => {
-      const spinner = createSpinner();
-
-      try {
-        const cwd = process.cwd();
-
-        // Show welcome banner
-        console.log(renderWelcomeBanner());
-
-        // Step 1: Create directories
-        spinner.start('Creating project structure...');
-        await mkdir(join(cwd, 'scenarios'), { recursive: true });
-        await mkdir(join(cwd, 'artemis-runs'), { recursive: true });
-        await mkdir(join(cwd, 'artemis-output'), { recursive: true });
-        spinner.succeed('Created project structure');
-
-        // Step 2: Write config file
-        const configPath = join(cwd, 'artemis.config.yaml');
-        const configExists = existsSync(configPath);
-
-        if (configExists && !options.force) {
-          spinner.info('Config file already exists (use --force to overwrite)');
-        } else {
-          spinner.start('Writing configuration...');
-          await writeFile(configPath, DEFAULT_CONFIG);
-          spinner.succeed(
-            configExists ? 'Overwrote artemis.config.yaml' : 'Created artemis.config.yaml'
-          );
-        }
+    .option('-i, --interactive', 'Run interactive setup wizard')
+    .option('-y, --yes', 'Use defaults without prompts (non-interactive)')
+    .action(
+      async (options: {
+        force?: boolean;
+        skipEnv?: boolean;
+        interactive?: boolean;
+        yes?: boolean;
+      }) => {
+        const spinner = createSpinner();
+
+        try {
+          const cwd = process.cwd();
+
+          // Show welcome banner
+          console.log(renderWelcomeBanner());
+
+          // Determine if we should run interactive wizard
+          const shouldRunWizard =
+            options.interactive || (isInteractive() && !options.yes && !options.force);
+
+          let configContent = DEFAULT_CONFIG;
+          let createExample = true;
+
+          // Run interactive wizard if applicable
+          if (shouldRunWizard) {
+            try {
+              const wizardResult = await runInitWizard();
+              configContent = generateConfigFromWizard(wizardResult);
+              createExample = wizardResult.createExample;
+              console.log(''); // Add spacing after wizard
+            } catch (wizardError) {
+              // If wizard fails (e.g., user cancels), fall back to defaults
+              if ((wizardError as Error).message?.includes('closed')) {
+                console.log(chalk.yellow('\n  Setup cancelled. Using defaults.\n'));
+              } else {
+                throw wizardError;
+              }
+            }
+          }
 
-        // Step 3: Write example scenario
-        const scenarioPath = join(cwd, 'scenarios', 'example.yaml');
-        const scenarioExists = existsSync(scenarioPath);
-
-        if (scenarioExists && !options.force) {
-          spinner.info('Example scenario already exists (use --force to overwrite)');
-        } else {
-          spinner.start('Creating example scenario...');
-          await writeFile(scenarioPath, DEFAULT_SCENARIO);
-          spinner.succeed(
-            scenarioExists ? 'Overwrote scenarios/example.yaml' : 'Created scenarios/example.yaml'
-          );
-        }
+          // Step 1: Create directories
+          spinner.start('Creating project structure...');
+          await mkdir(join(cwd, 'scenarios'), { recursive: true });
+          await mkdir(join(cwd, 'artemis-runs'), { recursive: true });
+          await mkdir(join(cwd, 'artemis-output'), { recursive: true });
+          spinner.succeed('Created project structure');
 
-        // Step 4: Update .env file
-        if (!options.skipEnv) {
-          spinner.start('Updating .env file...');
-          const { added, skipped } = await appendEnvKeys(cwd);
-
-          if (added.length > 0) {
-            spinner.succeed(`Added ${added.length} environment variable(s) to .env`);
-            if (skipped.length > 0) {
-              console.log(
-                chalk.dim(
-                  `  ${icons.info} Skipped ${skipped.length} existing key(s): ${skipped.join(', ')}`
-                )
+          // Step 2: Write config file
+          const configPath = join(cwd, 'artemis.config.yaml');
+          const configExists = existsSync(configPath);
+
+          if (configExists && !options.force) {
+            spinner.info('Config file already exists (use --force to overwrite)');
+          } else {
+            spinner.start('Writing configuration...');
+            await writeFile(configPath, configContent);
+            spinner.succeed(
+              configExists ? 'Overwrote artemis.config.yaml' : 'Created artemis.config.yaml'
+            );
+          }
+
+          // Step 3: Write example scenario (if requested)
+          if (createExample) {
+            const scenarioPath = join(cwd, 'scenarios', 'example.yaml');
+            const scenarioExists = existsSync(scenarioPath);
+
+            if (scenarioExists && !options.force) {
+              spinner.info('Example scenario already exists (use --force to overwrite)');
+            } else {
+              spinner.start('Creating example scenario...');
+              await writeFile(scenarioPath, DEFAULT_SCENARIO);
+              spinner.succeed(
+                scenarioExists
+                  ? 'Overwrote scenarios/example.yaml'
+                  : 'Created scenarios/example.yaml'
               );
             }
-          } else if (skipped.length > 0) {
-            spinner.info('All environment variables already exist in .env');
-          } else {
-            spinner.succeed('Created .env with environment variables');
           }
-        }
 
-        // Show success panel
-        console.log(renderSuccessPanel());
+          // Step 4: Update .env file
+          if (!options.skipEnv) {
+            spinner.start('Updating .env file...');
+            const { added, skipped } = await appendEnvKeys(cwd);
+
+            if (added.length > 0) {
+              spinner.succeed(`Added ${added.length} environment variable(s) to .env`);
+              if (skipped.length > 0) {
+                console.log(
+                  chalk.dim(
+                    `  ${icons.info} Skipped ${skipped.length} existing key(s): ${skipped.join(', ')}`
+                  )
+                );
+              }
+            } else if (skipped.length > 0) {
+              spinner.info('All environment variables already exist in .env');
+            } else {
+              spinner.succeed('Created .env with environment variables');
+            }
+          }
 
-        // Non-blocking update check (fire and forget)
-        checkForUpdateAndNotify();
-      } catch (error) {
-        spinner.fail('Error');
-        console.error(chalk.red(`\n${icons.failed} ${(error as Error).message}`));
-        process.exit(1);
+          // Show success panel
+          console.log(renderSuccessPanel());
+
+          // Non-blocking update check (fire and forget)
+          checkForUpdateAndNotify();
+        } catch (error) {
+          spinner.fail('Error');
+          console.error(chalk.red(`\n${icons.failed} ${(error as Error).message}`));
+          process.exit(1);
+        }
       }
-    });
+    );
 
   return cmd;
 }

package/src/commands/redteam.ts

@@ -19,14 +19,18 @@ import {
   parseScenarioFile,
 } from '@artemiskit/core';
 import {
+  type ConversationTurn,
   CotInjectionMutation,
+  EncodingMutation,
   InstructionFlipMutation,
+  MultiTurnMutation,
   type Mutation,
   RedTeamGenerator,
   RoleSpoofMutation,
   SeverityMapper,
   TypoMutation,
   UnsafeResponseDetector,
+  loadCustomAttacks,
 } from '@artemiskit/redteam';
 import { generateJSONReport, generateRedTeamHTMLReport } from '@artemiskit/reports';
 import chalk from 'chalk';
@@ -55,6 +59,7 @@ interface RedteamOptions {
   model?: string;
   mutations?: string[];
   count?: number;
+  customAttacks?: string;
   save?: boolean;
   output?: string;
   verbose?: boolean;
@@ -73,9 +78,10 @@ export function redteamCommand(): Command {
     .option('-m, --model <model>', 'Model to use')
     .option(
       '--mutations <mutations...>',
-      'Mutations to apply (typo, role-spoof, instruction-flip, cot-injection)'
+      'Mutations to apply (typo, role-spoof, instruction-flip, cot-injection, encoding, multi-turn)'
     )
     .option('-c, --count <number>', 'Number of mutated prompts per case', '5')
+    .option('--custom-attacks <path>', 'Path to custom attacks YAML file')
     .option('--save', 'Save results to storage')
     .option('-o, --output <dir>', 'Output directory for reports')
     .option('-v, --verbose', 'Verbose output')
@@ -131,7 +137,7 @@ export function redteamCommand(): Command {
         spinner.succeed(`Connected to ${provider}`);
 
         // Set up mutations
-        const mutations = selectMutations(options.mutations);
+        const mutations = selectMutations(options.mutations, options.customAttacks);
         const generator = new RedTeamGenerator(mutations);
         const detector = new UnsafeResponseDetector();
         const count = Number.parseInt(String(options.count)) || 5;
@@ -179,12 +185,47 @@ export function redteamCommand(): Command {
         for (const testCase of scenario.cases) {
           console.log(chalk.bold(`Testing case: ${testCase.id}`));
 
-          const originalPrompt =
-            typeof testCase.prompt === 'string'
-              ? testCase.prompt
-              : testCase.prompt.map((m) => m.content).join('\n');
+          // Handle both string and array prompts (consistent with run command)
+          // For array prompts: last user message is the attack target, rest is context
+          let attackPrompt: string;
+          let conversationPrefix: ConversationTurn[] | undefined;
+
+          if (typeof testCase.prompt === 'string') {
+            // Simple string prompt - use directly
+            attackPrompt = testCase.prompt;
+          } else {
+            // Array prompt - extract last user message as attack, rest as context
+            const messages = testCase.prompt;
+            const lastUserIndex = messages.map((m) => m.role).lastIndexOf('user');
+
+            if (lastUserIndex === -1) {
+              // No user message found - use concatenated content
+              attackPrompt = messages.map((m) => m.content).join('\n');
+            } else {
+              // Extract attack prompt (last user message)
+              attackPrompt = messages[lastUserIndex].content;
+
+              // Extract conversation prefix (everything before the last user message)
+              if (lastUserIndex > 0) {
+                conversationPrefix = messages.slice(0, lastUserIndex).map((m) => ({
+                  role: m.role as 'user' | 'assistant' | 'system',
+                  content: m.content,
+                }));
+              }
+            }
+          }
+
+          // Clear any previous prefix and set new one if applicable
+          for (const mutation of mutations) {
+            if (mutation instanceof MultiTurnMutation) {
+              mutation.clearConversationPrefix();
+              if (conversationPrefix && conversationPrefix.length > 0) {
+                mutation.setConversationPrefix(conversationPrefix);
+              }
+            }
+          }
 
-          const mutatedPrompts = generator.generate(originalPrompt, count);
+          const mutatedPrompts = generator.generate(attackPrompt, count);
 
           for (const mutated of mutatedPrompts) {
             completedTests++;
@@ -474,19 +515,31 @@ export function redteamCommand(): Command {
   return cmd;
 }
 
-function selectMutations(names?: string[]): Mutation[] {
+function selectMutations(names?: string[], customAttacksPath?: string): Mutation[] {
   const allMutations: Record<string, Mutation> = {
     typo: new TypoMutation(),
     'role-spoof': new RoleSpoofMutation(),
     'instruction-flip': new InstructionFlipMutation(),
     'cot-injection': new CotInjectionMutation(),
+    encoding: new EncodingMutation(),
+    'multi-turn': new MultiTurnMutation(),
   };
 
+  let mutations: Mutation[];
+
   if (!names || names.length === 0) {
-    return Object.values(allMutations);
+    mutations = Object.values(allMutations);
+  } else {
+    mutations = names.filter((name) => name in allMutations).map((name) => allMutations[name]);
+  }
+
+  // Load custom attacks if path provided
+  if (customAttacksPath) {
+    const customMutations = loadCustomAttacks(customAttacksPath);
+    mutations.push(...customMutations);
   }
 
-  return names.filter((name) => name in allMutations).map((name) => allMutations[name]);
+  return mutations;
 }
 
 /**