Skip to main content

tower_llm/validation/
mutate.rs

1//! Mutators to introduce specific violations into otherwise valid conversations.
2
3use async_openai::types::*;
4
5#[derive(Debug, Clone, Copy)]
6pub enum MutationKind {
7    AssistantBeforeUser,
8    RepeatedUser,
9    MissingOneToolResponse,
10    UnknownToolResponse,
11    ReorderToolResponses,
12    DuplicateToolResponse,
13    SystemNotFirst,
14    DuplicateToolCallIdsInAssistant,
15    EmptyToolCallIdInAssistant,
16    EmptyToolMessageId,
17    ToolBeforeAssistant,
18    ToolResponsesNotContiguous,
19    RemoveAllUsers,
20}
21
22/// Apply a mutation in-place. Returns true if mutation was applied.
23pub fn apply_violation(
24    messages: &mut Vec<ChatCompletionRequestMessage>,
25    kind: MutationKind,
26) -> bool {
27    match kind {
28        MutationKind::AssistantBeforeUser => {
29            // Find first user; insert assistant at beginning if user exists.
30            if messages
31                .iter()
32                .any(|m| matches!(m, ChatCompletionRequestMessage::User(_)))
33            {
34                let asst = ChatCompletionRequestAssistantMessageArgs::default()
35                    .content("early")
36                    .build()
37                    .unwrap();
38                messages.insert(0, asst.into());
39                return true;
40            }
41            false
42        }
43        MutationKind::RepeatedUser => {
44            // Duplicate first user consecutively
45            if let Some((idx, _)) = messages
46                .iter()
47                .enumerate()
48                .find(|(_, m)| matches!(m, ChatCompletionRequestMessage::User(_)))
49            {
50                if let ChatCompletionRequestMessage::User(u) = messages[idx].clone() {
51                    messages.insert(idx + 1, ChatCompletionRequestMessage::User(u));
52                    return true;
53                }
54            }
55            false
56        }
57        MutationKind::MissingOneToolResponse => {
58            // Remove the first tool message if any
59            if let Some((idx, _)) = messages
60                .iter()
61                .enumerate()
62                .find(|(_, m)| matches!(m, ChatCompletionRequestMessage::Tool(_)))
63            {
64                messages.remove(idx);
65                return true;
66            }
67            false
68        }
69        MutationKind::UnknownToolResponse => {
70            let t = ChatCompletionRequestToolMessageArgs::default()
71                .tool_call_id("unknown_id")
72                .content("{}")
73                .build()
74                .unwrap();
75            // Insert near the end
76            messages.push(t.into());
77            true
78        }
79        MutationKind::ReorderToolResponses => {
80            // Look for a block of two consecutive tool messages and swap them
81            for i in 0..messages.len().saturating_sub(1) {
82                if matches!(messages[i], ChatCompletionRequestMessage::Tool(_))
83                    && matches!(messages[i + 1], ChatCompletionRequestMessage::Tool(_))
84                {
85                    messages.swap(i, i + 1);
86                    return true;
87                }
88            }
89            false
90        }
91        MutationKind::DuplicateToolResponse => {
92            // Duplicate the first tool message and re-insert after it.
93            if let Some((idx, m)) = messages
94                .iter()
95                .enumerate()
96                .find(|(_, m)| matches!(m, ChatCompletionRequestMessage::Tool(_)))
97            {
98                messages.insert(idx + 1, m.clone());
99                return true;
100            }
101            false
102        }
103        MutationKind::SystemNotFirst => {
104            // Move first system to the middle, if any
105            if let Some((idx, m)) = messages
106                .iter()
107                .enumerate()
108                .find(|(_, m)| matches!(m, ChatCompletionRequestMessage::System(_)))
109            {
110                let sys = m.clone();
111                messages.remove(idx);
112                let insert_at = (messages.len() / 2).max(1);
113                messages.insert(insert_at, sys);
114                return true;
115            }
116            false
117        }
118        MutationKind::DuplicateToolCallIdsInAssistant => {
119            for msg in messages.iter_mut() {
120                if let ChatCompletionRequestMessage::Assistant(asst) = msg {
121                    if let Some(calls) = asst.tool_calls.as_mut() {
122                        if calls.len() >= 2 {
123                            let id = calls[0].id.clone();
124                            calls[1].id = id;
125                            return true;
126                        }
127                    }
128                }
129            }
130            false
131        }
132        MutationKind::EmptyToolCallIdInAssistant => {
133            for msg in messages.iter_mut() {
134                if let ChatCompletionRequestMessage::Assistant(asst) = msg {
135                    if let Some(calls) = asst.tool_calls.as_mut() {
136                        if !calls.is_empty() {
137                            calls[0].id = String::new();
138                            return true;
139                        }
140                    }
141                }
142            }
143            false
144        }
145        MutationKind::EmptyToolMessageId => {
146            for m in messages.iter_mut() {
147                if let ChatCompletionRequestMessage::Tool(tmsg) = m {
148                    tmsg.tool_call_id = String::new();
149                    return true;
150                }
151            }
152            false
153        }
154        MutationKind::ToolBeforeAssistant => {
155            // Insert a tool message at the beginning with a fresh id
156            let tool = ChatCompletionRequestToolMessageArgs::default()
157                .tool_call_id("pre_tool")
158                .content("{}")
159                .build()
160                .unwrap();
161            messages.insert(0, ChatCompletionRequestMessage::Tool(tool));
162            true
163        }
164        MutationKind::ToolResponsesNotContiguous => {
165            // Find an assistant with tool_calls and interrupt with a user before the first tool response
166            for i in 0..messages.len() {
167                if let ChatCompletionRequestMessage::Assistant(asst) = &messages[i] {
168                    if asst
169                        .tool_calls
170                        .as_ref()
171                        .map(|v| !v.is_empty())
172                        .unwrap_or(false)
173                        && i + 1 < messages.len()
174                    {
175                        // Insert a user immediately after assistant
176                        let u = ChatCompletionRequestUserMessageArgs::default()
177                            .content("interrupt")
178                            .build()
179                            .unwrap();
180                        messages.insert(i + 1, ChatCompletionRequestMessage::User(u));
181                        return true;
182                    }
183                }
184            }
185            false
186        }
187        MutationKind::RemoveAllUsers => {
188            let had_users = messages
189                .iter()
190                .any(|m| matches!(m, ChatCompletionRequestMessage::User(_)));
191            messages.retain(|m| !matches!(m, ChatCompletionRequestMessage::User(_)));
192            had_users
193        }
194    }
195}