Skip to main content

aagt_core/skills/tool/
code_interpreter.rs

1//! Code Interpreter Tool
2
3use async_trait::async_trait;
4use serde::{Deserialize, Serialize};
5use schemars::JsonSchema;
6use std::sync::Arc;
7use tokio::sync::Mutex;
8
9use crate::skills::tool::Tool;
10use crate::error::Error;
11use crate::skills::capabilities::Sidecar;
12
13/// Arguments for the Code Interpreter tool
14#[derive(Debug, Deserialize, Serialize, JsonSchema)]
15pub struct CodeArgs {
16    /// The Python code to execute
17    pub code: String,
18}
19
20/// A tool that executes Python code in a stateful sidecar
21pub struct CodeInterpreter {
22    sidecar: Arc<Mutex<Sidecar>>,
23}
24
25impl CodeInterpreter {
26    /// Create a new CodeInterpreter connected to the given sidecar
27    pub fn new(sidecar: Arc<Mutex<Sidecar>>) -> Self {
28        Self { sidecar }
29    }
30}
31
32#[async_trait]
33impl Tool for CodeInterpreter {
34    fn name(&self) -> String {
35        "code_interpreter".to_string()
36    }
37
38    async fn definition(&self) -> crate::skills::tool::ToolDefinition {
39        crate::skills::tool::ToolDefinition {
40            name: self.name(),
41            description: "Executes Python code in a stateful shell. Use this for data analysis, math, and plotting.".to_string(),
42            parameters: serde_json::json!({
43
44                "type": "object",
45                "properties": {
46                    "code": {
47                        "type": "string",
48                        "description": "Python code to execute"
49                    }
50                },
51                "required": ["code"]
52            }),
53            parameters_ts: Some("interface CodeArgs {\n  code: string; // Python code to execute\n}".to_string()),
54            is_binary: false,
55            is_verified: true,
56            usage_guidelines: Some("Use this for data analysis, complex mathematical calculations, or file processing that cannot be done with simple search. Avoid using it for simple string manipulations.".to_string()),
57        }
58    }
59
60    async fn call(&self, arguments: &str) -> anyhow::Result<String> {
61        let args: CodeArgs = serde_json::from_str(arguments)
62            .map_err(|e| Error::ToolArguments {
63                tool_name: self.name(),
64                message: format!("Invalid JSON arguments: {}", e),
65            })?;
66
67        let mut sidecar = self.sidecar.lock().await;
68        let result = sidecar.execute(args.code).await?;
69
70        let mut output = result.stdout;
71        if !result.stderr.is_empty() {
72            output.push_str("\n--- Stderr ---\n");
73            output.push_str(&result.stderr);
74        }
75
76        if !result.images.is_empty() {
77            output.push_str(&format!("\n(Note: Generated {} image(s))", result.images.len()));
78            // In a real scenario, we might want to save these to files or return them as part of multi-modal message
79        }
80
81        Ok(output)
82    }
83}