sqz_engine/
url_indexer.rs

1use std::time::Duration;
2
3use chrono::{DateTime, Utc};
4use rusqlite::{params, Connection};
5use serde::{Deserialize, Serialize};
6
7use crate::error::{Result, SqzError};
8
9// ── Public types ──────────────────────────────────────────────────────────────
10
11/// Trait abstracting HTTP fetching so callers can provide their own
12/// implementation. This keeps sqz fully offline (Requirement 23).
13pub trait ContentFetcher: Send + Sync {
14    fn fetch(&self, url: &str) -> Result<String>;
15}
16
17/// A single indexed chunk (one heading section of a page).
18#[derive(Debug, Clone, Serialize, Deserialize)]
19pub struct IndexedChunk {
20    pub url: String,
21    pub heading: String,
22    pub body: String,
23    pub rank: f64,
24}
25
26/// Result of a `fetch_and_index` call.
27#[derive(Debug, Clone)]
28pub struct IndexResult {
29    pub url: String,
30    pub chunks_indexed: usize,
31    pub cached: bool,
32}
33
34// ── Schema ────────────────────────────────────────────────────────────────────
35
36const URL_SCHEMA: &str = r#"
37CREATE TABLE IF NOT EXISTS url_cache (
38    url          TEXT PRIMARY KEY,
39    fetched_at   TEXT NOT NULL,
40    ttl_secs     INTEGER NOT NULL,
41    markdown     TEXT NOT NULL
42);
43
44CREATE TABLE IF NOT EXISTS url_chunks (
45    id       INTEGER PRIMARY KEY AUTOINCREMENT,
46    url      TEXT NOT NULL,
47    heading  TEXT NOT NULL,
48    body     TEXT NOT NULL
49);
50
51CREATE VIRTUAL TABLE IF NOT EXISTS url_chunks_fts USING fts5(
52    url,
53    heading,
54    body,
55    content='url_chunks',
56    content_rowid='id',
57    tokenize='porter ascii'
58);
59
60CREATE TRIGGER IF NOT EXISTS url_chunks_ai AFTER INSERT ON url_chunks BEGIN
61    INSERT INTO url_chunks_fts(rowid, url, heading, body)
62    VALUES (new.id, new.url, new.heading, new.body);
63END;
64
65CREATE TRIGGER IF NOT EXISTS url_chunks_ad AFTER DELETE ON url_chunks BEGIN
66    INSERT INTO url_chunks_fts(url_chunks_fts, rowid, url, heading, body)
67    VALUES ('delete', old.id, old.url, old.heading, old.body);
68END;
69"#;
70
71pub(crate) fn apply_url_schema(conn: &Connection) -> rusqlite::Result<()> {
72    conn.execute_batch(URL_SCHEMA)
73}
74
75// ── HTML → Markdown conversion ────────────────────────────────────────────────
76
77/// Simple regex-free HTML to markdown converter. Handles common tags without
78/// pulling in any external dependency.
79fn html_to_markdown(html: &str) -> String {
80    let mut out = String::with_capacity(html.len());
81    let mut chars = html.chars().peekable();
82    let mut in_tag = false;
83    let mut tag_buf = String::new();
84    let mut skip_content = false;
85
86    while let Some(ch) = chars.next() {
87        if ch == '<' {
88            in_tag = true;
89            tag_buf.clear();
90            continue;
91        }
92        if in_tag {
93            if ch == '>' {
94                in_tag = false;
95                let tag_lower = tag_buf.to_ascii_lowercase();
96                let tag_name = tag_lower.split_whitespace().next().unwrap_or("");
97
98                match tag_name {
99                    "h1" => out.push_str("\n# "),
100                    "h2" => out.push_str("\n## "),
101                    "h3" => out.push_str("\n### "),
102                    "h4" => out.push_str("\n#### "),
103                    "h5" => out.push_str("\n##### "),
104                    "h6" => out.push_str("\n###### "),
105                    "/h1" | "/h2" | "/h3" | "/h4" | "/h5" | "/h6" => {
106                        out.push('\n');
107                    }
108                    "p" | "/p" | "br" | "br/" => out.push('\n'),
109                    "li" => out.push_str("\n- "),
110                    "pre" | "code" => out.push_str("\n```\n"),
111                    "/pre" | "/code" => out.push_str("\n```\n"),
112                    "strong" | "b" => out.push_str("**"),
113                    "/strong" | "/b" => out.push_str("**"),
114                    "em" | "i" => out.push('*'),
115                    "/em" | "/i" => out.push('*'),
116                    "script" | "style" | "noscript" => skip_content = true,
117                    "/script" | "/style" | "/noscript" => skip_content = false,
118                    _ => {} // ignore other tags
119                }
120            } else {
121                tag_buf.push(ch);
122            }
123            continue;
124        }
125        if skip_content {
126            continue;
127        }
128        // Decode common HTML entities
129        if ch == '&' {
130            let mut entity = String::new();
131            for ec in chars.by_ref() {
132                if ec == ';' {
133                    break;
134                }
135                entity.push(ec);
136                if entity.len() > 8 {
137                    break;
138                }
139            }
140            match entity.as_str() {
141                "amp" => out.push('&'),
142                "lt" => out.push('<'),
143                "gt" => out.push('>'),
144                "quot" => out.push('"'),
145                "apos" => out.push('\''),
146                "nbsp" => out.push(' '),
147                _ => {
148                    out.push('&');
149                    out.push_str(&entity);
150                    out.push(';');
151                }
152            }
153            continue;
154        }
155        out.push(ch);
156    }
157
158    // Collapse excessive blank lines
159    collapse_blank_lines(&out)
160}
161
162fn collapse_blank_lines(s: &str) -> String {
163    let mut result = String::with_capacity(s.len());
164    let mut blank_count = 0u32;
165    for line in s.lines() {
166        if line.trim().is_empty() {
167            blank_count += 1;
168            if blank_count <= 2 {
169                result.push('\n');
170            }
171        } else {
172            blank_count = 0;
173            result.push_str(line);
174            result.push('\n');
175        }
176    }
177    result.trim().to_string()
178}
179
180
181// ── Chunking by headings ──────────────────────────────────────────────────────
182
183/// A chunk is a heading + the body text until the next heading of equal or
184/// higher level.
185#[derive(Debug, Clone)]
186struct MdChunk {
187    heading: String,
188    body: String,
189}
190
191fn chunk_by_headings(markdown: &str) -> Vec<MdChunk> {
192    let mut chunks: Vec<MdChunk> = Vec::new();
193    let mut current_heading = String::new();
194    let mut current_body = String::new();
195
196    for line in markdown.lines() {
197        let trimmed = line.trim_start();
198        if trimmed.starts_with('#') {
199            // Flush previous chunk
200            if !current_heading.is_empty() || !current_body.trim().is_empty() {
201                chunks.push(MdChunk {
202                    heading: current_heading.clone(),
203                    body: current_body.trim().to_string(),
204                });
205            }
206            // Strip leading '#' chars and whitespace for the heading text
207            current_heading = trimmed.trim_start_matches('#').trim().to_string();
208            current_body.clear();
209        } else {
210            current_body.push_str(line);
211            current_body.push('\n');
212        }
213    }
214
215    // Flush last chunk
216    if !current_heading.is_empty() || !current_body.trim().is_empty() {
217        chunks.push(MdChunk {
218            heading: current_heading,
219            body: current_body.trim().to_string(),
220        });
221    }
222
223    chunks
224}
225
226// ── UrlIndexer ────────────────────────────────────────────────────────────────
227
228/// Fetches URLs, converts HTML to markdown, chunks by headings, and indexes
229/// into SQLite FTS5 for BM25-ranked section search.
230pub struct UrlIndexer {
231    db: Connection,
232    ttl: Duration,
233}
234
235impl UrlIndexer {
236    /// Create a new `UrlIndexer` backed by an in-memory SQLite database.
237    /// Useful for testing.
238    #[cfg(test)]
239    pub(crate) fn in_memory() -> Result<Self> {
240        let conn = Connection::open_in_memory()?;
241        apply_url_schema(&conn)?;
242        Ok(Self {
243            db: conn,
244            ttl: Duration::from_secs(24 * 3600),
245        })
246    }
247
248    /// Create a `UrlIndexer` from an existing connection (e.g. shared with
249    /// `SessionStore`). Applies the URL-specific schema tables.
250    pub fn from_connection(conn: Connection, ttl: Duration) -> Result<Self> {
251        apply_url_schema(&conn)?;
252        Ok(Self { db: conn, ttl })
253    }
254
255    /// Create a `UrlIndexer` with a custom TTL from an in-memory database.
256    #[cfg(test)]
257    pub(crate) fn in_memory_with_ttl(ttl: Duration) -> Result<Self> {
258        let conn = Connection::open_in_memory()?;
259        apply_url_schema(&conn)?;
260        Ok(Self { db: conn, ttl })
261    }
262
263    /// Fetch a URL via the provided `fetcher`, convert HTML to markdown,
264    /// chunk by headings, and index into FTS5.
265    ///
266    /// - If the URL was fetched within the TTL window and `force` is false,
267    ///   the cached version is used (no fetch).
268    /// - If `force` is true, the cache is bypassed and the URL is re-fetched.
269    pub fn fetch_and_index(
270        &self,
271        url: &str,
272        force: bool,
273        fetcher: &dyn ContentFetcher,
274    ) -> Result<IndexResult> {
275        // Check cache
276        if !force {
277            if let Some(fetched_at) = self.get_cache_time(url)? {
278                let age = Utc::now().signed_duration_since(fetched_at);
279                if age.num_seconds() < self.ttl.as_secs() as i64 {
280                    let count = self.chunk_count(url)?;
281                    return Ok(IndexResult {
282                        url: url.to_string(),
283                        chunks_indexed: count,
284                        cached: true,
285                    });
286                }
287            }
288        }
289
290        // Fetch
291        let html = fetcher.fetch(url)?;
292        let markdown = html_to_markdown(&html);
293        let chunks = chunk_by_headings(&markdown);
294
295        // Clear old data for this URL
296        self.delete_url(url)?;
297
298        // Store cache metadata
299        let now = Utc::now().to_rfc3339();
300        self.db.execute(
301            "INSERT OR REPLACE INTO url_cache (url, fetched_at, ttl_secs, markdown) VALUES (?1, ?2, ?3, ?4)",
302            params![url, now, self.ttl.as_secs() as i64, markdown],
303        )?;
304
305        // Index chunks
306        let chunk_count = chunks.len();
307        for chunk in &chunks {
308            self.db.execute(
309                "INSERT INTO url_chunks (url, heading, body) VALUES (?1, ?2, ?3)",
310                params![url, chunk.heading, chunk.body],
311            )?;
312        }
313
314        Ok(IndexResult {
315            url: url.to_string(),
316            chunks_indexed: chunk_count,
317            cached: false,
318        })
319    }
320
321    /// Search indexed URL content. Returns only matching sections (not full
322    /// documents), ranked by BM25.
323    pub fn search(&self, query: &str) -> Result<Vec<IndexedChunk>> {
324        let mut stmt = self.db.prepare(
325            r#"SELECT c.url, c.heading, c.body, rank
326               FROM url_chunks c
327               JOIN url_chunks_fts f ON c.id = f.rowid
328               WHERE url_chunks_fts MATCH ?1
329               ORDER BY rank"#,
330        )?;
331
332        let rows = stmt.query_map(params![query], |row| {
333            Ok(IndexedChunk {
334                url: row.get(0)?,
335                heading: row.get(1)?,
336                body: row.get(2)?,
337                rank: row.get(3)?,
338            })
339        })?;
340
341        let mut results = Vec::new();
342        for row in rows {
343            results.push(row?);
344        }
345        Ok(results)
346    }
347
348    // ── Helpers ───────────────────────────────────────────────────────────────
349
350    fn get_cache_time(&self, url: &str) -> Result<Option<DateTime<Utc>>> {
351        let result: rusqlite::Result<String> = self.db.query_row(
352            "SELECT fetched_at FROM url_cache WHERE url = ?1",
353            params![url],
354            |row| row.get(0),
355        );
356        match result {
357            Ok(ts) => {
358                let dt = ts
359                    .parse::<DateTime<Utc>>()
360                    .map_err(|e| SqzError::Other(format!("invalid timestamp: {e}")))?;
361                Ok(Some(dt))
362            }
363            Err(rusqlite::Error::QueryReturnedNoRows) => Ok(None),
364            Err(e) => Err(SqzError::SessionStore(e)),
365        }
366    }
367
368    fn chunk_count(&self, url: &str) -> Result<usize> {
369        let count: i64 = self.db.query_row(
370            "SELECT COUNT(*) FROM url_chunks WHERE url = ?1",
371            params![url],
372            |row| row.get(0),
373        )?;
374        Ok(count as usize)
375    }
376
377    fn delete_url(&self, url: &str) -> Result<()> {
378        // Delete chunks (triggers will clean FTS)
379        self.db.execute(
380            "DELETE FROM url_chunks WHERE url = ?1",
381            params![url],
382        )?;
383        self.db.execute(
384            "DELETE FROM url_cache WHERE url = ?1",
385            params![url],
386        )?;
387        Ok(())
388    }
389}
390
391
392// ── Tests ─────────────────────────────────────────────────────────────────────
393
394#[cfg(test)]
395mod tests {
396    use super::*;
397    use std::collections::HashMap;
398    use std::sync::atomic::{AtomicUsize, Ordering};
399    use std::time::Duration;
400
401    /// Mock fetcher that returns pre-configured HTML for given URLs.
402    struct MockFetcher {
403        pages: HashMap<String, String>,
404        call_count: AtomicUsize,
405    }
406
407    impl MockFetcher {
408        fn new() -> Self {
409            Self {
410                pages: HashMap::new(),
411                call_count: AtomicUsize::new(0),
412            }
413        }
414
415        fn add_page(&mut self, url: &str, html: &str) {
416            self.pages.insert(url.to_string(), html.to_string());
417        }
418
419        fn calls(&self) -> usize {
420            self.call_count.load(Ordering::SeqCst)
421        }
422    }
423
424    impl ContentFetcher for MockFetcher {
425        fn fetch(&self, url: &str) -> Result<String> {
426            self.call_count.fetch_add(1, Ordering::SeqCst);
427            self.pages
428                .get(url)
429                .cloned()
430                .ok_or_else(|| SqzError::Other(format!("URL not found: {url}")))
431        }
432    }
433
434    // ── html_to_markdown tests ────────────────────────────────────────────────
435
436    #[test]
437    fn test_html_to_markdown_headings() {
438        let html = "<h1>Title</h1><h2>Subtitle</h2><p>Body text</p>";
439        let md = html_to_markdown(html);
440        assert!(md.contains("# Title"));
441        assert!(md.contains("## Subtitle"));
442        assert!(md.contains("Body text"));
443    }
444
445    #[test]
446    fn test_html_to_markdown_strips_script_and_style() {
447        let html = "<p>visible</p><script>alert('x')</script><style>.a{}</style><p>also visible</p>";
448        let md = html_to_markdown(html);
449        assert!(md.contains("visible"));
450        assert!(md.contains("also visible"));
451        assert!(!md.contains("alert"));
452        assert!(!md.contains(".a{}"));
453    }
454
455    #[test]
456    fn test_html_to_markdown_entities() {
457        let html = "<p>A &amp; B &lt; C &gt; D</p>";
458        let md = html_to_markdown(html);
459        assert!(md.contains("A & B < C > D"));
460    }
461
462    #[test]
463    fn test_html_to_markdown_lists() {
464        let html = "<ul><li>one</li><li>two</li></ul>";
465        let md = html_to_markdown(html);
466        assert!(md.contains("- one"));
467        assert!(md.contains("- two"));
468    }
469
470    #[test]
471    fn test_html_to_markdown_bold_italic() {
472        let html = "<strong>bold</strong> and <em>italic</em>";
473        let md = html_to_markdown(html);
474        assert!(md.contains("**bold**"));
475        assert!(md.contains("*italic*"));
476    }
477
478    // ── chunk_by_headings tests ───────────────────────────────────────────────
479
480    #[test]
481    fn test_chunk_by_headings_basic() {
482        let md = "# Intro\nSome text\n## Details\nMore text\n## Conclusion\nFinal words";
483        let chunks = chunk_by_headings(md);
484        assert_eq!(chunks.len(), 3);
485        assert_eq!(chunks[0].heading, "Intro");
486        assert!(chunks[0].body.contains("Some text"));
487        assert_eq!(chunks[1].heading, "Details");
488        assert_eq!(chunks[2].heading, "Conclusion");
489    }
490
491    #[test]
492    fn test_chunk_by_headings_no_headings() {
493        let md = "Just plain text\nwith multiple lines";
494        let chunks = chunk_by_headings(md);
495        assert_eq!(chunks.len(), 1);
496        assert!(chunks[0].heading.is_empty());
497        assert!(chunks[0].body.contains("Just plain text"));
498    }
499
500    // ── UrlIndexer integration tests ──────────────────────────────────────────
501
502    #[test]
503    fn test_fetch_and_index_basic() {
504        let indexer = UrlIndexer::in_memory().unwrap();
505        let mut fetcher = MockFetcher::new();
506        fetcher.add_page(
507            "https://example.com/docs",
508            "<h1>API Reference</h1><p>This is the API docs.</p>\
509             <h2>Authentication</h2><p>Use bearer tokens for auth.</p>\
510             <h2>Endpoints</h2><p>GET /users returns user list.</p>",
511        );
512
513        let result = indexer
514            .fetch_and_index("https://example.com/docs", false, &fetcher)
515            .unwrap();
516
517        assert_eq!(result.url, "https://example.com/docs");
518        assert_eq!(result.chunks_indexed, 3);
519        assert!(!result.cached);
520        assert_eq!(fetcher.calls(), 1);
521    }
522
523    #[test]
524    fn test_cache_prevents_refetch() {
525        let indexer = UrlIndexer::in_memory().unwrap();
526        let mut fetcher = MockFetcher::new();
527        fetcher.add_page("https://example.com", "<h1>Hello</h1><p>World</p>");
528
529        // First fetch
530        let r1 = indexer
531            .fetch_and_index("https://example.com", false, &fetcher)
532            .unwrap();
533        assert!(!r1.cached);
534        assert_eq!(fetcher.calls(), 1);
535
536        // Second fetch — should be cached
537        let r2 = indexer
538            .fetch_and_index("https://example.com", false, &fetcher)
539            .unwrap();
540        assert!(r2.cached);
541        assert_eq!(fetcher.calls(), 1); // no additional fetch
542    }
543
544    #[test]
545    fn test_force_bypasses_cache() {
546        let indexer = UrlIndexer::in_memory().unwrap();
547        let mut fetcher = MockFetcher::new();
548        fetcher.add_page("https://example.com", "<h1>Hello</h1><p>World</p>");
549
550        indexer
551            .fetch_and_index("https://example.com", false, &fetcher)
552            .unwrap();
553        assert_eq!(fetcher.calls(), 1);
554
555        // Force re-fetch
556        let r2 = indexer
557            .fetch_and_index("https://example.com", true, &fetcher)
558            .unwrap();
559        assert!(!r2.cached);
560        assert_eq!(fetcher.calls(), 2);
561    }
562
563    #[test]
564    fn test_expired_ttl_refetches() {
565        // Use a 0-second TTL so everything is immediately expired
566        let indexer = UrlIndexer::in_memory_with_ttl(Duration::from_secs(0)).unwrap();
567        let mut fetcher = MockFetcher::new();
568        fetcher.add_page("https://example.com", "<h1>Hello</h1><p>World</p>");
569
570        indexer
571            .fetch_and_index("https://example.com", false, &fetcher)
572            .unwrap();
573        assert_eq!(fetcher.calls(), 1);
574
575        // TTL is 0s, so this should re-fetch
576        let r2 = indexer
577            .fetch_and_index("https://example.com", false, &fetcher)
578            .unwrap();
579        assert!(!r2.cached);
580        assert_eq!(fetcher.calls(), 2);
581    }
582
583    #[test]
584    fn test_search_returns_matching_sections_only() {
585        let indexer = UrlIndexer::in_memory().unwrap();
586        let mut fetcher = MockFetcher::new();
587        fetcher.add_page(
588            "https://example.com/docs",
589            "<h1>Overview</h1><p>General introduction to the system.</p>\
590             <h2>Authentication</h2><p>Use bearer tokens for authentication.</p>\
591             <h2>Rate Limiting</h2><p>Requests are limited to 100 per minute.</p>",
592        );
593
594        indexer
595            .fetch_and_index("https://example.com/docs", false, &fetcher)
596            .unwrap();
597
598        let results = indexer.search("authentication").unwrap();
599        // Should return only the authentication section, not the full doc
600        assert!(!results.is_empty());
601        assert!(results.iter().any(|r| r.heading == "Authentication"));
602        // The overview section should NOT match "authentication"
603        assert!(results.iter().all(|r| r.heading != "Overview"
604            || r.body.to_lowercase().contains("authentication")));
605    }
606
607    #[test]
608    fn test_search_across_multiple_urls() {
609        let indexer = UrlIndexer::in_memory().unwrap();
610        let mut fetcher = MockFetcher::new();
611        fetcher.add_page(
612            "https://a.com",
613            "<h1>Rust Guide</h1><p>Learn about ownership and borrowing.</p>",
614        );
615        fetcher.add_page(
616            "https://b.com",
617            "<h1>Python Guide</h1><p>Learn about decorators and generators.</p>",
618        );
619
620        indexer
621            .fetch_and_index("https://a.com", false, &fetcher)
622            .unwrap();
623        indexer
624            .fetch_and_index("https://b.com", false, &fetcher)
625            .unwrap();
626
627        let results = indexer.search("ownership").unwrap();
628        assert_eq!(results.len(), 1);
629        assert_eq!(results[0].url, "https://a.com");
630
631        let results = indexer.search("decorators").unwrap();
632        assert_eq!(results.len(), 1);
633        assert_eq!(results[0].url, "https://b.com");
634    }
635
636    #[test]
637    fn test_search_empty_query_returns_empty() {
638        let indexer = UrlIndexer::in_memory().unwrap();
639        let mut fetcher = MockFetcher::new();
640        fetcher.add_page("https://example.com", "<h1>Hello</h1><p>World</p>");
641        indexer
642            .fetch_and_index("https://example.com", false, &fetcher)
643            .unwrap();
644
645        // FTS5 with empty string — should error or return empty
646        let results = indexer.search("nonexistentxyzterm");
647        assert!(results.is_ok());
648        assert!(results.unwrap().is_empty());
649    }
650
651    #[test]
652    fn test_reindex_replaces_old_chunks() {
653        let indexer = UrlIndexer::in_memory().unwrap();
654        let mut fetcher = MockFetcher::new();
655        fetcher.add_page(
656            "https://example.com",
657            "<h1>Version 1</h1><p>Old content about widgets.</p>",
658        );
659
660        indexer
661            .fetch_and_index("https://example.com", false, &fetcher)
662            .unwrap();
663
664        // Update the page content
665        fetcher.pages.insert(
666            "https://example.com".to_string(),
667            "<h1>Version 2</h1><p>New content about gadgets.</p>".to_string(),
668        );
669
670        // Force re-index
671        indexer
672            .fetch_and_index("https://example.com", true, &fetcher)
673            .unwrap();
674
675        // Old content should be gone
676        let old = indexer.search("widgets").unwrap();
677        assert!(old.is_empty());
678
679        // New content should be found
680        let new = indexer.search("gadgets").unwrap();
681        assert_eq!(new.len(), 1);
682        assert!(new[0].body.contains("gadgets"));
683    }
684
685    #[test]
686    fn test_fetch_error_propagates() {
687        let indexer = UrlIndexer::in_memory().unwrap();
688        let fetcher = MockFetcher::new(); // no pages registered
689
690        let result = indexer.fetch_and_index("https://missing.com", false, &fetcher);
691        assert!(result.is_err());
692    }
693}
sqz_engine/url_indexer.rs

sqz_engine/
url_indexer.rs
