lc_rag/graph_rag/
extractor.rs

1// src/retrieval/graph_rag/extractor.rs
2//! LLM-based entity and relation extraction from document text.
3//!
4//! Sends a structured prompt to the LLM, parses the JSON response into
5//! [`ExtractedEntity`] and [`ExtractedRelation`] lists.
6
7use lc_core::language_models::{BaseChatModel, LLMResult};
8use lc_schema::Message;
9use serde::Deserialize;
10
11/// An entity extracted from text by the LLM.
12#[derive(Debug, Clone, Deserialize)]
13pub struct ExtractedEntity {
14    pub name: String,
15    #[serde(rename = "type")]
16    pub entity_type: String,
17    #[serde(default)]
18    pub description: String,
19}
20
21/// A relation extracted from text by the LLM.
22#[derive(Debug, Clone, Deserialize)]
23pub struct ExtractedRelation {
24    pub source: String,
25    pub target: String,
26    #[serde(rename = "type")]
27    pub relation_type: String,
28    #[serde(default)]
29    pub description: String,
30}
31
32/// The full LLM extraction response.
33#[derive(Debug, Deserialize)]
34pub struct ExtractionResult {
35    #[serde(default)]
36    pub entities: Vec<ExtractedEntity>,
37    #[serde(default)]
38    pub relations: Vec<ExtractedRelation>,
39}
40
41const EXTRACTION_PROMPT: &str = r#"You are a knowledge graph extraction assistant. Given the following text, extract entities and their relations.
42
43Return a JSON object with exactly two keys:
44- "entities": an array of objects, each with keys "name", "type", "description"
45- "relations": an array of objects, each with keys "source", "target", "type", "description"
46
47Rules:
48- "source" and "target" in relations must match entity "name" values exactly.
49- Keep entity types simple: Person, Organization, Location, Technology, Concept, Event, etc.
50- Keep relation types simple: works_at, located_in, uses, created, part_of, related_to, etc.
51- Extract at most {max_entities} entities and {max_relations} relations.
52- Return ONLY the JSON object, no other text.
53
54Example:
55Text: "Alice works at Google as a software engineer. She uses Python and TensorFlow."
56Output:
57{
58  "entities": [
59    {"name": "Alice", "type": "Person", "description": "A software engineer at Google"},
60    {"name": "Google", "type": "Organization", "description": "A technology company"},
61    {"name": "Python", "type": "Technology", "description": "A programming language"},
62    {"name": "TensorFlow", "type": "Technology", "description": "A machine learning framework"}
63  ],
64  "relations": [
65    {"source": "Alice", "target": "Google", "type": "works_at", "description": "Alice is employed at Google"},
66    {"source": "Alice", "target": "Python", "type": "uses", "description": "Alice uses Python"},
67    {"source": "Alice", "target": "TensorFlow", "type": "uses", "description": "Alice uses TensorFlow"}
68  ]
69}
70
71Text:
72{text}"#;
73
74/// Extracts entities and relations from text using the LLM.
75pub async fn extract<M: BaseChatModel>(
76    llm: &M,
77    text: &str,
78    max_entities: usize,
79    max_relations: usize,
80) -> Result<ExtractionResult, super::GraphRAGError> {
81    let prompt = {
82        use lc_prompts::PromptTemplate;
83        let template = PromptTemplate::new(EXTRACTION_PROMPT);
84        let max_entities_str = max_entities.to_string();
85        let max_relations_str = max_relations.to_string();
86        let mut vars: std::collections::HashMap<&str, &str> = std::collections::HashMap::new();
87        vars.insert("max_entities", &max_entities_str);
88        vars.insert("max_relations", &max_relations_str);
89        vars.insert("text", text);
90        template
91            .format(&vars)
92            .unwrap_or_else(|_| EXTRACTION_PROMPT.to_string())
93    };
94
95    let messages = vec![Message::human(prompt)];
96
97    let response: LLMResult = llm
98        .chat(messages, None)
99        .await
100        .map_err(|e| super::GraphRAGError::LLMError(e.to_string()))?;
101
102    parse_extraction(&response.content)
103}
104
105/// Parses the LLM JSON response into an `ExtractionResult`.
106pub fn parse_extraction(raw: &str) -> Result<ExtractionResult, super::GraphRAGError> {
107    lc_core::json_parse::parse_llm_json::<ExtractionResult>(raw).map_err(|e| {
108        super::GraphRAGError::ExtractionError(format!("Failed to parse extraction JSON: {}", e))
109    })
110}
111
112/// Attempts to extract a JSON object from text that may contain markdown fences.
113#[cfg(test)]
114fn extract_json(text: &str) -> String {
115    let trimmed = text.trim();
116
117    // Case 1: wrapped in ```json ... ```
118    if let Some(rest) = trimmed.strip_prefix("```json") {
119        if let Some(end) = rest.find("```") {
120            return rest[..end].trim().to_string();
121        }
122    }
123
124    // Case 2: wrapped in ``` ... ```
125    if let Some(rest) = trimmed.strip_prefix("```") {
126        if let Some(end) = rest.find("```") {
127            return rest[..end].trim().to_string();
128        }
129    }
130
131    // Case 3: find first { ... last }
132    if let Some(start) = trimmed.find('{') {
133        if let Some(end) = trimmed.rfind('}') {
134            if end > start {
135                return trimmed[start..=end].to_string();
136            }
137        }
138    }
139
140    trimmed.to_string()
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[test]
148    fn test_parse_extraction_valid() {
149        let raw = r#"{"entities":[{"name":"Alice","type":"Person","description":"A developer"}],"relations":[{"source":"Alice","target":"Rust","type":"uses","description":"Alice uses Rust"}]}"#;
150        let result = parse_extraction(raw).unwrap();
151        assert_eq!(result.entities.len(), 1);
152        assert_eq!(result.entities[0].name, "Alice");
153        assert_eq!(result.relations.len(), 1);
154        assert_eq!(result.relations[0].source, "Alice");
155    }
156
157    #[test]
158    fn test_parse_extraction_markdown_wrapped() {
159        let raw = r#"```json
160{"entities":[{"name":"Bob","type":"Person","description":"A manager"}],"relations":[]}
161```"#;
162        let result = parse_extraction(raw).unwrap();
163        assert_eq!(result.entities.len(), 1);
164        assert_eq!(result.entities[0].name, "Bob");
165    }
166
167    #[test]
168    fn test_parse_extraction_empty_arrays() {
169        let raw = r#"{"entities":[],"relations":[]}"#;
170        let result = parse_extraction(raw).unwrap();
171        assert!(result.entities.is_empty());
172        assert!(result.relations.is_empty());
173    }
174
175    #[test]
176    fn test_parse_extraction_invalid() {
177        let raw = "not json at all";
178        assert!(parse_extraction(raw).is_err());
179    }
180
181    #[test]
182    fn test_extract_json_plain() {
183        let input = r#"{"key": "value"}"#;
184        assert_eq!(extract_json(input), input);
185    }
186
187    #[test]
188    fn test_extract_json_code_fence() {
189        let input = "```json\n{\"key\": \"value\"}\n```";
190        assert_eq!(extract_json(input), "{\"key\": \"value\"}");
191    }
192
193    #[test]
194    fn test_extract_json_with_surrounding_text() {
195        let input = "Here is the result:\n{\"key\": \"value\"}\nDone.";
196        assert_eq!(extract_json(input), "{\"key\": \"value\"}");
197    }
198
199    /// Verify the extraction prompt contains a few-shot example.
200    #[test]
201    fn test_extraction_prompt_contains_few_shot_example() {
202        assert!(
203            EXTRACTION_PROMPT.contains("Example:"),
204            "extraction prompt should contain few-shot example"
205        );
206        assert!(
207            EXTRACTION_PROMPT.contains("Alice"),
208            "extraction prompt example should contain entity 'Alice'"
209        );
210        assert!(
211            EXTRACTION_PROMPT.contains("works_at"),
212            "extraction prompt example should contain relation type 'works_at'"
213        );
214    }
215}
lc_rag/graph_rag/extractor.rs

lc_rag/graph_rag/
extractor.rs
