agentshield/adapter/
autogen.rs

1use std::path::Path;
2
3use once_cell::sync::Lazy;
4use regex::Regex;
5
6use crate::config::ScanPathFilter;
7use crate::error::Result;
8use crate::ir::*;
9
10static AUTOGEN_REGISTER_FN_RE: Lazy<Regex> = Lazy::new(|| {
11    Regex::new(r#"(?s)(?:register_function|@user_proxy\.register_for_execution|@assistant\.register_for_llm)\s*\(\s*(?:(?:f=)?(\w+))?(?:[^\)]*?description\s*=\s*["'`]([^"'`]+)["'`])?"#)
12        .expect("static regex pattern is valid")
13});
14
15static AUTOGEN_DEF_DOC_RE: Lazy<Regex> = Lazy::new(|| {
16    Regex::new(r#"(?m)^\s*(?:async\s+)?def\s+(\w+)\s*\([^)]*\)\s*(?:->[^:]+)?:\s*\n\s*["'`]([\s\S]*?)["'`]"#)
17        .expect("static regex pattern is valid")
18});
19
20/// Microsoft AutoGen framework adapter (`autogen` and `pyautogen`).
21///
22/// Detects AutoGen projects by looking for:
23/// - `pyproject.toml` / `requirements.txt` with `autogen` or `pyautogen`
24/// - Python files importing `from autogen` or `import autogen`
25pub struct AutoGenAdapter;
26
27impl super::Adapter for AutoGenAdapter {
28    fn framework(&self) -> Framework {
29        Framework::AutoGen
30    }
31
32    fn detect(&self, root: &Path) -> bool {
33        super::pipeline::detect_python_framework(
34            root,
35            &["autogen", "pyautogen"],
36            |trimmed| {
37                trimmed.starts_with("autogen")
38                    || trimmed.starts_with("pyautogen")
39                    || trimmed.starts_with("autogen-agentchat")
40                    || trimmed.starts_with("autogen-ext")
41            },
42            &[
43                "from autogen",
44                "import autogen",
45                "from autogen.agentchat",
46                "from autogen_agentchat",
47            ],
48        )
49    }
50
51    fn load(&self, root: &Path, ignore_tests: bool) -> Result<Vec<ScanTarget>> {
52        let filter = ScanPathFilter::for_ignore_tests(ignore_tests);
53        self.load_with_filter(root, &filter)
54    }
55
56    fn load_with_filter(&self, root: &Path, filter: &ScanPathFilter) -> Result<Vec<ScanTarget>> {
57        super::pipeline::load_framework_target(
58            root,
59            filter,
60            Framework::AutoGen,
61            "autogen-project",
62            |lang| lang == Language::Python,
63            Some(extract_autogen_tools_from_source),
64        )
65    }
66}
67
68pub(crate) fn extract_autogen_tools_from_source(path: &Path, content: &str) -> Vec<ToolSurface> {
69    let mut tools = Vec::new();
70    let mut doc_map = std::collections::HashMap::new();
71
72    for cap in AUTOGEN_DEF_DOC_RE.captures_iter(content) {
73        let func_name = cap[1].to_string();
74        let doc = cap[2].trim().to_string();
75        doc_map.insert(func_name, doc);
76    }
77
78    for cap in AUTOGEN_REGISTER_FN_RE.captures_iter(content) {
79        let func_name = cap
80            .get(1)
81            .map(|m| m.as_str().to_string())
82            .unwrap_or_default();
83        if func_name.is_empty() {
84            continue;
85        }
86
87        let explicit_desc = cap.get(2).map(|m| m.as_str().to_string());
88        let description = explicit_desc.or_else(|| doc_map.get(&func_name).cloned());
89
90        let match_start = cap.get(0).map(|m| m.start()).unwrap_or(0);
91        let line = content[..match_start].lines().count() + 1;
92
93        tools.push(super::pipeline::build_extracted_tool(
94            func_name,
95            description,
96            path,
97            line,
98        ));
99    }
100
101    tools
102}
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107    use crate::adapter::Adapter;
108
109    #[test]
110    fn detects_autogen_requirements() {
111        let dir = tempfile::tempdir().unwrap();
112        let req = dir.path().join("requirements.txt");
113        std::fs::write(&req, "pyautogen>=0.2.0\n").unwrap();
114
115        let adapter = AutoGenAdapter;
116        assert!(adapter.detect(dir.path()));
117    }
118
119    #[test]
120    fn detects_autogen_import() {
121        let dir = tempfile::tempdir().unwrap();
122        let src = dir.path().join("agent.py");
123        std::fs::write(&src, "from autogen import AssistantAgent\n").unwrap();
124
125        let adapter = AutoGenAdapter;
126        assert!(adapter.detect(dir.path()));
127    }
128
129    #[test]
130    fn extracts_autogen_register_function() {
131        let content = r#"
132from autogen import register_function
133
134def execute_command(cmd: str) -> str:
135    """Run a shell command on host."""
136    import subprocess
137    return subprocess.check_output(cmd, shell=True).decode()
138
139register_function(
140    execute_command,
141    caller=assistant,
142    executor=user_proxy,
143    name="execute_command",
144    description="Run shell commands securely"
145)
146"#;
147        let tools = extract_autogen_tools_from_source(Path::new("agent.py"), content);
148        assert_eq!(tools.len(), 1);
149        assert_eq!(tools[0].name, "execute_command");
150        assert_eq!(
151            tools[0].description.as_deref(),
152            Some("Run shell commands securely")
153        );
154    }
155}
agentshield/adapter/autogen.rs

agentshield/adapter/
autogen.rs
