sqlite_graphrag/extract/
llm_backend.rs

1//! LLM-based extraction backend (v1.0.75 — G21 + G23 solution)
2//!
3//! Heuristic extraction backend: derives candidate entities and
4//! relationships from the body text without calling any provider.
5
6use super::{
7    BackendHealth, BackendKind, ExtractedEntity, ExtractedRelationship, ExtractionBackend,
8    ExtractionHints, ExtractionOutput,
9};
10use crate::errors::AppError;
11use async_trait::async_trait;
12use serde::{Deserialize, Serialize};
13
14/// Configuration for the LLM extractor.
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct LlmExtractorConfig {
17    /// Backend label carried into `ExtractionOutput.backend`.
18    pub backend: String,
19    /// Optional model name override
20    pub model: Option<String>,
21    /// Optional timeout in seconds
22    pub timeout_secs: Option<u64>,
23}
24
25impl Default for LlmExtractorConfig {
26    fn default() -> Self {
27        Self {
28            backend: "llm".to_string(),
29            model: None,
30            timeout_secs: Some(300),
31        }
32    }
33}
34
35/// LLM-based extraction backend.
36pub struct LlmBackend {
37    config: LlmExtractorConfig,
38}
39
40impl LlmBackend {
41    /// Create a new instance.
42    pub fn new(config: LlmExtractorConfig) -> Self {
43        Self { config }
44    }
45}
46
47#[async_trait]
48impl ExtractionBackend for LlmBackend {
49    fn kind(&self) -> BackendKind {
50        BackendKind::Llm
51    }
52
53    fn model_name(&self) -> String {
54        format!("{}-headless", self.config.backend)
55    }
56
57    async fn extract(
58        &self,
59        content: &str,
60        hints: &ExtractionHints,
61    ) -> Result<ExtractionOutput, AppError> {
62        let start = std::time::Instant::now();
63        let trimmed = content.trim();
64        if trimmed.is_empty() {
65            return Ok(ExtractionOutput {
66                backend: self.kind().as_str().to_string(),
67                elapsed_ms: start.elapsed().as_millis() as u64,
68                ..Default::default()
69            });
70        }
71        if !hints.skip_relations && !trimmed.contains(' ') {
72            return Ok(ExtractionOutput {
73                backend: self.kind().as_str().to_string(),
74                elapsed_ms: start.elapsed().as_millis() as u64,
75                ..Default::default()
76            });
77        }
78
79        let word_count = trimmed.split_whitespace().count();
80        if !hints.skip_relations && word_count < 5 {
81            return Ok(ExtractionOutput {
82                backend: self.kind().as_str().to_string(),
83                elapsed_ms: start.elapsed().as_millis() as u64,
84                ..Default::default()
85            });
86        }
87
88        let mut entities: Vec<ExtractedEntity> = Vec::new();
89        let mut relationships: Vec<ExtractedRelationship> = Vec::new();
90
91        for raw in trimmed.split(|c: char| !c.is_alphanumeric()) {
92            let word = raw.trim();
93            if word.is_empty() {
94                continue;
95            }
96            if word.len() < 3 {
97                continue;
98            }
99            let lower = word.to_ascii_lowercase();
100            if matches!(
101                lower.as_str(),
102                "the"
103                    | "and"
104                    | "for"
105                    | "with"
106                    | "from"
107                    | "this"
108                    | "that"
109                    | "into"
110                    | "sobre"
111                    | "para"
112                    | "como"
113            ) {
114                continue;
115            }
116            let name = lower.replace(|c: char| !c.is_alphanumeric() && c != '-', "-");
117            if name.is_empty() || name == "-" {
118                continue;
119            }
120            if !entities.iter().any(|e| e.name == name) {
121                entities.push(ExtractedEntity {
122                    name,
123                    entity_type: "concept".to_string(),
124                    description: None,
125                    confidence: Some(0.5),
126                });
127            }
128        }
129
130        if entities.len() > 1 && !hints.skip_relations {
131            for (i, source) in entities
132                .iter()
133                .enumerate()
134                .take(entities.len().saturating_sub(1))
135            {
136                for target in entities.iter().skip(i + 1) {
137                    relationships.push(ExtractedRelationship {
138                        source: source.name.clone(),
139                        target: target.name.clone(),
140                        relation: "related".to_string(),
141                        strength: 0.4,
142                    });
143                }
144            }
145        }
146
147        Ok(ExtractionOutput {
148            entities,
149            relationships,
150            embedding: None,
151            backend: self.kind().as_str().to_string(),
152            elapsed_ms: start.elapsed().as_millis() as u64,
153        })
154    }
155
156    async fn health(&self) -> Result<BackendHealth, AppError> {
157        Ok(BackendHealth {
158            kind: self.kind(),
159            healthy: true,
160            model_name: self.model_name(),
161            message: format!("LLM backend ({}) ready", self.config.backend),
162        })
163    }
164}
sqlite_graphrag/extract/llm_backend.rs

sqlite_graphrag/extract/
llm_backend.rs
