gemini_adk_rs/code_mode/
mod.rs

1//! Code mode: one tool through which the model runs a short Python program
2//! that calls the session's own tools.
3//!
4//! Function calling makes the model do the arithmetic, sorting and joining in
5//! its head, one call at a time, with every result added to the conversation.
6//! In code mode the model writes a program instead: `run_code` runs it in a
7//! sandbox, each function the program calls is the session's tool of that
8//! name, and the model reads only the program's final value. Measured on
9//! Gemini 3.8 Live over 120 payments, the answer was right in 5 of 5 runs
10//! against 2 of 5 with function calls, and the prompt held 77% fewer tokens.
11//! On a single lookup, writing a program cost the caller up to 700 ms before
12//! the first action, so by default the functions stay callable directly as
13//! well, and the tool's description says when to use which.
14//!
15//! # Every call is a governed call
16//!
17//! A program cannot reach around the session's governance. Each call it makes
18//! pauses the interpreter and runs down the same path as a direct call: the
19//! phase's tool list, the decision gate, the flow gate (with extraction of the
20//! turn in progress), middleware, confirmation, the tool itself, and the flow
21//! advance on its result. A refused call raises `RuntimeError` in the program
22//! with the reason, and every call and its full result is recorded as a tool
23//! execution. A program cannot write state or move the flow except through
24//! the calls it makes.
25//!
26//! Commit tools (a booking, a payment) are better left out of the functions:
27//! a commit should be one visible call, not a line inside a program.
28//!
29//! # Sandboxes
30//!
31//! [`InProcess`] runs [Monty](https://github.com/pydantic/monty), a Python
32//! interpreter written in Rust, on a blocking thread of this process. It is
33//! fast and needs nothing installed, but it is not crash-isolated: a program
34//! that asks for a huge allocation (`'a' * 10**10`) can abort the process or
35//! run past its time limit while the allocation completes. Use it for
36//! development, tests and trusted deployments. [`MontyPool`] (feature
37//! `code-pool`) runs programs in Monty worker processes with a memory cap and
38//! hard deadlines, so a runaway program kills only its worker.
39//!
40//! ```rust,no_run
41//! use gemini_adk_rs::code_mode::{CodeLimits, CodeMode};
42//!
43//! let code = CodeMode::new()
44//!     .function("lookup_payment_history")
45//!     .function("check_availability")
46//!     .limits(CodeLimits {
47//!         calls: 10,
48//!         ..CodeLimits::default()
49//!     });
50//! ```
51
52mod sandbox;
53mod stubs;
54
55#[cfg(feature = "code-mode")]
56mod monty;
57#[cfg(feature = "code-pool")]
58mod pool;
59
60use std::collections::BTreeSet;
61use std::sync::Arc;
62use std::time::Duration;
63
64use gemini_genai_rs::prelude::FunctionDeclaration;
65use serde_json::{Value, json};
66
67#[cfg(feature = "code-mode")]
68pub use monty::InProcess;
69#[cfg(feature = "code-pool")]
70pub use pool::MontyPool;
71pub use sandbox::{CodeSandbox, Execution, Suspension, Unavailable};
72pub use stubs::{arguments, parameters, python_stub, without_redefinitions};
73
74/// The default name of the code tool.
75pub const DEFAULT_TOOL: &str = "run_code";
76
77/// Bounds on one program.
78#[derive(Debug, Clone, Copy, PartialEq, Eq)]
79pub struct CodeLimits {
80    /// Function calls one program may make. The call past the limit ends it.
81    pub calls: usize,
82    /// Interpreter time, summed over the program and excluding the time its
83    /// calls take. A program over it ends with `TimeoutError`.
84    pub cpu: Duration,
85    /// Wall-clock time for the whole program, its calls included. The caller
86    /// is waiting on the line for all of it.
87    pub wall: Duration,
88    /// Characters of the program's result (as JSON) and of its printed
89    /// output that the model reads. A larger result is refused with a note
90    /// to return less, rather than cut mid-value.
91    pub output_chars: usize,
92    /// Memory a program may allocate, in bytes. Enforced by
93    /// [`MontyPool`](self#sandboxes), whose workers have their own allocator;
94    /// the in-process sandbox cannot bound memory.
95    pub memory: usize,
96}
97
98impl Default for CodeLimits {
99    fn default() -> Self {
100        Self {
101            calls: 20,
102            cpu: Duration::from_secs(2),
103            wall: Duration::from_secs(30),
104            output_chars: 2_000,
105            memory: 64 * 1024 * 1024,
106        }
107    }
108}
109
110/// A function a program can call: a tool of the session, by name.
111#[derive(Debug, Clone, PartialEq, Eq)]
112pub struct CodeFunction {
113    /// The tool's name, which is also the function's name in the program.
114    pub name: String,
115    /// What the tool returns, shown in the function's docstring, such as
116    /// `{"payments": [{"paid": str, "due": str, "amount": float}]}`. Without
117    /// it the model guesses the result's keys, and a wrong guess costs it a
118    /// failed program.
119    pub returns: Option<String>,
120    /// Whether the tool stays declared on its own too; `None` follows
121    /// [`CodeMode::direct`]. A tool that returns many rows is better reached
122    /// only through programs: given it directly, Gemini 3.8 Live fetched 120
123    /// rows and then pasted them into a program.
124    pub direct: Option<bool>,
125}
126
127/// The code tool of a session: its name, the tools a program can call, its
128/// limits and the sandbox it runs in. See the [module docs](self).
129#[derive(Clone)]
130pub struct CodeMode {
131    tool: String,
132    functions: Vec<CodeFunction>,
133    limits: CodeLimits,
134    direct: bool,
135    sandbox: Arc<dyn CodeSandbox>,
136}
137
138impl std::fmt::Debug for CodeMode {
139    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
140        f.debug_struct("CodeMode")
141            .field("tool", &self.tool)
142            .field("functions", &self.functions)
143            .field("limits", &self.limits)
144            .field("direct", &self.direct)
145            .finish_non_exhaustive()
146    }
147}
148
149impl Default for CodeMode {
150    fn default() -> Self {
151        Self::new()
152    }
153}
154
155impl CodeMode {
156    /// A code tool named `run_code` with no functions yet, the default
157    /// limits, and the in-process sandbox (or, built without the `code-mode`
158    /// feature, a sandbox that refuses every program).
159    pub fn new() -> Self {
160        Self {
161            tool: DEFAULT_TOOL.to_string(),
162            functions: Vec::new(),
163            limits: CodeLimits::default(),
164            direct: true,
165            sandbox: default_sandbox(),
166        }
167    }
168
169    /// Let programs call the tool `name`.
170    #[must_use]
171    pub fn function(mut self, name: impl Into<String>) -> Self {
172        self.push(CodeFunction {
173            name: name.into(),
174            returns: None,
175            direct: None,
176        });
177        self
178    }
179
180    /// Let programs call the tool `name`, and declare it only through the
181    /// code tool, whatever [`direct`](Self::direct) says. For a tool whose
182    /// result the model should never read whole, such as a long history.
183    #[must_use]
184    pub fn only_in_code(mut self, name: impl Into<String>) -> Self {
185        self.push(CodeFunction {
186            name: name.into(),
187            returns: None,
188            direct: Some(false),
189        });
190        self
191    }
192
193    /// Let programs call the tool `name`, documenting what it returns.
194    #[must_use]
195    pub fn function_returning(
196        mut self,
197        name: impl Into<String>,
198        returns: impl Into<String>,
199    ) -> Self {
200        self.push(CodeFunction {
201            name: name.into(),
202            returns: Some(returns.into()),
203            direct: None,
204        });
205        self
206    }
207
208    /// Let programs call each of `names`.
209    #[must_use]
210    pub fn functions<I, S>(mut self, names: I) -> Self
211    where
212        I: IntoIterator<Item = S>,
213        S: Into<String>,
214    {
215        for name in names {
216            self = self.function(name);
217        }
218        self
219    }
220
221    /// Name the code tool `name` instead of `run_code`.
222    #[must_use]
223    pub fn tool_name(mut self, name: impl Into<String>) -> Self {
224        self.tool = name.into();
225        self
226    }
227
228    /// Bound each program by `limits`.
229    #[must_use]
230    pub fn limits(mut self, limits: CodeLimits) -> Self {
231        self.limits = limits;
232        self
233    }
234
235    /// Whether the functions stay declared as tools of their own (the
236    /// default). With `false` the model reaches them only through programs.
237    #[must_use]
238    pub fn direct(mut self, direct: bool) -> Self {
239        self.direct = direct;
240        self
241    }
242
243    /// Run programs in `sandbox`.
244    #[must_use]
245    pub fn sandbox(mut self, sandbox: Arc<dyn CodeSandbox>) -> Self {
246        self.sandbox = sandbox;
247        self
248    }
249
250    fn push(&mut self, function: CodeFunction) {
251        match self.functions.iter_mut().find(|f| f.name == function.name) {
252            Some(existing) => {
253                if function.returns.is_some() {
254                    existing.returns = function.returns;
255                }
256                if function.direct.is_some() {
257                    existing.direct = function.direct;
258                }
259            }
260            None => self.functions.push(function),
261        }
262    }
263
264    /// The code tool's name.
265    pub fn tool(&self) -> &str {
266        &self.tool
267    }
268
269    /// The functions programs can call.
270    pub fn function_list(&self) -> &[CodeFunction] {
271        &self.functions
272    }
273
274    /// The limits on each program.
275    pub fn limit(&self) -> &CodeLimits {
276        &self.limits
277    }
278
279    /// Whether the functions stay declared as tools of their own.
280    pub fn is_direct(&self) -> bool {
281        self.direct
282    }
283
284    /// The sandbox programs run in.
285    pub fn sandbox_ref(&self) -> &Arc<dyn CodeSandbox> {
286        &self.sandbox
287    }
288
289    /// Whether programs can call the tool `name`.
290    pub fn exposes(&self, name: &str) -> bool {
291        self.functions.iter().any(|f| f.name == name)
292    }
293
294    /// Whether the model reaches the tool `name` only through programs.
295    pub fn hides(&self, name: &str) -> bool {
296        self.functions
297            .iter()
298            .find(|f| f.name == name)
299            .is_some_and(|f| !f.direct.unwrap_or(self.direct))
300    }
301
302    /// The code tool's declaration, listing the functions in `catalog` that
303    /// programs can call and `offered` admits. `None` when there is none,
304    /// in which case the code tool is not declared either.
305    pub fn declaration<'a>(
306        &self,
307        catalog: impl IntoIterator<Item = &'a FunctionDeclaration>,
308        offered: impl Fn(&str) -> bool,
309    ) -> Option<FunctionDeclaration> {
310        let catalog: Vec<&FunctionDeclaration> = catalog.into_iter().collect();
311        let listed: Vec<&CodeFunction> = self
312            .functions
313            .iter()
314            .filter(|f| offered(&f.name))
315            .filter(|f| catalog.iter().any(|d| d.name == f.name))
316            .collect();
317        let stubs: Vec<String> = listed
318            .iter()
319            .filter_map(|f| {
320                catalog
321                    .iter()
322                    .find(|d| d.name == f.name)
323                    .map(|d| python_stub(d, f.returns.as_deref()))
324            })
325            .collect();
326        if stubs.is_empty() {
327            return None;
328        }
329        let when = if listed.iter().any(|f| !self.hides(&f.name)) {
330            "Use it whenever the answer needs arithmetic, dates, counting, sorting \
331             or comparing over what a function returns, or several calls: call \
332             the function inside the program rather than calling the tool first \
333             and copying its result in. Call a tool directly only when its result \
334             is the answer as it stands."
335        } else {
336            "Use it to call the functions below and to work out arithmetic, \
337             dates, counting or sorting over what they return."
338        };
339        Some(FunctionDeclaration {
340            name: self.tool.clone(),
341            description: format!(
342                "Run a short Python program and get back the value of its last \
343                 expression, and anything it prints. {when} Do the whole request \
344                 in one program. Each function's docstring says what it returns: \
345                 rely on it rather than printing results to look at them. End the \
346                 program with an expression holding only the values you need to \
347                 answer. The functions are already defined: call them directly; \
348                 never import, define or mock them. Keep judgments of meaning out \
349                 of the program: when a choice depends on what something is (which \
350                 medication treats a condition), return the candidates and decide \
351                 yourself. date.weekday() counts Monday as 0. A call the conversation does \
352                 not allow yet raises RuntimeError with the reason: tell the \
353                 caller what is needed instead of retrying. The program can \
354                 import datetime, math, json and re.\n\n{}",
355                stubs.join("\n")
356            ),
357            parameters: Some(json!({
358                "type": "object",
359                "properties": {
360                    "code": {"type": "string", "description": "The Python program."}
361                },
362                "required": ["code"],
363            })),
364            behavior: None,
365        })
366    }
367
368    /// The function names of `catalog` that the code tool's declaration
369    /// would list when `offered` admits them, sorted.
370    pub fn offered_functions<'a>(
371        &self,
372        catalog: impl IntoIterator<Item = &'a FunctionDeclaration>,
373        offered: impl Fn(&str) -> bool,
374    ) -> BTreeSet<String> {
375        catalog
376            .into_iter()
377            .filter(|d| self.exposes(&d.name) && offered(&d.name))
378            .map(|d| d.name.clone())
379            .collect()
380    }
381
382    /// What the model reads for a program that ended with `result` and
383    /// printed `printed`: `{"result": ..}` or `{"error": ..}`, with `printed`
384    /// when it printed anything.
385    pub fn response(&self, result: Result<Value, String>, printed: &str) -> Value {
386        let limit = self.limits.output_chars;
387        let mut response = match result {
388            Ok(value) => {
389                let size = value.to_string().chars().count();
390                if size > limit {
391                    json!({
392                        "error": format!(
393                            "The program's result is {size} characters, over the \
394                             limit of {limit}. End it with only the values you \
395                             need to answer."
396                        )
397                    })
398                } else {
399                    json!({ "result": value })
400                }
401            }
402            Err(error) => json!({ "error": error }),
403        };
404        // A printout over the limit is refused whole: cut, it reads as the
405        // full data, and in live runs the model answered from what was left.
406        let printed = printed.trim();
407        let size = printed.chars().count();
408        if size > limit {
409            return json!({
410                "error": format!(
411                    "The program printed {size} characters, over the limit of {limit}, \
412                     so none of it is shown. Print only the values you need: compute \
413                     the answer in the program."
414                )
415            });
416        }
417        if !printed.is_empty() {
418            response["printed"] = Value::String(printed.to_string());
419        }
420        response
421    }
422}
423
424fn default_sandbox() -> Arc<dyn CodeSandbox> {
425    #[cfg(feature = "code-mode")]
426    {
427        Arc::new(InProcess)
428    }
429    #[cfg(not(feature = "code-mode"))]
430    {
431        Arc::new(Unavailable)
432    }
433}
434
435impl CodeMode {
436    /// Start `program` in the sandbox, with the functions this mode exposes
437    /// defined as globals and named by their declarations in `catalog`.
438    ///
439    /// The caller drives it: while [`Program::call`] names a call, run that
440    /// call (down the governed path, in a session) and pass its result to
441    /// [`Program::answer`]; then [`Program::response`] is what the model
442    /// reads. Dropping a [`Program`] before it finishes cancels it.
443    pub async fn start(&self, program: &str, catalog: &[FunctionDeclaration]) -> Program {
444        let names: Vec<String> = self.functions.iter().map(|f| f.name.clone()).collect();
445        let declarations = catalog
446            .iter()
447            .filter(|d| self.exposes(&d.name))
448            .cloned()
449            .collect();
450        let (execution, first) = match without_redefinitions(program, &names) {
451            Ok(source) => self.sandbox.start(&source, &names, &self.limits).await,
452            Err(reason) => (sandbox::ended(), Suspension::failed(reason)),
453        };
454        let mut program = Program {
455            mode: self.clone(),
456            execution,
457            declarations,
458            pending: None,
459            finished: None,
460            calls: Vec::new(),
461        };
462        program.advance(first).await;
463        program
464    }
465}
466
467/// A program running in a sandbox, suspended at a call or finished. See
468/// [`CodeMode::start`].
469pub struct Program {
470    mode: CodeMode,
471    execution: Box<dyn Execution>,
472    declarations: Vec<FunctionDeclaration>,
473    /// The call the program waits on: the tool and its named arguments.
474    pending: Option<(String, Value)>,
475    /// The program's end: its value or error, and what it printed.
476    finished: Option<(Result<Value, String>, String)>,
477    /// The calls answered so far: tool name and whether it ran.
478    calls: Vec<(String, bool)>,
479}
480
481impl std::fmt::Debug for Program {
482    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
483        f.debug_struct("Program")
484            .field("pending", &self.pending)
485            .field("finished", &self.finished)
486            .field("calls", &self.calls)
487            .finish_non_exhaustive()
488    }
489}
490
491impl Program {
492    /// The call the program waits on: the tool's name and its arguments,
493    /// named as its declaration names them. `None` once it has finished.
494    pub fn call(&self) -> Option<(&str, &Value)> {
495        self.pending
496            .as_ref()
497            .map(|(name, args)| (name.as_str(), args))
498    }
499
500    /// Answer the pending call with the tool's result, or with the reason it
501    /// failed or was refused (raised in the program as `RuntimeError`), and
502    /// run the program to its next call or its end.
503    pub async fn answer(&mut self, result: Result<Value, String>) {
504        let Some((name, _)) = self.pending.take() else {
505            return;
506        };
507        self.calls.push((name, result.is_ok()));
508        let next = self.execution.resume(result).await;
509        self.advance(next).await;
510    }
511
512    /// End the program with `error`, as when its time runs out.
513    pub fn abort(&mut self, error: impl Into<String>) {
514        self.pending = None;
515        self.finished = Some((Err(error.into()), String::new()));
516    }
517
518    /// Whether the program has finished.
519    pub fn is_finished(&self) -> bool {
520        self.finished.is_some()
521    }
522
523    /// The calls the program made, in order: tool name and whether it ran.
524    pub fn calls(&self) -> &[(String, bool)] {
525        &self.calls
526    }
527
528    /// What the model reads: `{"result": ..}` or `{"error": ..}`, with
529    /// `printed` when the program printed anything. `None` until it has
530    /// finished.
531    pub fn response(&self) -> Option<Value> {
532        self.finished
533            .as_ref()
534            .map(|(result, printed)| self.mode.response(result.clone(), printed))
535    }
536
537    /// Follow suspensions until the program waits on a call this session
538    /// should run, or ends. Calls with arguments that do not fit the
539    /// declaration are answered here with the reason, as Python would.
540    async fn advance(&mut self, mut next: Suspension) {
541        loop {
542            match next {
543                Suspension::Finished { result, printed } => {
544                    self.finished = Some((result, printed));
545                    return;
546                }
547                Suspension::Call {
548                    function,
549                    positional,
550                    keywords,
551                } => {
552                    let limit = self.mode.limits.calls;
553                    if self.calls.len() >= limit {
554                        self.abort(format!(
555                            "The program made more than {limit} function calls. \
556                             Do the work in fewer calls."
557                        ));
558                        return;
559                    }
560                    let named = match self.declarations.iter().find(|d| d.name == function) {
561                        Some(declaration) => arguments(declaration, positional, keywords),
562                        None => Err(format!("NameError: name '{function}' is not defined")),
563                    };
564                    match named {
565                        Ok(args) => {
566                            self.pending = Some((function, args));
567                            return;
568                        }
569                        Err(reason) => {
570                            self.calls.push((function, false));
571                            next = self.execution.resume(Err(reason)).await;
572                        }
573                    }
574                }
575            }
576        }
577    }
578}
579
580#[cfg(test)]
581mod tests {
582    use super::*;
583
584    #[test]
585    fn an_over_limit_printout_is_refused_whole() {
586        let mode = CodeMode::new().limits(CodeLimits {
587            output_chars: 10,
588            ..CodeLimits::default()
589        });
590        let response = mode.response(Ok(Value::Null), "0123456789abc");
591        assert!(
592            response.get("printed").is_none(),
593            "nothing partial is shown"
594        );
595        let error = response["error"].as_str().unwrap();
596        assert!(error.contains("printed 13 characters"), "{error}");
597        assert_eq!(
598            mode.response(Ok(json!(1)), "short"),
599            json!({ "result": 1, "printed": "short" })
600        );
601    }
602}
gemini_adk_rs/code_mode/mod.rs

gemini_adk_rs/code_mode/
mod.rs
