gemini_adk_fluent_rs/spec/
code.rs

1//! `code`: a tool through which the model runs a short Python program whose
2//! functions are the spec's tools.
3//!
4//! ```json
5//! "tools": [
6//!   { "name": "lookup_payment_history",
7//!     "returns": { "type": "object", "properties": { "payments": { "type": "array",
8//!       "items": { "type": "object", "properties": {
9//!         "paid": { "type": "string" }, "due": { "type": "string" }, "amount": { "type": "number" } } } } } } },
10//!   { "name": "book_table" }
11//! ],
12//! "code": { "functions": ["lookup_payment_history"] }
13//! ```
14//!
15//! Without `functions`, a program can call every declared tool except the
16//! commits (a conversation stage's `commit`, a flow's `commit`) and the
17//! background tools. Each call it makes is governed like a direct call.
18//! Lowered to a [`CodeMode`]; programs run in
19//! [`SpecResources::code_sandbox`](super::SpecResources::code_sandbox), or in
20//! this process when it is unset.
21
22use std::collections::BTreeSet;
23use std::sync::Arc;
24use std::time::Duration;
25
26use serde::{Deserialize, Serialize};
27use serde_json::Value;
28
29use gemini_adk_rs::code_mode::{CodeLimits, CodeMode, CodeSandbox, DEFAULT_TOOL};
30use gemini_adk_rs::flow::Constraint;
31
32use super::{SessionSpec, ToolSpec};
33
34/// The code tool of a spec.
35#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
36pub struct CodeSpec {
37    /// The code tool's name. Default `run_code`.
38    #[serde(default, skip_serializing_if = "Option::is_none")]
39    pub tool: Option<String>,
40    /// The declared tools a program can call. Empty: every declared tool
41    /// except commits and background tools.
42    #[serde(default, skip_serializing_if = "Vec::is_empty")]
43    pub functions: Vec<String>,
44    /// Keep the functions declared as tools of their own too, so the model
45    /// can make a single lookup without writing a program. Default `true`.
46    #[serde(default = "yes", skip_serializing_if = "is_yes")]
47    pub direct: bool,
48    /// Tools a program can call that are declared only through the code
49    /// tool, whatever `direct` says: those whose results the model should
50    /// not read whole, such as a long history. Listed tools join
51    /// `functions`.
52    #[serde(default, skip_serializing_if = "Vec::is_empty")]
53    pub only_in_code: Vec<String>,
54    /// Bounds on each program.
55    #[serde(default, skip_serializing_if = "Option::is_none")]
56    pub limits: Option<CodeLimitsSpec>,
57}
58
59impl Default for CodeSpec {
60    fn default() -> Self {
61        Self {
62            tool: None,
63            functions: Vec::new(),
64            direct: true,
65            only_in_code: Vec::new(),
66            limits: None,
67        }
68    }
69}
70
71fn yes() -> bool {
72    true
73}
74
75#[allow(
76    clippy::trivially_copy_pass_by_ref,
77    reason = "serde's skip_serializing_if passes a reference"
78)]
79fn is_yes(value: &bool) -> bool {
80    *value
81}
82
83/// Bounds on each program; unset fields keep the defaults.
84#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize, schemars::JsonSchema)]
85pub struct CodeLimitsSpec {
86    /// Function calls per program. Default 20.
87    #[serde(default, skip_serializing_if = "Option::is_none")]
88    pub calls: Option<usize>,
89    /// Interpreter time per program, excluding its calls, in ms. Default 2000.
90    #[serde(default, skip_serializing_if = "Option::is_none")]
91    pub cpu_ms: Option<u64>,
92    /// Wall-clock time per program, its calls included, in ms. Default 30000.
93    #[serde(default, skip_serializing_if = "Option::is_none")]
94    pub wall_ms: Option<u64>,
95    /// Characters of a program's result (and of its printed output) the
96    /// model reads. Default 2000.
97    #[serde(default, skip_serializing_if = "Option::is_none")]
98    pub output_chars: Option<usize>,
99    /// Memory per program in MiB, enforced by worker-process sandboxes.
100    /// Default 64.
101    #[serde(default, skip_serializing_if = "Option::is_none")]
102    pub memory_mb: Option<usize>,
103}
104
105impl CodeLimitsSpec {
106    /// The runtime limits, defaults filled in.
107    pub fn to_limits(&self) -> CodeLimits {
108        let base = CodeLimits::default();
109        CodeLimits {
110            calls: self.calls.unwrap_or(base.calls),
111            cpu: self.cpu_ms.map_or(base.cpu, Duration::from_millis),
112            wall: self.wall_ms.map_or(base.wall, Duration::from_millis),
113            output_chars: self.output_chars.unwrap_or(base.output_chars),
114            memory: self.memory_mb.map_or(base.memory, |mb| mb * 1024 * 1024),
115        }
116    }
117}
118
119impl SessionSpec {
120    /// The code tool's name, when the spec has one.
121    pub fn code_tool(&self) -> Option<&str> {
122        self.code
123            .as_ref()
124            .map(|code| code.tool.as_deref().unwrap_or(DEFAULT_TOOL))
125    }
126
127    /// The tools that commit: a conversation stage's `commit` (main flow and
128    /// overlays) and a flow's `commit` steps.
129    pub fn commit_tools(&self) -> BTreeSet<String> {
130        let mut commits = BTreeSet::new();
131        if let Some(conversation) = &self.conversation {
132            let stages = conversation
133                .stages
134                .iter()
135                .chain(conversation.overlays.iter().flat_map(|o| &o.stages));
136            commits.extend(
137                stages
138                    .filter_map(|stage| stage.commit.as_ref())
139                    .map(|commit| commit.tool.clone()),
140            );
141        } else if let Ok(flow) = self.effective_flow() {
142            commits.extend(flow.constraints.iter().filter_map(|c| match c {
143                Constraint::NeverUntil { tool, .. } => Some(tool.clone()),
144                _ => None,
145            }));
146        }
147        commits
148    }
149
150    /// The tools a program can call: `code.functions`, or every declared
151    /// tool except commits and background tools.
152    pub fn code_functions(&self) -> Vec<String> {
153        let Some(code) = &self.code else {
154            return Vec::new();
155        };
156        let mut functions: Vec<String> = if code.functions.is_empty() {
157            let commits = self.commit_tools();
158            self.tools
159                .iter()
160                .filter(|t| !commits.contains(&t.name) && !t.background && t.scheduling.is_none())
161                .map(|t| t.name.clone())
162                .collect()
163        } else {
164            code.functions.clone()
165        };
166        for name in &code.only_in_code {
167            if !functions.contains(name) {
168                functions.push(name.clone());
169            }
170        }
171        functions
172    }
173
174    /// The spec's code tool, running programs in `sandbox` (in this process
175    /// when `None`). `None` when the spec has no `code`.
176    pub fn compile_code_mode(&self, sandbox: Option<Arc<dyn CodeSandbox>>) -> Option<CodeMode> {
177        let code = self.code.as_ref()?;
178        let mut mode = CodeMode::new()
179            .tool_name(self.code_tool().unwrap_or(DEFAULT_TOOL))
180            .direct(code.direct)
181            .limits(code.limits.clone().unwrap_or_default().to_limits());
182        for name in self.code_functions() {
183            if code.only_in_code.contains(&name) {
184                mode = mode.only_in_code(name.clone());
185            }
186            mode = match self
187                .tools
188                .iter()
189                .find(|t| t.name == name)
190                .and_then(ToolSpec::returns_shape)
191            {
192                Some(shape) => mode.function_returning(name, shape),
193                None => mode.function(name),
194            };
195        }
196        if let Some(sandbox) = sandbox {
197            mode = mode.sandbox(sandbox);
198        }
199        Some(mode)
200    }
201
202    /// Errors and warnings for `code`.
203    pub(super) fn code_problems(&self) -> (Vec<String>, Vec<String>) {
204        let mut errors = Vec::new();
205        let mut warnings = Vec::new();
206        let Some(code) = &self.code else {
207            return (errors, warnings);
208        };
209        let tool = self.code_tool().unwrap_or(DEFAULT_TOOL);
210        if !self.skills.is_empty() {
211            errors
212                .push("code: a session of skills cannot have a session-wide code tool".to_string());
213        }
214        if self.tools.iter().any(|t| t.name == tool) {
215            errors.push(format!(
216                "code: the code tool's name '{tool}' is also a declared tool"
217            ));
218        }
219        for name in code.functions.iter().chain(&code.only_in_code) {
220            if !self.tools.iter().any(|t| &t.name == name) && self.mcp.is_empty() {
221                errors.push(format!("code: function '{name}' is not a declared tool"));
222            }
223        }
224        let functions = self.code_functions();
225        if functions.is_empty() && self.mcp.is_empty() {
226            warnings.push("code: no tool for a program to call".to_string());
227        }
228        let commits = self.commit_tools();
229        for name in functions.iter().filter(|n| commits.contains(*n)) {
230            warnings.push(format!(
231                "code: '{name}' commits; a commit is clearer as one visible call than as a \
232                 line in a program, so consider leaving it out of code.functions"
233            ));
234        }
235        for name in &functions {
236            if let Some(t) = self.tools.iter().find(|t| &t.name == name)
237                && (t.background || t.scheduling.is_some())
238            {
239                warnings.push(format!(
240                    "code: '{name}' is a background tool; in a program it runs inline, \
241                     since the program needs its result"
242                ));
243            }
244        }
245        (errors, warnings)
246    }
247}
248
249impl SessionSpec {
250    /// A skill's definition with the skill's `code`, its functions filled in
251    /// when unset: the skill's read tools that are not commits.
252    pub(super) fn with_skill_code(mut self, skill: &super::SkillSpec) -> Self {
253        let Some(mut code) = skill.code.clone() else {
254            return self;
255        };
256        if code.functions.is_empty() {
257            let commits = self.commit_tools();
258            code.functions = skill
259                .tools
260                .iter()
261                .filter(|t| matches!(t.effect, gemini_adk_rs::tasks::TaskToolEffect::Read))
262                .map(|t| t.tool.name.clone())
263                .filter(|name| !commits.contains(name))
264                .collect();
265        }
266        self.code = Some(code);
267        self
268    }
269}
270
271impl ToolSpec {
272    /// What the tool returns, as a program's docstring shows it: from
273    /// `returns`, or the shape of the canned `response` of a mock tool.
274    pub fn returns_shape(&self) -> Option<String> {
275        match (&self.returns, &self.response) {
276            (Some(schema), _) => Some(schema_shape(schema)),
277            (None, Some(example)) if self.http.is_none() && self.mcp.is_none() => {
278                Some(value_shape(example))
279            }
280            _ => None,
281        }
282    }
283}
284
285/// A JSON Schema as a compact shape: `{"payments": [{"amount": float}]}`.
286pub fn schema_shape(schema: &Value) -> String {
287    let kind = schema
288        .get("type")
289        .and_then(Value::as_str)
290        .map(str::to_ascii_lowercase);
291    match kind.as_deref() {
292        Some("object") => match schema.get("properties").and_then(Value::as_object) {
293            Some(properties) if !properties.is_empty() => format!(
294                "{{{}}}",
295                properties
296                    .iter()
297                    .map(|(k, v)| format!("\"{k}\": {}", schema_shape(v)))
298                    .collect::<Vec<_>>()
299                    .join(", ")
300            ),
301            _ => "dict".to_string(),
302        },
303        Some("array") => match schema.get("items") {
304            Some(items) => format!("[{}]", schema_shape(items)),
305            None => "list".to_string(),
306        },
307        Some("string") => "str".to_string(),
308        Some("integer") => "int".to_string(),
309        Some("number") => "float".to_string(),
310        Some("boolean") => "bool".to_string(),
311        _ => "object".to_string(),
312    }
313}
314
315/// The shape of an example value: `{"times": [str]}`.
316pub fn value_shape(value: &Value) -> String {
317    match value {
318        Value::Null => "None".to_string(),
319        Value::Bool(_) => "bool".to_string(),
320        Value::Number(n) if n.is_i64() || n.is_u64() => "int".to_string(),
321        Value::Number(_) => "float".to_string(),
322        Value::String(_) => "str".to_string(),
323        Value::Array(items) => match items.first() {
324            Some(first) => format!("[{}]", value_shape(first)),
325            None => "list".to_string(),
326        },
327        Value::Object(map) => format!(
328            "{{{}}}",
329            map.iter()
330                .map(|(k, v)| format!("\"{k}\": {}", value_shape(v)))
331                .collect::<Vec<_>>()
332                .join(", ")
333        ),
334    }
335}
336
337#[cfg(test)]
338mod tests {
339    use super::*;
340    use serde_json::json;
341
342    fn spec(doc: Value) -> SessionSpec {
343        serde_json::from_value(doc).expect("spec parses")
344    }
345
346    fn booking() -> Value {
347        json!({
348            "name": "booking",
349            "tools": [
350                { "name": "check_availability",
351                  "response": { "date": "2026-10-12", "times": ["19:00", "20:30"] } },
352                { "name": "lookup_payment_history",
353                  "returns": { "type": "object", "properties": { "payments": {
354                      "type": "array", "items": { "type": "object", "properties": {
355                          "amount": { "type": "number" }, "paid": { "type": "string" } } } } } } },
356                { "name": "book_table" },
357                { "name": "notify", "background": true }
358            ],
359            "conversation": { "name": "book", "stages": [
360                { "id": "confirm", "commit": { "tool": "book_table", "when": { "decided": "confirmed" } } }
361            ] },
362            "code": {}
363        })
364    }
365
366    #[test]
367    fn by_default_a_program_calls_every_tool_but_commits_and_background_tools() {
368        let spec = spec(booking());
369        assert_eq!(
370            spec.code_functions(),
371            ["check_availability", "lookup_payment_history"]
372        );
373        let mode = spec.compile_code_mode(None).expect("code mode");
374        assert_eq!(mode.tool(), "run_code");
375        let returns: Vec<Option<&str>> = mode
376            .function_list()
377            .iter()
378            .map(|f| f.returns.as_deref())
379            .collect();
380        assert_eq!(
381            returns,
382            [
383                Some(r#"{"date": str, "times": [str]}"#),
384                Some(r#"{"payments": [{"amount": float, "paid": str}]}"#),
385            ]
386        );
387    }
388
389    #[test]
390    fn listing_a_commit_warns_and_an_unknown_function_is_an_error() {
391        let mut doc = booking();
392        doc["code"] = json!({ "functions": ["book_table", "refund"] });
393        let (errors, warnings) = spec(doc).code_problems();
394        assert_eq!(errors, ["code: function 'refund' is not a declared tool"]);
395        assert!(warnings[0].contains("'book_table' commits"), "{warnings:?}");
396    }
397
398    #[test]
399    fn the_code_tool_cannot_take_a_declared_tool_name() {
400        let mut doc = booking();
401        doc["code"] = json!({ "tool": "book_table" });
402        let (errors, _) = spec(doc).code_problems();
403        assert!(errors[0].contains("'book_table' is also a declared tool"));
404    }
405
406    #[test]
407    fn a_long_history_can_be_reached_only_in_code() {
408        let mut doc = booking();
409        doc["code"] = json!({ "functions": ["check_availability"],
410                              "only_in_code": ["lookup_payment_history"] });
411        let spec = spec(doc);
412        assert_eq!(
413            spec.code_functions(),
414            ["check_availability", "lookup_payment_history"]
415        );
416        let mode = spec.compile_code_mode(None).unwrap();
417        assert!(mode.hides("lookup_payment_history"));
418        assert!(!mode.hides("check_availability"));
419        assert!(
420            spec.to_rust()
421                .contains(".only_in_code(\"lookup_payment_history\")")
422        );
423    }
424
425    #[test]
426    fn limits_fill_in_defaults() {
427        let limits = CodeLimitsSpec {
428            calls: Some(5),
429            memory_mb: Some(16),
430            ..Default::default()
431        }
432        .to_limits();
433        assert_eq!(limits.calls, 5);
434        assert_eq!(limits.memory, 16 * 1024 * 1024);
435        assert_eq!(limits.cpu, CodeLimits::default().cpu);
436    }
437}
gemini_adk_fluent_rs/spec/code.rs

gemini_adk_fluent_rs/spec/
code.rs
