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

lc_rag/graph_rag/
extractor.rs
