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

sqlite_graphrag/
extraction.rs
