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::capability::{project_declared_description, project_declared_permissions};
9use crate::ir::taint_builder::build_data_surface;
10use crate::ir::*;
11
12static LLAMA_FUNCTION_TOOL_RE: Lazy<Regex> = Lazy::new(|| {
13    Regex::new(r#"(?s)FunctionTool\.from_defaults\s*\(\s*(?:fn\s*=\s*)?(\w+)(?:[^\)]*?name\s*=\s*["'`]([^"'`]+)["'`])?(?:[^\)]*?description\s*=\s*["'`]([^"'`]+)["'`])?"#)
14        .expect("static regex pattern is valid")
15});
16
17static LLAMA_TOOL_METADATA_RE: Lazy<Regex> = Lazy::new(|| {
18    Regex::new(r#"(?s)ToolMetadata\s*\(\s*(?:name\s*=\s*["'`]([^"'`]+)["'`])?[^\)]*?description\s*=\s*["'`]([^"'`]+)["'`]"#)
19        .expect("static regex pattern is valid")
20});
21
22static LLAMA_DEF_DOC_RE: Lazy<Regex> = Lazy::new(|| {
23    Regex::new(r#"(?m)^\s*(?:async\s+)?def\s+(\w+)\s*\([^)]*\)\s*(?:->[^:]+)?:\s*\n\s*["'`]([\s\S]*?)["'`]"#)
24        .expect("static regex pattern is valid")
25});
26
27/// LlamaIndex framework adapter (`llama-index` and `llama-index-core`).
28///
29/// Detects LlamaIndex projects by looking for:
30/// - `pyproject.toml` / `requirements.txt` with `llama-index` or `llama_index`
31/// - Python files importing `from llama_index` or `import llama_index`
32pub struct LlamaIndexAdapter;
33
34impl super::Adapter for LlamaIndexAdapter {
35    fn framework(&self) -> Framework {
36        Framework::LlamaIndex
37    }
38
39    fn detect(&self, root: &Path) -> bool {
40        let pyproject = root.join("pyproject.toml");
41        if pyproject.exists() {
42            if let Some(content) = super::read_file_capped(&pyproject) {
43                if content.contains("llama-index") || content.contains("llama_index") {
44                    return true;
45                }
46            }
47        }
48
49        let requirements = root.join("requirements.txt");
50        if requirements.exists() {
51            if let Some(content) = super::read_file_capped(&requirements) {
52                if content.lines().any(|l| {
53                    let trimmed = l.trim();
54                    trimmed.starts_with("llama-index")
55                        || trimmed.starts_with("llama_index")
56                        || trimmed.starts_with("llama-index-core")
57                }) {
58                    return true;
59                }
60            }
61        }
62
63        if super::mcp::has_recursive_python_import(
64            root,
65            &[
66                "from llama_index",
67                "import llama_index",
68                "from llama_index.core",
69            ],
70        ) {
71            return true;
72        }
73
74        false
75    }
76
77    fn load(&self, root: &Path, ignore_tests: bool) -> Result<Vec<ScanTarget>> {
78        let filter = ScanPathFilter::for_ignore_tests(ignore_tests);
79        self.load_with_filter(root, &filter)
80    }
81
82    fn load_with_filter(&self, root: &Path, filter: &ScanPathFilter) -> Result<Vec<ScanTarget>> {
83        let name = root
84            .file_name()
85            .map(|n| n.to_string_lossy().to_string())
86            .unwrap_or_else(|| "llama-index-project".into());
87
88        let mut source_files = Vec::new();
89        super::mcp::collect_source_files_with_filter(root, filter, &mut source_files)?;
90
91        source_files.retain(|sf| sf.language == Language::Python);
92
93        let mut tools = Vec::new();
94        for sf in &source_files {
95            tools.extend(extract_llama_index_tools_from_source(&sf.path, &sf.content));
96        }
97
98        for tool in &mut tools {
99            project_declared_permissions(tool);
100            project_declared_description(tool);
101        }
102
103        let execution = super::pipeline::build_execution_surface(&source_files);
104        let data = build_data_surface(&tools, &execution);
105        let dependencies = super::mcp::parse_dependencies(root, filter);
106        let provenance = super::mcp::parse_provenance(root, filter);
107
108        Ok(vec![ScanTarget {
109            name,
110            framework: Framework::LlamaIndex,
111            root_path: root.to_path_buf(),
112            tools,
113            execution,
114            data,
115            dependencies,
116            provenance,
117            source_files,
118        }])
119    }
120}
121
122pub(crate) fn extract_llama_index_tools_from_source(path: &Path, content: &str) -> Vec<ToolSurface> {
123    let mut tools = Vec::new();
124    let mut doc_map = std::collections::HashMap::new();
125
126    for cap in LLAMA_DEF_DOC_RE.captures_iter(content) {
127        let func_name = cap[1].to_string();
128        let doc = cap[2].trim().to_string();
129        doc_map.insert(func_name, doc);
130    }
131
132    for cap in LLAMA_FUNCTION_TOOL_RE.captures_iter(content) {
133        let fn_ident = cap.get(1).map(|m| m.as_str().to_string()).unwrap_or_default();
134        let explicit_name = cap.get(2).map(|m| m.as_str().to_string());
135        let tool_name = explicit_name.unwrap_or_else(|| fn_ident.clone());
136
137        if tool_name.is_empty() {
138            continue;
139        }
140
141        let explicit_desc = cap.get(3).map(|m| m.as_str().to_string());
142        let description = explicit_desc.or_else(|| doc_map.get(&fn_ident).cloned());
143
144        let match_start = cap.get(0).map(|m| m.start()).unwrap_or(0);
145        let line = content[..match_start].lines().count() + 1;
146
147        tools.push(ToolSurface {
148            name: tool_name,
149            description,
150            input_schema: None,
151            output_schema: None,
152            declared_permissions: Vec::new(),
153            declared_capabilities: std::collections::BTreeSet::new(),
154            capability_declarations: Vec::new(),
155            observed_capabilities: std::collections::BTreeSet::new(),
156            capability_evidence: Vec::new(),
157            capability_observation_complete: false,
158            defined_at: Some(SourceLocation {
159                file: path.to_path_buf(),
160                line,
161                column: 0,
162                end_line: Some(line),
163                end_column: None,
164            }),
165        });
166    }
167
168    for cap in LLAMA_TOOL_METADATA_RE.captures_iter(content) {
169        let name = cap.get(1).map(|m| m.as_str().to_string()).unwrap_or_else(|| "query_tool".into());
170        let description = cap.get(2).map(|m| m.as_str().to_string());
171
172        let match_start = cap.get(0).map(|m| m.start()).unwrap_or(0);
173        let line = content[..match_start].lines().count() + 1;
174
175        tools.push(ToolSurface {
176            name,
177            description,
178            input_schema: None,
179            output_schema: None,
180            declared_permissions: Vec::new(),
181            declared_capabilities: std::collections::BTreeSet::new(),
182            capability_declarations: Vec::new(),
183            observed_capabilities: std::collections::BTreeSet::new(),
184            capability_evidence: Vec::new(),
185            capability_observation_complete: false,
186            defined_at: Some(SourceLocation {
187                file: path.to_path_buf(),
188                line,
189                column: 0,
190                end_line: Some(line),
191                end_column: None,
192            }),
193        });
194    }
195
196    tools
197}
198
199#[cfg(test)]
200mod tests {
201    use super::*;
202    use crate::adapter::Adapter;
203
204    #[test]
205    fn detects_llama_index_requirements() {
206        let dir = tempfile::tempdir().unwrap();
207        let req = dir.path().join("requirements.txt");
208        std::fs::write(&req, "llama-index-core>=0.10.0\n").unwrap();
209
210        let adapter = LlamaIndexAdapter;
211        assert!(adapter.detect(dir.path()));
212    }
213
214    #[test]
215    fn detects_llama_index_import() {
216        let dir = tempfile::tempdir().unwrap();
217        let src = dir.path().join("rag.py");
218        std::fs::write(&src, "from llama_index.core.tools import FunctionTool\n").unwrap();
219
220        let adapter = LlamaIndexAdapter;
221        assert!(adapter.detect(dir.path()));
222    }
223
224    #[test]
225    fn extracts_llama_index_function_tool() {
226        let content = r#"
227from llama_index.core.tools import FunctionTool
228
229def fetch_document(doc_id: str) -> str:
230    """Fetch a document from the filesystem."""
231    with open(f"/data/{doc_id}") as f:
232        return f.read()
233
234doc_tool = FunctionTool.from_defaults(
235    fn=fetch_document,
236    name="doc_fetcher",
237    description="Fetch document content by ID"
238)
239"#;
240        let tools = extract_llama_index_tools_from_source(Path::new("rag.py"), content);
241        assert_eq!(tools.len(), 1);
242        assert_eq!(tools[0].name, "doc_fetcher");
243        assert_eq!(
244            tools[0].description.as_deref(),
245            Some("Fetch document content by ID")
246        );
247    }
248}
agentshield/adapter/llama_index.rs

agentshield/adapter/
llama_index.rs
