agentshield/adapter/
llama_index.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 LLAMA_FUNCTION_TOOL_RE: Lazy<Regex> = Lazy::new(|| {
11    Regex::new(r#"(?s)FunctionTool\.from_defaults\s*\(\s*(?:fn\s*=\s*)?(\w+)(?:[^\)]*?name\s*=\s*["'`]([^"'`]+)["'`])?(?:[^\)]*?description\s*=\s*["'`]([^"'`]+)["'`])?"#)
12        .expect("static regex pattern is valid")
13});
14
15static LLAMA_TOOL_METADATA_RE: Lazy<Regex> = Lazy::new(|| {
16    Regex::new(r#"(?s)ToolMetadata\s*\(\s*(?:name\s*=\s*["'`]([^"'`]+)["'`])?[^\)]*?description\s*=\s*["'`]([^"'`]+)["'`]"#)
17        .expect("static regex pattern is valid")
18});
19
20static LLAMA_DEF_DOC_RE: Lazy<Regex> = Lazy::new(|| {
21    Regex::new(r#"(?m)^\s*(?:async\s+)?def\s+(\w+)\s*\([^)]*\)\s*(?:->[^:]+)?:\s*\n\s*["'`]([\s\S]*?)["'`]"#)
22        .expect("static regex pattern is valid")
23});
24
25/// LlamaIndex framework adapter (`llama-index` and `llama-index-core`).
26///
27/// Detects LlamaIndex projects by looking for:
28/// - `pyproject.toml` / `requirements.txt` with `llama-index` or `llama_index`
29/// - Python files importing `from llama_index` or `import llama_index`
30pub struct LlamaIndexAdapter;
31
32impl super::Adapter for LlamaIndexAdapter {
33    fn framework(&self) -> Framework {
34        Framework::LlamaIndex
35    }
36
37    fn detect(&self, root: &Path) -> bool {
38        super::pipeline::detect_python_framework(
39            root,
40            &["llama-index", "llama_index"],
41            |trimmed| {
42                trimmed.starts_with("llama-index")
43                    || trimmed.starts_with("llama_index")
44                    || trimmed.starts_with("llama-index-core")
45            },
46            &[
47                "from llama_index",
48                "import llama_index",
49                "from llama_index.core",
50            ],
51        )
52    }
53
54    fn load(&self, root: &Path, ignore_tests: bool) -> Result<Vec<ScanTarget>> {
55        let filter = ScanPathFilter::for_ignore_tests(ignore_tests);
56        self.load_with_filter(root, &filter)
57    }
58
59    fn load_with_filter(&self, root: &Path, filter: &ScanPathFilter) -> Result<Vec<ScanTarget>> {
60        super::pipeline::load_framework_target(
61            root,
62            filter,
63            Framework::LlamaIndex,
64            "llama-index-project",
65            |lang| lang == Language::Python,
66            Some(extract_llama_index_tools_from_source),
67        )
68    }
69}
70
71pub(crate) fn extract_llama_index_tools_from_source(
72    path: &Path,
73    content: &str,
74) -> Vec<ToolSurface> {
75    let mut tools = Vec::new();
76    let mut doc_map = std::collections::HashMap::new();
77
78    for cap in LLAMA_DEF_DOC_RE.captures_iter(content) {
79        let func_name = cap[1].to_string();
80        let doc = cap[2].trim().to_string();
81        doc_map.insert(func_name, doc);
82    }
83
84    for cap in LLAMA_FUNCTION_TOOL_RE.captures_iter(content) {
85        let fn_ident = cap
86            .get(1)
87            .map(|m| m.as_str().to_string())
88            .unwrap_or_default();
89        let explicit_name = cap.get(2).map(|m| m.as_str().to_string());
90        let tool_name = explicit_name.unwrap_or_else(|| fn_ident.clone());
91
92        if tool_name.is_empty() {
93            continue;
94        }
95
96        let explicit_desc = cap.get(3).map(|m| m.as_str().to_string());
97        let description = explicit_desc.or_else(|| doc_map.get(&fn_ident).cloned());
98
99        let match_start = cap.get(0).map(|m| m.start()).unwrap_or(0);
100        let line = content[..match_start].lines().count() + 1;
101
102        tools.push(super::pipeline::build_extracted_tool(
103            tool_name,
104            description,
105            path,
106            line,
107        ));
108    }
109
110    for cap in LLAMA_TOOL_METADATA_RE.captures_iter(content) {
111        let name = cap
112            .get(1)
113            .map(|m| m.as_str().to_string())
114            .unwrap_or_else(|| "query_tool".into());
115        let description = cap.get(2).map(|m| m.as_str().to_string());
116
117        let match_start = cap.get(0).map(|m| m.start()).unwrap_or(0);
118        let line = content[..match_start].lines().count() + 1;
119
120        tools.push(super::pipeline::build_extracted_tool(
121            name,
122            description,
123            path,
124            line,
125        ));
126    }
127
128    tools
129}
130
131#[cfg(test)]
132mod tests {
133    use super::*;
134    use crate::adapter::Adapter;
135
136    #[test]
137    fn detects_llama_index_requirements() {
138        let dir = tempfile::tempdir().unwrap();
139        let req = dir.path().join("requirements.txt");
140        std::fs::write(&req, "llama-index-core>=0.10.0\n").unwrap();
141
142        let adapter = LlamaIndexAdapter;
143        assert!(adapter.detect(dir.path()));
144    }
145
146    #[test]
147    fn detects_llama_index_import() {
148        let dir = tempfile::tempdir().unwrap();
149        let src = dir.path().join("rag.py");
150        std::fs::write(&src, "from llama_index.core.tools import FunctionTool\n").unwrap();
151
152        let adapter = LlamaIndexAdapter;
153        assert!(adapter.detect(dir.path()));
154    }
155
156    #[test]
157    fn extracts_llama_index_function_tool() {
158        let content = r#"
159from llama_index.core.tools import FunctionTool
160
161def fetch_document(doc_id: str) -> str:
162    """Fetch a document from the filesystem."""
163    with open(f"/data/{doc_id}") as f:
164        return f.read()
165
166doc_tool = FunctionTool.from_defaults(
167    fn=fetch_document,
168    name="doc_fetcher",
169    description="Fetch document content by ID"
170)
171"#;
172        let tools = extract_llama_index_tools_from_source(Path::new("rag.py"), content);
173        assert_eq!(tools.len(), 1);
174        assert_eq!(tools[0].name, "doc_fetcher");
175        assert_eq!(
176            tools[0].description.as_deref(),
177            Some("Fetch document content by ID")
178        );
179    }
180}
agentshield/adapter/llama_index.rs

agentshield/adapter/
llama_index.rs
