Skip to main content

link_r/extract/
mod.rs

1// Author: support@lituus.io
2// Copyright: lituus-io, all rights reserved
3//
4//! Extraction: distill a fetched document into the small descriptor the index
5//! stores.
6//!
7//! An [`Extractor`] turns raw bytes into a [`Descriptor`]: a title, keywords/topics,
8//! tags, a display snippet, the distilled `embed_text` that gets embedded, the
9//! BM25 terms, and (for HTML) the outbound links the crawler follows. The body
10//! itself never enters the index. [`AutoExtractor`] routes by [`ResourceKind`].
11
12#[cfg(feature = "html")]
13pub mod html;
14pub mod markdown;
15
16use crate::error::Result;
17use crate::resource::{Resource, ResourceKind};
18use crate::text;
19use crate::url_key::UrlKeyExt as _;
20use compact_str::CompactString;
21use smallvec::SmallVec;
22use std::collections::HashMap;
23
24/// Maximum number of BM25 terms retained per document (caps huge pages).
25const MAX_BM25_TERMS: usize = 4096;
26/// Maximum keywords surfaced per document.
27const MAX_KEYWORDS: usize = 8;
28/// Minimum token length to be considered a keyword.
29const MIN_KEYWORD_LEN: usize = 3;
30/// Snippet length in characters.
31const SNIPPET_CHARS: usize = 180;
32/// Maximum persisted outbound edges (knowledge-graph paths) per document.
33const MAX_EDGES_PER_DOC: usize = 64;
34/// Body characters folded into the embed text. Headings and keywords already
35/// front-load the high-signal terms; a wider body window improves recall for a
36/// knowledge base without materially growing the single per-document vector.
37const EMBED_BODY_CHARS: usize = 1000;
38/// Maximum headings retained per document.
39const MAX_HEADINGS: usize = 8;
40
41/// The distilled, indexable form of a document.
42#[derive(Clone, Debug, Default)]
43pub struct Descriptor {
44    /// Display title, if found.
45    pub title: Option<CompactString>,
46    /// Top keywords/topics (most frequent discriminative terms).
47    pub keywords: SmallVec<[CompactString; MAX_KEYWORDS]>,
48    /// Section headings (H1–H3 / ATX), high-signal for retrieval.
49    pub headings: SmallVec<[CompactString; MAX_HEADINGS]>,
50    /// Heading depths parallel to `headings` (1 = H1 …); empty when the
51    /// extractor cannot tell (code/plain text), in which case consumers fall
52    /// back to position-derived depths.
53    pub heading_levels: SmallVec<[u8; MAX_HEADINGS]>,
54    /// Structured tags (currently URL path components) for filtering.
55    pub tags: SmallVec<[CompactString; 8]>,
56    /// Detected language tag, if any.
57    pub lang: Option<CompactString>,
58    /// A short display snippet.
59    pub snippet: CompactString,
60    /// The distilled text handed to the embedder.
61    pub embed_text: String,
62    /// Normalized BM25 terms (also defines document length).
63    pub bm25_terms: Vec<CompactString>,
64    /// Outbound link hrefs (raw, unresolved) discovered in the document.
65    pub links: Vec<CompactString>,
66}
67
68impl Descriptor {
69    /// Whether this descriptor carries no indexable content (no body/title terms).
70    /// URL-path tags alone do not count as content.
71    #[must_use]
72    pub fn is_empty(&self) -> bool {
73        self.bm25_terms.is_empty()
74    }
75
76    /// Convert into an index [`Document`](crate::index::Document), pairing with a
77    /// resolved URL, content hash, and embedding vector.
78    ///
79    /// Freshness fields (`fetched_at_ms`, `etag`, `pinned`) default to
80    /// unset/unpinned; callers that maintain a persisted knowledge base (the
81    /// facade) set them on the returned document before upserting.
82    #[must_use]
83    pub fn into_document(
84        self,
85        url: url::Url,
86        kind: ResourceKind,
87        content_hash: u64,
88        vector: Vec<f32>,
89    ) -> crate::index::Document {
90        // Resolve outbound hrefs against the page URL into canonical-URL keys — the
91        // persisted knowledge-graph edges. Deduplicated in document order and
92        // capped *before* sorting, so a hub page keeps its first (most
93        // prominent) links rather than a hash-random subset, then sorted for
94        // deterministic storage.
95        let mut seen =
96            std::collections::HashSet::with_capacity(self.links.len().min(MAX_EDGES_PER_DOC));
97        let mut edges: Vec<crate::url_key::UrlKey> = Vec::new();
98        for href in &self.links {
99            let Ok(target) = url.join(href) else { continue };
100            if target.scheme() != "http" && target.scheme() != "https" {
101                continue;
102            }
103            let key = crate::url_key::UrlKey::from_url(&target);
104            if seen.insert(key) {
105                edges.push(key);
106                if edges.len() == MAX_EDGES_PER_DOC {
107                    break;
108                }
109            }
110        }
111        edges.sort_unstable();
112
113        crate::index::Document {
114            url,
115            kind,
116            content_hash,
117            title: self.title,
118            snippet: self.snippet,
119            lang: self.lang,
120            tags: self.tags.into_iter().take(4).collect(),
121            terms: self.bm25_terms,
122            vector,
123            edges,
124            fetched_at_ms: 0,
125            etag: None,
126            pinned: false,
127        }
128    }
129}
130
131/// Distill an extracted document into a [`Descriptor`]. Shared by every
132/// content-type extractor so keywords/snippets/embed-text are produced uniformly.
133/// `headings` (H1–H3 / ATX) and `lang` are optional high-signal inputs; code and
134/// plain text pass an empty heading slice and `None`.
135#[must_use]
136pub fn distill(
137    resource: &Resource,
138    title: Option<CompactString>,
139    headings: &[CompactString],
140    body: &str,
141    links: Vec<CompactString>,
142    lang: Option<CompactString>,
143) -> Descriptor {
144    // BM25 terms: normalized title + heading tokens (high signal) then body, capped.
145    let mut bm25_terms: Vec<CompactString> = Vec::new();
146    if let Some(t) = &title {
147        bm25_terms.extend(text::normalized_tokens(t));
148    }
149    for h in headings {
150        bm25_terms.extend(text::normalized_tokens(h));
151    }
152    bm25_terms.extend(text::normalized_tokens(body).take(MAX_BM25_TERMS));
153    bm25_terms.truncate(MAX_BM25_TERMS);
154
155    // Keywords: term frequency over non-stopword, sufficiently-long tokens.
156    let mut freq: HashMap<CompactString, u32> = HashMap::new();
157    for term in &bm25_terms {
158        if term.len() >= MIN_KEYWORD_LEN && !text::is_stopword(term) {
159            *freq.entry(term.clone()).or_insert(0) += 1;
160        }
161    }
162    let mut ranked: Vec<(CompactString, u32)> = freq.into_iter().collect();
163    // Deterministic: by descending frequency, then term ascending.
164    ranked.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| a.0.cmp(&b.0)));
165    let keywords: SmallVec<[CompactString; MAX_KEYWORDS]> = ranked
166        .into_iter()
167        .take(MAX_KEYWORDS)
168        .map(|(t, _)| t)
169        .collect();
170
171    // Tags: URL path components.
172    let tags = path_tags(resource);
173
174    // Snippet: first SNIPPET_CHARS of whitespace-collapsed body.
175    let snippet = make_snippet(body);
176
177    // Headings retained for the descriptor (also folded into embed text below).
178    let heading_vec: SmallVec<[CompactString; MAX_HEADINGS]> =
179        headings.iter().take(MAX_HEADINGS).cloned().collect();
180
181    // Embed text: title + headings + keywords + path tail + body head.
182    let mut embed_text = String::new();
183    if let Some(t) = &title {
184        embed_text.push_str(t);
185        embed_text.push(' ');
186    }
187    for h in &heading_vec {
188        embed_text.push_str(h);
189        embed_text.push(' ');
190    }
191    for kw in &keywords {
192        embed_text.push_str(kw);
193        embed_text.push(' ');
194    }
195    for tag in &tags {
196        embed_text.push_str(tag);
197        embed_text.push(' ');
198    }
199    push_char_prefix(&mut embed_text, body, EMBED_BODY_CHARS);
200
201    Descriptor {
202        title,
203        keywords,
204        headings: heading_vec,
205        heading_levels: SmallVec::new(),
206        tags,
207        lang,
208        snippet,
209        embed_text,
210        bm25_terms,
211        links,
212    }
213}
214
215fn path_tags(resource: &Resource) -> SmallVec<[CompactString; 8]> {
216    let mut tags: SmallVec<[CompactString; 8]> = SmallVec::new();
217    for seg in resource.url.path().split('/') {
218        let seg = seg.trim();
219        if seg.is_empty() {
220            continue;
221        }
222        // strip an extension from the last segment
223        let base = seg.rsplit_once('.').map_or(seg, |(b, _)| b);
224        if base.is_empty() {
225            continue;
226        }
227        let norm = CompactString::from(text::normalize(base));
228        if norm.len() >= 2 && !tags.contains(&norm) {
229            tags.push(norm);
230        }
231    }
232    tags.truncate(8);
233    tags
234}
235
236fn make_snippet(body: &str) -> CompactString {
237    let mut snippet = String::with_capacity(SNIPPET_CHARS);
238    let mut prev_space = false;
239    for ch in body.chars() {
240        if snippet.chars().count() >= SNIPPET_CHARS {
241            break;
242        }
243        if ch.is_whitespace() {
244            if !prev_space && !snippet.is_empty() {
245                snippet.push(' ');
246                prev_space = true;
247            }
248        } else {
249            snippet.push(ch);
250            prev_space = false;
251        }
252    }
253    CompactString::from(snippet.trim())
254}
255
256fn push_char_prefix(out: &mut String, body: &str, max_chars: usize) {
257    let mut prev_space = false;
258    let mut count = 0;
259    for ch in body.chars() {
260        if count >= max_chars {
261            break;
262        }
263        if ch.is_whitespace() {
264            if !prev_space {
265                out.push(' ');
266                prev_space = true;
267                count += 1;
268            }
269        } else {
270            out.push(ch);
271            prev_space = false;
272            count += 1;
273        }
274    }
275}
276
277/// Extracts a [`Descriptor`] from raw document bytes.
278pub trait Extractor: Send + Sync {
279    /// Whether this extractor handles the given content kind.
280    fn handles(&self, kind: ResourceKind) -> bool;
281
282    /// Extract a descriptor from `bytes`.
283    fn extract(&self, resource: &Resource, bytes: &[u8]) -> Result<Descriptor>;
284}
285
286/// The default extractor: routes by [`ResourceKind`] to the right content handler.
287#[derive(Clone, Copy, Debug, Default)]
288pub struct AutoExtractor;
289
290impl Extractor for AutoExtractor {
291    fn handles(&self, _kind: ResourceKind) -> bool {
292        true
293    }
294
295    fn extract(&self, resource: &Resource, bytes: &[u8]) -> Result<Descriptor> {
296        match resource.kind {
297            #[cfg(feature = "html")]
298            ResourceKind::Html => html::extract(resource, bytes),
299            ResourceKind::Markdown => markdown::extract_markdown(resource, bytes),
300            ResourceKind::Code => markdown::extract_code(resource, bytes),
301            _ => markdown::extract_plain(resource, bytes),
302        }
303    }
304}
305
306#[cfg(test)]
307mod tests {
308    use super::*;
309    use url::Url;
310
311    fn resource(url: &str) -> Resource {
312        Resource::new(Url::parse(url).unwrap())
313    }
314
315    #[test]
316    fn distill_extracts_keywords_and_snippet() {
317        let r = resource("https://x.dev/docs/guide.md");
318        let body = "Kubernetes Autopilot clusters run Kubernetes without node management. \
319                    Autopilot manages nodes for you.";
320        let d = distill(
321            &r,
322            Some(CompactString::from("GKE Autopilot")),
323            &[],
324            body,
325            Vec::new(),
326            None,
327        );
328        // "autopilot" and "kubernetes" are the most frequent discriminative terms.
329        assert!(d.keywords.iter().any(|k| k == "autopilot"));
330        assert!(d.keywords.iter().any(|k| k == "kubernetes"));
331        // stopwords excluded
332        assert!(!d.keywords.iter().any(|k| k == "for" || k == "you"));
333        assert!(!d.snippet.is_empty());
334        assert!(d.embed_text.contains("GKE Autopilot"));
335    }
336
337    #[test]
338    fn distill_derives_path_tags() {
339        let r = resource("https://x.dev/docs/networking/psc.md");
340        let d = distill(&r, None, &[], "private service connect", Vec::new(), None);
341        assert!(d.tags.iter().any(|t| t == "docs"));
342        assert!(d.tags.iter().any(|t| t == "networking"));
343        assert!(d.tags.iter().any(|t| t == "psc")); // extension stripped
344    }
345
346    #[test]
347    fn empty_body_is_empty_descriptor() {
348        let r = resource("https://x.dev/empty");
349        let d = distill(&r, None, &[], "   ", Vec::new(), None);
350        assert!(d.is_empty());
351    }
352
353    #[test]
354    fn into_document_carries_fields() {
355        let r = resource("https://x.dev/a");
356        let d = distill(
357            &r,
358            Some(CompactString::from("T")),
359            &[],
360            "cat dog cat",
361            Vec::new(),
362            None,
363        );
364        let doc = d.into_document(r.url.clone(), ResourceKind::Markdown, 99, vec![0.1, 0.2]);
365        assert_eq!(doc.content_hash, 99);
366        assert_eq!(doc.title.as_deref(), Some("T"));
367        assert!(doc.terms.contains(&CompactString::from("cat")));
368    }
369}