1use crate::config::LlmConfig;
2use crate::error::{CodeSynapseError, Result};
3use crate::extract::make_id;
4use crate::types::{Edge, ExtractionFragment, Node};
5use serde::Deserialize;
6use std::collections::HashMap;
7use std::path::Path;
8
9pub trait LlmExtractor: Send + Sync {
10 fn extract(&self, source: &[u8], path: &Path) -> Result<ExtractionFragment>;
11}
12
13const SYSTEM_PROMPT: &str = "You are a knowledge graph extractor. Given text, extract entities as nodes and relationships as edges. Respond with ONLY a JSON object — no prose, no markdown fences. Schema: {\"nodes\":[{\"id\":\"slug\",\"label\":\"Human Label\"}],\"edges\":[{\"source\":\"id1\",\"target\":\"id2\",\"relation\":\"relationship_type\"}]}";
14
15fn extraction_user_prompt(text: &str) -> String {
16 format!(
17 "Extract knowledge graph from:\n\n{}",
18 &text[..text.len().min(8000)]
19 )
20}
21
22#[derive(Deserialize)]
23struct LlmNode {
24 id: String,
25 label: String,
26}
27
28#[derive(Deserialize)]
29struct LlmEdge {
30 source: String,
31 target: String,
32 relation: String,
33}
34
35#[derive(Deserialize)]
36struct LlmResponse {
37 nodes: Vec<LlmNode>,
38 edges: Vec<LlmEdge>,
39}
40
41pub fn parse_llm_response(text: &str, path: &Path) -> Result<ExtractionFragment> {
42 let json_str = strip_fences(text);
43 let resp: LlmResponse = serde_json::from_str(json_str)
44 .map_err(|e| CodeSynapseError::Parse(format!("LLM response parse error: {e}")))?;
45
46 let source_file = path.to_string_lossy().to_string();
47 let nodes = resp
48 .nodes
49 .into_iter()
50 .map(|n| Node {
51 id: n.id,
52 label: n.label,
53 file_type: "llm".to_string(),
54 source_file: source_file.clone(),
55 source_location: None,
56 community: None,
57 rationale: None,
58 docstring: None,
59 metadata: HashMap::new(),
60 })
61 .collect();
62
63 let edges = resp
64 .edges
65 .into_iter()
66 .map(|e| Edge {
67 source: e.source,
68 target: e.target,
69 relation: e.relation,
70 confidence: "high".to_string(),
71 source_file: Some(source_file.clone()),
72 weight: 1.0,
73 context: None,
74 })
75 .collect();
76
77 Ok(ExtractionFragment { nodes, edges })
78}
79
80fn strip_fences(text: &str) -> &str {
81 let trimmed = text.trim();
82 if let Some(inner) = trimmed
83 .strip_prefix("```json")
84 .or_else(|| trimmed.strip_prefix("```"))
85 {
86 if let Some(end) = inner.rfind("```") {
87 return inner[..end].trim();
88 }
89 }
90 trimmed
91}
92
93fn fallback_fragment(path: &Path) -> ExtractionFragment {
94 let file_id = make_id(&[path
95 .file_stem()
96 .unwrap_or_default()
97 .to_string_lossy()
98 .as_ref()]);
99 ExtractionFragment {
100 nodes: vec![Node {
101 id: file_id,
102 label: path
103 .file_name()
104 .unwrap_or_default()
105 .to_string_lossy()
106 .to_string(),
107 file_type: "llm".to_string(),
108 source_file: path.to_string_lossy().to_string(),
109 source_location: None,
110 community: None,
111 rationale: None,
112 docstring: None,
113 metadata: HashMap::new(),
114 }],
115 edges: vec![],
116 }
117}
118
119#[allow(clippy::result_large_err)]
120fn send_json(
121 req: ureq::Request,
122 body: serde_json::Value,
123) -> std::result::Result<ureq::Response, ureq::Error> {
124 let s = body.to_string();
125 req.send_string(&s)
126}
127
128pub struct AnthropicLlmExtractor {
129 pub api_key: String,
130 pub model: String,
131}
132
133impl LlmExtractor for AnthropicLlmExtractor {
134 fn extract(&self, source: &[u8], path: &Path) -> Result<ExtractionFragment> {
135 let text = String::from_utf8_lossy(source);
136 let body = serde_json::json!({
137 "model": self.model,
138 "max_tokens": 1024,
139 "system": SYSTEM_PROMPT,
140 "messages": [{"role": "user", "content": extraction_user_prompt(&text)}]
141 });
142 let req = ureq::post("https://api.anthropic.com/v1/messages")
143 .set("x-api-key", &self.api_key)
144 .set("anthropic-version", "2023-06-01")
145 .set("content-type", "application/json");
146 let response = send_json(req, body)
147 .map_err(|e| CodeSynapseError::Other(format!("Anthropic API error: {e}")))?;
148 let body = response
149 .into_string()
150 .map_err(|e| CodeSynapseError::Other(format!("Anthropic response read error: {e}")))?;
151 let json: serde_json::Value = serde_json::from_str(&body)
152 .map_err(|e| CodeSynapseError::Other(format!("Anthropic response parse error: {e}")))?;
153 let content = json["content"][0]["text"]
154 .as_str()
155 .unwrap_or("")
156 .to_string();
157 parse_llm_response(&content, path).or_else(|_| Ok(fallback_fragment(path)))
158 }
159}
160
161pub struct OpenAiLlmExtractor {
162 pub api_key: String,
163 pub model: String,
164 pub base_url: String,
165}
166
167impl LlmExtractor for OpenAiLlmExtractor {
168 fn extract(&self, source: &[u8], path: &Path) -> Result<ExtractionFragment> {
169 let text = String::from_utf8_lossy(source);
170 let body = serde_json::json!({
171 "model": self.model,
172 "max_tokens": 1024,
173 "messages": [
174 {"role": "system", "content": SYSTEM_PROMPT},
175 {"role": "user", "content": extraction_user_prompt(&text)}
176 ]
177 });
178 let url = format!("{}/chat/completions", self.base_url.trim_end_matches('/'));
179 let mut req = ureq::post(&url).set("content-type", "application/json");
180 if !self.api_key.is_empty() {
181 req = req.set("authorization", &format!("Bearer {}", self.api_key));
182 }
183 let response = send_json(req, body)
184 .map_err(|e| CodeSynapseError::Other(format!("OpenAI API error: {e}")))?;
185 let raw = response
186 .into_string()
187 .map_err(|e| CodeSynapseError::Other(format!("OpenAI response read error: {e}")))?;
188 let json: serde_json::Value = serde_json::from_str(&raw)
189 .map_err(|e| CodeSynapseError::Other(format!("OpenAI response parse error: {e}")))?;
190 let content = json["choices"][0]["message"]["content"]
191 .as_str()
192 .unwrap_or("")
193 .to_string();
194 parse_llm_response(&content, path).or_else(|_| Ok(fallback_fragment(path)))
195 }
196}
197
198pub struct OllamaLlmExtractor {
199 pub base_url: String,
200 pub model: String,
201}
202
203impl LlmExtractor for OllamaLlmExtractor {
204 fn extract(&self, source: &[u8], path: &Path) -> Result<ExtractionFragment> {
205 let text = String::from_utf8_lossy(source);
206 let body = serde_json::json!({
207 "model": self.model,
208 "stream": false,
209 "messages": [
210 {"role": "system", "content": SYSTEM_PROMPT},
211 {"role": "user", "content": extraction_user_prompt(&text)}
212 ]
213 });
214 let url = format!("{}/api/chat", self.base_url.trim_end_matches('/'));
215 let req = ureq::post(&url).set("content-type", "application/json");
216 let response = send_json(req, body)
217 .map_err(|e| CodeSynapseError::Other(format!("Ollama API error: {e}")))?;
218 let raw = response
219 .into_string()
220 .map_err(|e| CodeSynapseError::Other(format!("Ollama response read error: {e}")))?;
221 let json: serde_json::Value = serde_json::from_str(&raw)
222 .map_err(|e| CodeSynapseError::Other(format!("Ollama response parse error: {e}")))?;
223 let content = json["message"]["content"]
224 .as_str()
225 .unwrap_or("")
226 .to_string();
227 parse_llm_response(&content, path).or_else(|_| Ok(fallback_fragment(path)))
228 }
229}
230
231pub fn build_extractor(config: &LlmConfig) -> Result<Box<dyn LlmExtractor>> {
232 match config.provider.as_deref().unwrap_or("anthropic") {
233 "anthropic" => {
234 let api_key = config
235 .api_key
236 .clone()
237 .or_else(|| std::env::var("ANTHROPIC_API_KEY").ok())
238 .unwrap_or_default();
239 Ok(Box::new(AnthropicLlmExtractor {
240 api_key,
241 model: config
242 .model
243 .clone()
244 .unwrap_or_else(|| "claude-haiku-4-5-20251001".to_string()),
245 }))
246 }
247 "openai" => {
248 let api_key = config
249 .api_key
250 .clone()
251 .or_else(|| std::env::var("OPENAI_API_KEY").ok())
252 .unwrap_or_default();
253 Ok(Box::new(OpenAiLlmExtractor {
254 api_key,
255 model: config
256 .model
257 .clone()
258 .unwrap_or_else(|| "gpt-4o-mini".to_string()),
259 base_url: config
260 .base_url
261 .clone()
262 .unwrap_or_else(|| "https://api.openai.com/v1".to_string()),
263 }))
264 }
265 "ollama" => Ok(Box::new(OllamaLlmExtractor {
266 model: config.model.clone().unwrap_or_else(|| "llama3".to_string()),
267 base_url: config
268 .base_url
269 .clone()
270 .unwrap_or_else(|| "http://localhost:11434".to_string()),
271 })),
272 "openai-compat" => Ok(Box::new(OpenAiLlmExtractor {
273 api_key: config.api_key.clone().unwrap_or_default(),
274 model: config.model.clone().unwrap_or_default(),
275 base_url: config.base_url.clone().unwrap_or_default(),
276 })),
277 other => Err(CodeSynapseError::Other(format!(
278 "Unknown LLM provider: {other}"
279 ))),
280 }
281}
282
283#[cfg(test)]
284mod tests {
285 use super::*;
286 use std::path::PathBuf;
287
288 fn test_path() -> PathBuf {
289 PathBuf::from("test/doc.md")
290 }
291
292 #[test]
293 fn test_strip_fences_plain_json() {
294 let input = r#"{"nodes":[],"edges":[]}"#;
295 assert_eq!(strip_fences(input), input);
296 }
297
298 #[test]
299 fn test_strip_fences_json_block() {
300 let input = "```json\n{\"nodes\":[],\"edges\":[]}\n```";
301 assert_eq!(strip_fences(input), "{\"nodes\":[],\"edges\":[]}");
302 }
303
304 #[test]
305 fn test_strip_fences_plain_block() {
306 let input = "```\n{\"nodes\":[],\"edges\":[]}\n```";
307 assert_eq!(strip_fences(input), "{\"nodes\":[],\"edges\":[]}");
308 }
309
310 #[test]
311 fn test_parse_llm_response_valid() {
312 let json = r#"{"nodes":[{"id":"foo","label":"Foo"}],"edges":[{"source":"foo","target":"bar","relation":"uses"}]}"#;
313 let fragment = parse_llm_response(json, &test_path()).unwrap();
314 assert_eq!(fragment.nodes.len(), 1);
315 assert_eq!(fragment.nodes[0].id, "foo");
316 assert_eq!(fragment.nodes[0].label, "Foo");
317 assert_eq!(fragment.nodes[0].file_type, "llm");
318 assert_eq!(fragment.edges.len(), 1);
319 assert_eq!(fragment.edges[0].relation, "uses");
320 }
321
322 #[test]
323 fn test_parse_llm_response_fenced() {
324 let json = "```json\n{\"nodes\":[{\"id\":\"a\",\"label\":\"A\"}],\"edges\":[]}\n```";
325 let fragment = parse_llm_response(json, &test_path()).unwrap();
326 assert_eq!(fragment.nodes.len(), 1);
327 assert_eq!(fragment.nodes[0].id, "a");
328 }
329
330 #[test]
331 fn test_parse_llm_response_sets_source_file() {
332 let json = r#"{"nodes":[{"id":"n","label":"N"}],"edges":[]}"#;
333 let path = PathBuf::from("/some/path/doc.txt");
334 let fragment = parse_llm_response(json, &path).unwrap();
335 assert_eq!(fragment.nodes[0].source_file, "/some/path/doc.txt");
336 }
337
338 #[test]
339 fn test_parse_llm_response_empty() {
340 let json = r#"{"nodes":[],"edges":[]}"#;
341 let fragment = parse_llm_response(json, &test_path()).unwrap();
342 assert!(fragment.nodes.is_empty());
343 assert!(fragment.edges.is_empty());
344 }
345
346 #[test]
347 fn test_parse_llm_response_invalid_json() {
348 let result = parse_llm_response("not json at all", &test_path());
349 assert!(result.is_err());
350 let msg = result.unwrap_err().to_string();
351 assert!(msg.contains("parse error") || msg.contains("Parse error"));
352 }
353
354 #[test]
355 fn test_parse_llm_response_edge_has_source_file() {
356 let json = r#"{"nodes":[],"edges":[{"source":"a","target":"b","relation":"calls"}]}"#;
357 let path = PathBuf::from("/tmp/note.md");
358 let fragment = parse_llm_response(json, &path).unwrap();
359 assert_eq!(
360 fragment.edges[0].source_file,
361 Some("/tmp/note.md".to_string())
362 );
363 assert_eq!(fragment.edges[0].confidence, "high");
364 }
365
366 #[test]
367 fn test_build_extractor_anthropic() {
368 let config = LlmConfig {
369 provider: Some("anthropic".to_string()),
370 model: Some("claude-haiku-4-5-20251001".to_string()),
371 api_key: Some("test-key".to_string()),
372 base_url: None,
373 };
374 assert!(build_extractor(&config).is_ok());
375 }
376
377 #[test]
378 fn test_build_extractor_anthropic_default() {
379 let config = LlmConfig {
380 provider: None,
381 model: None,
382 api_key: None,
383 base_url: None,
384 };
385 assert!(build_extractor(&config).is_ok());
386 }
387
388 #[test]
389 fn test_build_extractor_openai() {
390 let config = LlmConfig {
391 provider: Some("openai".to_string()),
392 model: Some("gpt-4o-mini".to_string()),
393 api_key: Some("sk-test".to_string()),
394 base_url: None,
395 };
396 assert!(build_extractor(&config).is_ok());
397 }
398
399 #[test]
400 fn test_build_extractor_openai_compat() {
401 let config = LlmConfig {
402 provider: Some("openai-compat".to_string()),
403 model: Some("custom-model".to_string()),
404 api_key: Some("key".to_string()),
405 base_url: Some("http://localhost:8080/v1".to_string()),
406 };
407 assert!(build_extractor(&config).is_ok());
408 }
409
410 #[test]
411 fn test_build_extractor_ollama() {
412 let config = LlmConfig {
413 provider: Some("ollama".to_string()),
414 model: Some("llama3".to_string()),
415 api_key: None,
416 base_url: Some("http://localhost:11434".to_string()),
417 };
418 assert!(build_extractor(&config).is_ok());
419 }
420
421 #[test]
422 fn test_build_extractor_unknown_provider() {
423 let config = LlmConfig {
424 provider: Some("fakeprovider".to_string()),
425 model: None,
426 api_key: None,
427 base_url: None,
428 };
429 let result = build_extractor(&config);
430 assert!(result.is_err());
431 let err = result.err().unwrap();
432 let msg = err.to_string();
433 assert!(msg.contains("Unknown LLM provider: fakeprovider"));
434 }
435
436 #[test]
437 fn test_fallback_fragment_has_one_node() {
438 let path = PathBuf::from("/tmp/readme.md");
439 let fragment = fallback_fragment(&path);
440 assert_eq!(fragment.nodes.len(), 1);
441 assert_eq!(fragment.edges.len(), 0);
442 assert_eq!(fragment.nodes[0].file_type, "llm");
443 }
444}