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;
8use lc_core::tools::ToolDefinition;
9use lc_schema::Message;
10use serde::Deserialize;
11use serde_json::json;
12
13use crate::structured::{chat_structured, StructuredChatResult};
14
15/// An entity extracted from text by the LLM.
16#[derive(Debug, Clone, Deserialize)]
17pub struct ExtractedEntity {
18    /// Entity name.
19    pub name: String,
20    /// Entity type (e.g. Person, Organization, Technology).
21    #[serde(rename = "type")]
22    pub entity_type: String,
23    /// Optional description of the entity.
24    #[serde(default)]
25    pub description: String,
26}
27
28/// A relation extracted from text by the LLM.
29#[derive(Debug, Clone, Deserialize)]
30pub struct ExtractedRelation {
31    /// Source entity name.
32    pub source: String,
33    /// Target entity name.
34    pub target: String,
35    /// Relation type (e.g. works_at, uses, part_of).
36    #[serde(rename = "type")]
37    pub relation_type: String,
38    /// Optional description of the relation.
39    #[serde(default)]
40    pub description: String,
41}
42
43/// The full LLM extraction response.
44#[derive(Debug, Deserialize)]
45pub struct ExtractionResult {
46    /// Extracted entities.
47    #[serde(default)]
48    pub entities: Vec<ExtractedEntity>,
49    /// Extracted relations.
50    #[serde(default)]
51    pub relations: Vec<ExtractedRelation>,
52}
53
54const EXTRACTION_PROMPT: &str = r#"You are a knowledge graph extraction assistant. Given the following text, extract entities and their relations.
55
56Return a JSON object with exactly two keys:
57- "entities": an array of objects, each with keys "name", "type", "description"
58- "relations": an array of objects, each with keys "source", "target", "type", "description"
59
60Rules:
61- "source" and "target" in relations must match entity "name" values exactly.
62- Keep entity types simple: Person, Organization, Location, Technology, Concept, Event, etc.
63- Keep relation types simple: works_at, located_in, uses, created, part_of, related_to, etc.
64- Extract at most {max_entities} entities and {max_relations} relations.
65- Return ONLY the JSON object, no other text.
66
67Example:
68Text: "Alice works at Google as a software engineer. She uses Python and TensorFlow."
69Output:
70{
71  "entities": [
72    {"name": "Alice", "type": "Person", "description": "A software engineer at Google"},
73    {"name": "Google", "type": "Organization", "description": "A technology company"},
74    {"name": "Python", "type": "Technology", "description": "A programming language"},
75    {"name": "TensorFlow", "type": "Technology", "description": "A machine learning framework"}
76  ],
77  "relations": [
78    {"source": "Alice", "target": "Google", "type": "works_at", "description": "Alice is employed at Google"},
79    {"source": "Alice", "target": "Python", "type": "uses", "description": "Alice uses Python"},
80    {"source": "Alice", "target": "TensorFlow", "type": "uses", "description": "Alice uses TensorFlow"}
81  ]
82}
83
84Text:
85{text}"#;
86
87/// Entity/relation extraction tool definition (P2-1): forces the LLM to output structured
88/// JSON arguments.
89fn extraction_tool() -> ToolDefinition {
90    ToolDefinition::new(
91        "extract_entities_relations",
92        "从文本中提取知识图谱实体与关系,返回 entities 与 relations 两个数组",
93    )
94    .with_parameters(json!({
95        "type": "object",
96        "properties": {
97            "entities": {
98                "type": "array",
99                "items": {
100                    "type": "object",
101                    "properties": {
102                        "name": { "type": "string" },
103                        "type": { "type": "string" },
104                        "description": { "type": "string" }
105                    },
106                    "required": ["name", "type", "description"]
107                }
108            },
109            "relations": {
110                "type": "array",
111                "items": {
112                    "type": "object",
113                    "properties": {
114                        "source": { "type": "string" },
115                        "target": { "type": "string" },
116                        "type": { "type": "string" },
117                        "description": { "type": "string" }
118                    },
119                    "required": ["source", "target", "type", "description"]
120                }
121            }
122        },
123        "required": ["entities", "relations"]
124    }))
125}
126
127/// Parses one structured call result: prefers tool_calls arguments, then text JSON.
128fn parse_structured(result: &StructuredChatResult) -> Option<ExtractionResult> {
129    if let Some(args) = &result.tool_args {
130        if let Ok(parsed) = serde_json::from_value::<ExtractionResult>(args.clone()) {
131            return Some(parsed);
132        }
133    }
134    parse_extraction(&result.content).ok()
135}
136
137/// Extracts entities and relations from text using the LLM.
138pub async fn extract<M: BaseChatModel>(
139    llm: &M,
140    text: &str,
141    max_entities: usize,
142    max_relations: usize,
143) -> Result<ExtractionResult, super::GraphRAGError> {
144    let prompt = {
145        use lc_prompts::PromptTemplate;
146        let template = PromptTemplate::new(EXTRACTION_PROMPT);
147        let max_entities_str = max_entities.to_string();
148        let max_relations_str = max_relations.to_string();
149        let mut vars: std::collections::HashMap<&str, &str> = std::collections::HashMap::new();
150        vars.insert("max_entities", &max_entities_str);
151        vars.insert("max_relations", &max_relations_str);
152        vars.insert("text", text);
153        template
154            .format(&vars)
155            .unwrap_or_else(|_| EXTRACTION_PROMPT.to_string())
156    };
157
158    extract_with_retry(llm, text, prompt).await
159}
160
161/// P2-1/P2-2: prefers native tool_calls structured output; when text JSON parsing fails,
162/// retries 1-2 times with a "return JSON only" hint, and only then returns an error.
163async fn extract_with_retry<M: BaseChatModel>(
164    llm: &M,
165    original_text: &str,
166    prompt: String,
167) -> Result<ExtractionResult, super::GraphRAGError> {
168    const MAX_RETRIES: usize = 2;
169    let mut current_prompt = prompt;
170
171    for attempt in 0..=MAX_RETRIES {
172        let result = chat_structured(
173            llm,
174            Some(extraction_tool()),
175            vec![Message::human(&current_prompt)],
176        )
177        .await
178        .map_err(|e| super::GraphRAGError::LLMError(e.to_string()))?;
179
180        if let Some(parsed) = parse_structured(&result) {
181            return Ok(parsed);
182        }
183
184        if attempt < MAX_RETRIES {
185            current_prompt = format!(
186                "上次的输出不是合法 JSON,无法解析。请重新从下面文本提取实体与关系,\
187                 只返回一个 JSON 对象(键为 entities 与 relations),不要包含任何解释、\
188                 编号、引号或代码块。\n\n文本:\n{}\n\n上次输出(无效):\n{}\n\n只输出 JSON 对象:",
189                original_text, result.content
190            );
191        }
192    }
193
194    Err(super::GraphRAGError::ExtractionError(
195        "LLM repeatedly returned invalid JSON; entity/relation extraction failed".to_string(),
196    ))
197}
198
199/// Parses the LLM JSON response into an `ExtractionResult`.
200pub fn parse_extraction(raw: &str) -> Result<ExtractionResult, super::GraphRAGError> {
201    lc_core::json_parse::parse_llm_json::<ExtractionResult>(raw).map_err(|e| {
202        super::GraphRAGError::ExtractionError(format!("Failed to parse extraction JSON: {}", e))
203    })
204}
205
206/// Attempts to extract a JSON object from text that may contain markdown fences.
207#[cfg(test)]
208fn extract_json(text: &str) -> String {
209    let trimmed = text.trim();
210
211    // Case 1: wrapped in ```json ... ```
212    if let Some(rest) = trimmed.strip_prefix("```json") {
213        if let Some(end) = rest.find("```") {
214            return rest[..end].trim().to_string();
215        }
216    }
217
218    // Case 2: wrapped in ``` ... ```
219    if let Some(rest) = trimmed.strip_prefix("```") {
220        if let Some(end) = rest.find("```") {
221            return rest[..end].trim().to_string();
222        }
223    }
224
225    // Case 3: find first { ... last }
226    if let Some(start) = trimmed.find('{') {
227        if let Some(end) = trimmed.rfind('}') {
228            if end > start {
229                return trimmed[start..=end].to_string();
230            }
231        }
232    }
233
234    trimmed.to_string()
235}
236
237#[cfg(test)]
238mod tests {
239    use super::*;
240
241    #[test]
242    fn test_parse_extraction_valid() {
243        let raw = r#"{"entities":[{"name":"Alice","type":"Person","description":"A developer"}],"relations":[{"source":"Alice","target":"Rust","type":"uses","description":"Alice uses Rust"}]}"#;
244        let result = parse_extraction(raw).unwrap();
245        assert_eq!(result.entities.len(), 1);
246        assert_eq!(result.entities[0].name, "Alice");
247        assert_eq!(result.relations.len(), 1);
248        assert_eq!(result.relations[0].source, "Alice");
249    }
250
251    #[test]
252    fn test_parse_extraction_markdown_wrapped() {
253        let raw = r#"```json
254{"entities":[{"name":"Bob","type":"Person","description":"A manager"}],"relations":[]}
255```"#;
256        let result = parse_extraction(raw).unwrap();
257        assert_eq!(result.entities.len(), 1);
258        assert_eq!(result.entities[0].name, "Bob");
259    }
260
261    #[test]
262    fn test_parse_extraction_empty_arrays() {
263        let raw = r#"{"entities":[],"relations":[]}"#;
264        let result = parse_extraction(raw).unwrap();
265        assert!(result.entities.is_empty());
266        assert!(result.relations.is_empty());
267    }
268
269    #[test]
270    fn test_parse_extraction_invalid() {
271        let raw = "not json at all";
272        assert!(parse_extraction(raw).is_err());
273    }
274
275    #[test]
276    fn test_extract_json_plain() {
277        let input = r#"{"key": "value"}"#;
278        assert_eq!(extract_json(input), input);
279    }
280
281    #[test]
282    fn test_extract_json_code_fence() {
283        let input = "```json\n{\"key\": \"value\"}\n```";
284        assert_eq!(extract_json(input), "{\"key\": \"value\"}");
285    }
286
287    #[test]
288    fn test_extract_json_with_surrounding_text() {
289        let input = "Here is the result:\n{\"key\": \"value\"}\nDone.";
290        assert_eq!(extract_json(input), "{\"key\": \"value\"}");
291    }
292
293    /// Verify the extraction prompt contains a few-shot example.
294    #[test]
295    fn test_extraction_prompt_contains_few_shot_example() {
296        assert!(
297            EXTRACTION_PROMPT.contains("Example:"),
298            "extraction prompt should contain few-shot example"
299        );
300        assert!(
301            EXTRACTION_PROMPT.contains("Alice"),
302            "extraction prompt example should contain entity 'Alice'"
303        );
304        assert!(
305            EXTRACTION_PROMPT.contains("works_at"),
306            "extraction prompt example should contain relation type 'works_at'"
307        );
308    }
309
310    /// P2-1: the extraction tool definition carries a full JSON Schema (entities + relations).
311    #[test]
312    fn test_extraction_tool_schema() {
313        let tool = extraction_tool();
314        assert_eq!(tool.function.name, "extract_entities_relations");
315        let params = tool.function.parameters.expect("parameters should exist");
316        assert!(params["properties"]["entities"].is_object());
317        assert!(params["properties"]["relations"].is_object());
318    }
319
320    /// P2-1: tool_calls arguments parse into an ExtractionResult.
321    #[test]
322    fn test_parse_structured_tool_args() {
323        let result = StructuredChatResult {
324            content: "".to_string(),
325            tool_args: Some(json!({
326                "entities": [{"name": "Alice", "type": "Person", "description": "dev"}],
327                "relations": []
328            })),
329        };
330        let parsed = parse_structured(&result).expect("tool_args should parse successfully");
331        assert_eq!(parsed.entities.len(), 1);
332        assert_eq!(parsed.entities[0].name, "Alice");
333        assert!(parsed.relations.is_empty());
334    }
335
336    /// P2-1: falls back to text JSON parsing when there are no tool_calls.
337    #[test]
338    fn test_parse_structured_text_fallback() {
339        let result = StructuredChatResult {
340            content: r#"{"entities": [{"name": "Bob", "type": "Person", "description": "mgr"}], "relations": []}"#.to_string(),
341            tool_args: None,
342        };
343        let parsed = parse_structured(&result).expect("text JSON should parse successfully");
344        assert_eq!(parsed.entities[0].name, "Bob");
345    }
346
347    /// P2-1: when tool_args deserialization fails and the text is not JSON -> None
348    /// (triggering a retry).
349    #[test]
350    fn test_parse_structured_none() {
351        let result = StructuredChatResult {
352            content: "not json".to_string(),
353            tool_args: Some(json!("not an object")),
354        };
355        assert!(parse_structured(&result).is_none());
356    }
357}
lc_rag/graph_rag/extractor.rs

lc_rag/graph_rag/
extractor.rs
