@artemiskit/cli 0.1.8 → 0.2.0

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;
 }
 
 /**