1pub mod base;
5pub mod built_in;
7pub mod container;
9pub mod types;
11pub mod unsafe_local;
13pub mod utils;
15pub mod vertex_ai;
17
18pub use base::{CodeExecutor, CodeExecutorError};
19pub use built_in::BuiltInCodeExecutor;
20pub use container::{ContainerCodeExecutor, ContainerCodeExecutorConfig};
21pub use types::{CodeExecutionInput, CodeExecutionResult, CodeFile};
22pub use unsafe_local::UnsafeLocalCodeExecutor;
23pub use vertex_ai::{VertexAiCodeExecutor, VertexAiCodeExecutorConfig};
24
25#[cfg(test)]
26mod tests {
27 use super::*;
28 use async_trait::async_trait;
29
30 #[test]
33 fn code_file_construction() {
34 let file = CodeFile {
35 name: "main.py".into(),
36 content: "print('hello')".into(),
37 mime_type: "text/x-python".into(),
38 };
39 assert_eq!(file.name, "main.py");
40 assert_eq!(file.content, "print('hello')");
41 assert_eq!(file.mime_type, "text/x-python");
42 }
43
44 #[test]
45 fn code_file_serde_roundtrip() {
46 let file = CodeFile {
47 name: "data.csv".into(),
48 content: "a,b\n1,2".into(),
49 mime_type: "text/csv".into(),
50 };
51 let json = serde_json::to_string(&file).unwrap();
52 let deserialized: CodeFile = serde_json::from_str(&json).unwrap();
53 assert_eq!(deserialized.name, file.name);
54 assert_eq!(deserialized.content, file.content);
55 assert_eq!(deserialized.mime_type, file.mime_type);
56 }
57
58 #[test]
61 fn code_execution_result_empty() {
62 let result = CodeExecutionResult::empty();
63 assert!(result.stdout.is_empty());
64 assert!(result.stderr.is_empty());
65 assert!(result.output_files.is_empty());
66 }
67
68 #[test]
71 fn code_execution_input_construction() {
72 let input = CodeExecutionInput {
73 code: "x = 1 + 2".into(),
74 input_files: vec![CodeFile {
75 name: "input.txt".into(),
76 content: "hello".into(),
77 mime_type: "text/plain".into(),
78 }],
79 execution_id: Some("exec-123".into()),
80 };
81 assert_eq!(input.code, "x = 1 + 2");
82 assert_eq!(input.input_files.len(), 1);
83 assert_eq!(input.execution_id.as_deref(), Some("exec-123"));
84 }
85
86 #[test]
87 fn code_execution_input_no_execution_id() {
88 let input = CodeExecutionInput {
89 code: "pass".into(),
90 input_files: Vec::new(),
91 execution_id: None,
92 };
93 assert!(input.execution_id.is_none());
94 assert!(input.input_files.is_empty());
95 }
96
97 struct StubExecutor;
101
102 #[async_trait]
103 impl CodeExecutor for StubExecutor {
104 async fn execute_code(
105 &self,
106 _input: CodeExecutionInput,
107 ) -> Result<CodeExecutionResult, CodeExecutorError> {
108 Ok(CodeExecutionResult::empty())
109 }
110 }
111
112 #[test]
113 fn default_code_block_delimiters() {
114 let exec = StubExecutor;
115 let delims = exec.code_block_delimiters();
116 assert_eq!(delims.len(), 2);
117 assert_eq!(
118 delims[0],
119 ("```tool_code\n".to_string(), "\n```".to_string())
120 );
121 assert_eq!(delims[1], ("```python\n".to_string(), "\n```".to_string()));
122 }
123
124 #[test]
125 fn default_execution_result_delimiters() {
126 let exec = StubExecutor;
127 let (open, close) = exec.execution_result_delimiters();
128 assert_eq!(open, "```tool_output\n");
129 assert_eq!(close, "\n```");
130 }
131
132 #[test]
133 fn default_error_retry_attempts() {
134 let exec = StubExecutor;
135 assert_eq!(exec.error_retry_attempts(), 2);
136 }
137
138 #[test]
139 fn default_stateful() {
140 let exec = StubExecutor;
141 assert!(!exec.stateful());
142 }
143
144 #[tokio::test]
145 async fn stub_executor_execute_code() {
146 let exec = StubExecutor;
147 let input = CodeExecutionInput {
148 code: "1 + 1".into(),
149 input_files: Vec::new(),
150 execution_id: None,
151 };
152 let result = exec.execute_code(input).await.unwrap();
153 assert!(result.stdout.is_empty());
154 }
155
156 fn _assert_object_safe(_: &dyn CodeExecutor) {}
159
160 #[test]
161 fn code_executor_is_object_safe() {
162 let exec = StubExecutor;
163 _assert_object_safe(&exec);
164 }
165
166 #[test]
169 fn built_in_process_llm_request_adds_code_execution_for_gemini2() {
170 let executor = built_in::BuiltInCodeExecutor;
171 let mut request = crate::llm::LlmRequest::from_text("hello");
172 executor
173 .process_llm_request(&mut request, "gemini-2.5-flash")
174 .unwrap();
175 assert_eq!(request.tools.len(), 1);
176 assert!(request.tools[0].code_execution.is_some());
177 }
178
179 #[test]
180 fn built_in_process_llm_request_rejects_non_gemini2() {
181 let executor = built_in::BuiltInCodeExecutor;
182 let mut request = crate::llm::LlmRequest::from_text("hello");
183 let err = executor
184 .process_llm_request(&mut request, "gemini-1.5-pro")
185 .unwrap_err();
186 assert!(
187 err.to_string().contains("Gemini 2.0+"),
188 "expected UnsupportedModel error, got: {}",
189 err
190 );
191 assert!(request.tools.is_empty());
192 }
193
194 #[tokio::test]
195 async fn built_in_execute_code_returns_empty() {
196 let executor = built_in::BuiltInCodeExecutor;
197 let input = CodeExecutionInput {
198 code: "print('hi')".into(),
199 input_files: Vec::new(),
200 execution_id: None,
201 };
202 let result = executor.execute_code(input).await.unwrap();
203 assert!(result.stdout.is_empty());
204 assert!(result.stderr.is_empty());
205 assert!(result.output_files.is_empty());
206 }
207
208 #[test]
211 fn extract_code_from_tool_code_block() {
212 let text = "Some text\n```tool_code\nprint('hello')\n```\nMore text";
213 let delimiters = vec![
214 ("```tool_code\n".to_string(), "\n```".to_string()),
215 ("```python\n".to_string(), "\n```".to_string()),
216 ];
217 let (code, remaining) = utils::extract_code_from_text(text, &delimiters).unwrap();
218 assert_eq!(code, "print('hello')");
219 assert_eq!(remaining, "Some text\n\nMore text");
220 }
221
222 #[test]
223 fn extract_code_from_python_block() {
224 let text = "Intro\n```python\nx = 42\n```\nDone";
225 let delimiters = vec![
226 ("```tool_code\n".to_string(), "\n```".to_string()),
227 ("```python\n".to_string(), "\n```".to_string()),
228 ];
229 let (code, remaining) = utils::extract_code_from_text(text, &delimiters).unwrap();
230 assert_eq!(code, "x = 42");
231 assert_eq!(remaining, "Intro\n\nDone");
232 }
233
234 #[test]
235 fn extract_code_returns_none_when_no_block() {
236 let text = "Just plain text, no code blocks here.";
237 let delimiters = vec![
238 ("```tool_code\n".to_string(), "\n```".to_string()),
239 ("```python\n".to_string(), "\n```".to_string()),
240 ];
241 assert!(utils::extract_code_from_text(text, &delimiters).is_none());
242 }
243
244 #[test]
245 fn build_executable_code_part_correct() {
246 let part = utils::build_executable_code_part("x = 1");
247 match part {
248 rs_genai::prelude::Part::ExecutableCode { executable_code } => {
249 assert_eq!(executable_code.language, "PYTHON");
250 assert_eq!(executable_code.code, "x = 1");
251 }
252 other => panic!("expected ExecutableCode, got: {:?}", other),
253 }
254 }
255
256 #[test]
257 fn build_code_execution_result_part_ok_outcome() {
258 let part = utils::build_code_execution_result_part("42\n", "");
259 match part {
260 rs_genai::prelude::Part::CodeExecutionResult {
261 code_execution_result,
262 } => {
263 assert_eq!(code_execution_result.outcome, "OK");
264 assert_eq!(code_execution_result.output.as_deref(), Some("42\n"));
265 }
266 other => panic!("expected CodeExecutionResult, got: {:?}", other),
267 }
268 }
269
270 #[test]
271 fn build_code_execution_result_part_failed_outcome() {
272 let part =
273 utils::build_code_execution_result_part("partial output", "NameError: x not defined");
274 match part {
275 rs_genai::prelude::Part::CodeExecutionResult {
276 code_execution_result,
277 } => {
278 assert_eq!(code_execution_result.outcome, "FAILED");
279 let output = code_execution_result.output.unwrap();
280 assert!(output.contains("partial output"));
281 assert!(output.contains("NameError: x not defined"));
282 }
283 other => panic!("expected CodeExecutionResult, got: {:?}", other),
284 }
285 }
286}