run_code_rmcp/python_runner/
python_runner.rs

1//通过 uv 命令,来运行 python脚本
2use crate::{
3    cache::CodeFileCache,
4    model::{
5        CodeExecutor, CodeScriptExecutionResult, CommandExecutor, LanguageScript, RunCode,
6        TokioHeapSize,
7    },
8    python_runner::parse_import,
9};
10use anyhow::{Context, Result};
11use log::{debug, error, info, warn};
12use serde_json;
13use tokio::process::Command;
14
15#[derive(Default)]
16pub struct PythonRunner;
17
18impl RunCode for PythonRunner {
19    async fn run_with_params(
20        &self,
21        code: &str,
22        params: Option<serde_json::Value>,
23        timeout_seconds: Option<u64>,
24    ) -> Result<CodeScriptExecutionResult> {
25        debug!("开始执行Python脚本...,执行参数: {:?}", params);
26        // 根据 code ,获取对应的hash, 对用户脚本代码,使用胶水代码处理后,缓存到文件系统里,下次使用如果hash相同,直接使用
27        let hash = CodeFileCache::obtain_code_hash(code);
28        let cache_exist =
29            CodeFileCache::check_code_file_cache_exisht(&hash, &LanguageScript::Python).await;
30
31        let run_code_script_file_tuple = if cache_exist {
32            // 从缓存中读取代码
33            let cache_code =
34                CodeFileCache::get_code_file_cache(&hash, &LanguageScript::Python).await;
35            debug!("从缓存中读取代码:hash值 {:?}", &hash);
36            cache_code?
37        } else {
38            // 分析用户python代码依赖
39            let dependencies = parse_import(code)?;
40            // 准备Python代码，添加日志捕获和handler函数执行逻辑
41            let wrapped_code = self.prepare_python_code(code, true);
42            // 缓存代码
43            CodeFileCache::save_code_file_cache(&hash, &wrapped_code, &LanguageScript::Python)
44                .await?;
45
46            //保存后，获取文件
47            let code_script_file_tuple =
48                CodeFileCache::get_code_file_cache(&hash, &LanguageScript::Python).await?;
49            let run_code_script_file_path = code_script_file_tuple.1.clone();
50            // 按照 uv命令的规范,添加用户脚本所需的依赖
51            // uv add --script example.py 'requests<3' 'rich'  这是添加依赖的参考命令,dependencies 是解析出来的依赖列表
52
53            // 执行添加依赖命令
54            if !dependencies.is_empty() {
55                info!("正在添加依赖: {:?}", dependencies);
56                let mut cmd = Command::new("uv");
57                cmd.arg("add")
58                    .arg("--script")
59                    .arg(&run_code_script_file_path);
60
61                // 为每个依赖添加一个参数
62                for dep in &dependencies {
63                    cmd.arg(dep);
64                }
65                // 打印 cmd 命令,可以直接复制执行的命令字符串
66                let cmd_str = format!("{:?}", &cmd);
67                info!("uv命令字符串: {}", cmd_str);
68
69                let cmd_output = match cmd.kill_on_drop(true).output().await {
70                    Ok(output) => output,
71                    Err(e) => {
72                        error!("安装Python依赖失败: {:?}", e);
73                        error!("失败的命令: {:?}", cmd);
74                        return Err(e).context("Failed to add dependencies with uv");
75                    }
76                };
77
78                let stdout = String::from_utf8_lossy(&cmd_output.stdout).to_string();
79                let stderr = String::from_utf8_lossy(&cmd_output.stderr).to_string();
80                info!("添加依赖结果 - stdout: {}", stdout);
81                info!("添加依赖结果 - stderr: {}", stderr);
82
83                if !cmd_output.status.success() {
84                    warn!("添加依赖失败，状态码: {}", cmd_output.status);
85                }
86            }
87            debug!("创建脚本缓存:hash值 {:?}", &hash);
88            code_script_file_tuple
89        };
90
91        let temp_path = run_code_script_file_tuple.1;
92
93        // 将参数序列化为JSON字符串
94        let params_json = match params {
95            Some(p) => serde_json::to_string(&p)?,
96            None => "{}".to_string(),
97        };
98
99        // 使用uv run命令执行Python脚本，提供隔离环境
100        let mut execute_command = Command::new("uv");
101        execute_command
102            .arg("run")
103            .arg("-s") // 明确指定作为脚本运行
104            .arg("-p")
105            .arg("3.13") // 指定Python解释器版本3.13
106            .env("INPUT_JSON", &params_json) // 通过环境变量传递参数
107            .arg(&temp_path)
108            .kill_on_drop(true);
109
110        info!("执行命令: {:?}", &execute_command);
111
112        // let tokio_child_command = TokioHeapSize::default();
113        // // 设置堆大小限制
114        // tokio_child_command
115        //     .with_heap_limit(&mut execute_command)
116        //     .await;
117
118        //限制command 的执行超时时间
119        let executor = match timeout_seconds {
120            Some(timeout) => CommandExecutor::with_timeout(execute_command.output(), timeout),
121            None => CommandExecutor::default(execute_command.output()),
122        };
123
124        let executor_result = executor.await;
125        let output = match executor_result {
126            Ok(cmd_result) => match cmd_result {
127                Ok(output) => output,
128                Err(e) => {
129                    error!("Python命令执行失败: {:?}", e);
130                    return Err(e.into());
131                }
132            },
133            Err(e) => {
134                error!("Python任务执行异常: {:?}", e);
135                return Err(e.into());
136            }
137        };
138        // 调试输出
139        let stdout = String::from_utf8_lossy(&output.stdout).to_string();
140        let stderr = String::from_utf8_lossy(&output.stderr).to_string();
141        debug!("Python stdout: {}", stdout);
142        debug!("Python stderr: {}", stderr);
143
144        // 解析输出
145        CodeExecutor::parse_execution_output(&output.stdout, &output.stderr).await
146    }
147}
148
149impl PythonRunner {
150    /// 准备Python代码，添加日志捕获和handler函数执行逻辑
151    fn prepare_python_code(&self, code: &str, show_logs: bool) -> String {
152        let show_logs_value = if show_logs { "True" } else { "False" };
153
154        let template = include_str!("../templates/python_template.py");
155
156        template
157            .replace("{{USER_CODE}}", code)
158            .replace("{{SHOW_LOGS}}", show_logs_value)
159    }
160}
run_code_rmcp/python_runner/python_runner.rs

run_code_rmcp/python_runner/
python_runner.rs
