sqlite_graphrag/
extraction.rs

1//! Entity and URL extraction pipeline (v1.0.76).
2//!
3//! v1.0.76: the default build is **LLM-only**. v1.0.79: the legacy GLiNER
4//! NER pipeline (`extraction_gliner.rs`, `ner-legacy` feature) was REMOVED
5//! entirely. The build extracts:
6//!
7//! - **URLs** via regex (always available, no model needed).
8//! - **Entities** via the `ExtractionBackend` trait (LLM headless).
9//!   The default backend is `LlmBackend` (claude / codex), which produces
10//!   structured entities and relationships via tool-use JSON.
11//!
12//! The `extract_graph_auto` function below is the entry point used by
13//! `remember`, `ingest`, and `enrich`. With the default feature set, it
14//! runs the regex URL pass; entity extraction happens through the LLM
15//! extraction backend selected at the command layer.
16
17use serde::{Deserialize, Serialize};
18
19/// One URL extracted from a body. Always produced by the regex path.
20#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
21pub struct ExtractedUrl {
22    /// URL value.
23    pub url: String,
24    /// Start offset (inclusive).
25    pub start: usize,
26    /// End offset (exclusive).
27    pub end: usize,
28}
29
30/// One named-entity mention. Produced via the LLM extraction backend
31/// (the legacy GLiNER NER build was removed in v1.0.79).
32#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
33pub struct ExtractedEntity {
34    /// Name of this item.
35    pub name: String,
36    /// Entity type label.
37    pub entity_type: String,
38    /// Start offset (inclusive).
39    pub start: usize,
40    /// End offset (exclusive).
41    pub end: usize,
42}
43
44/// Full extraction result: URLs (regex), entities (LLM), and the
45/// relationships between them. The LLM backend also returns typed
46/// relationships directly in `ExtractionOutput`; this struct is the
47/// regex-only baseline that `remember` and `ingest` consume.
48#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
49pub struct ExtractionResult {
50    /// Extracted entities.
51    pub entities: Vec<ExtractedEntity>,
52    /// Extracted URLs.
53    pub urls: Vec<ExtractedUrl>,
54    /// Wall-clock latency in milliseconds.
55    pub elapsed_ms: u64,
56}
57
58/// Trait abstraction for any extractor. Implemented by the LLM backend
59/// (the legacy GLiNER backend was removed in v1.0.79).
60pub trait Extractor: Send + Sync {
61    /// Return the name of this item.
62    fn name(&self) -> &'static str;
63    /// Extract structured data from the given body text.
64    fn extract(&self, body: &str) -> Result<ExtractionResult, crate::errors::AppError>;
65}
66
67/// Regex-only extractor: URLs and nothing else. Used as a fast
68/// pre-pass before the (slower) LLM extractor in `extract_graph_auto`.
69pub struct RegexExtractor;
70
71impl Extractor for RegexExtractor {
72    fn name(&self) -> &'static str {
73        "regex"
74    }
75    fn extract(&self, body: &str) -> Result<ExtractionResult, crate::errors::AppError> {
76        Ok(ExtractionResult {
77            entities: Vec::new(),
78            urls: extract_urls(body),
79            elapsed_ms: 0,
80        })
81    }
82}
83
84/// Extracts URLs from `body` using a substring scan. UTF-8 safe; offsets
85/// are byte indices into the input.
86pub fn extract_urls(body: &str) -> Vec<ExtractedUrl> {
87    let mut out = Vec::new();
88    let mut cursor = 0usize;
89    while cursor < body.len() {
90        let hay = &body[cursor..];
91        // Find the next URL boundary, considering both schemes.
92        let http_at = hay.find("http://");
93        let https_at = hay.find("https://");
94        let (rel_start, scheme_len) = match (http_at, https_at) {
95            (Some(a), Some(b)) => {
96                if a <= b {
97                    (a, 7)
98                } else {
99                    (b, 8)
100                }
101            }
102            (Some(a), None) => (a, 7),
103            (None, Some(b)) => (b, 8),
104            (None, None) => break,
105        };
106        let abs_start = cursor + rel_start;
107        let after_scheme = abs_start + scheme_len;
108        let mut end = after_scheme;
109        for (i, c) in body[after_scheme..].char_indices() {
110            if c.is_whitespace() || matches!(c, ')' | ']' | '}' | '"' | '\'' | '<') {
111                end = after_scheme + i;
112                break;
113            }
114            end = after_scheme + i + c.len_utf8();
115        }
116        out.push(ExtractedUrl {
117            url: body[abs_start..end].to_string(),
118            start: abs_start,
119            end,
120        });
121        cursor = end;
122    }
123    out
124}
125
126/// Top-level extraction entry point used by `remember`, `ingest`, and
127/// `enrich`. Runs the regex URL pass (always available); the legacy GLiNER
128/// delegation was removed in v1.0.79 together with the `ner-legacy` feature.
129pub fn extract_graph_auto(
130    body: &str,
131    _paths: &crate::paths::AppPaths,
132) -> Result<ExtractionResult, crate::errors::AppError> {
133    let start = std::time::Instant::now();
134    let urls = extract_urls(body);
135    Ok(ExtractionResult {
136        entities: Vec::new(),
137        urls,
138        elapsed_ms: start.elapsed().as_millis() as u64,
139    })
140}
141
142#[cfg(test)]
143mod tests {
144    use super::*;
145
146    #[test]
147    fn extract_urls_finds_http_and_https() {
148        let body = "see https://example.com/foo and http://bar.baz/qux end";
149        let urls = extract_urls(body);
150        assert_eq!(urls.len(), 2, "got {urls:?} for body {body:?}");
151        assert_eq!(urls[0].url, "https://example.com/foo");
152        assert_eq!(urls[1].url, "http://bar.baz/qux");
153    }
154
155    #[test]
156    fn extract_urls_handles_trailing_punctuation() {
157        let body = "see https://example.com/foo).";
158        let urls = extract_urls(body);
159        assert_eq!(urls.len(), 1);
160        assert_eq!(urls[0].url, "https://example.com/foo");
161    }
162
163    #[test]
164    fn extract_urls_empty_body() {
165        assert!(extract_urls("").is_empty());
166    }
167
168    #[test]
169    fn regex_extractor_returns_only_urls() {
170        let result = RegexExtractor.extract("see https://example.com").unwrap();
171        assert_eq!(result.entities.len(), 0);
172        assert_eq!(result.urls.len(), 1);
173    }
174}
sqlite_graphrag/extraction.rs

sqlite_graphrag/
extraction.rs
