Implement M3.6.1: DocCorpusSource with heading-boundary chunking (196 tests)

crates/mem-ingest/Cargo.toml

@@ -15,6 +15,8 @@ anyhow = { workspace = true }
 thiserror = { workspace = true }
 tracing = { workspace = true }
 time = { workspace = true }
+walkdir = "2.5"
+sha2 = { workspace = true }
 
 [dev-dependencies]
 time = { workspace = true }

crates/mem-ingest/src/doc_corpus.rs

@@ -0,0 +1,562 @@
+//! Document corpus source for ingesting markdown documentation.
+//!
+//! Reads a directory tree of markdown files, splits on ATX headings,
+//! and emits one Record per section with breadcrumb paths attached.
+
+use anyhow::{anyhow, Result};
+use mem_chunk::RecordSource;
+use mem_core::{Provenance, Record, Role};
+use futures::stream::{self, Stream};
+use sha2::{Sha256, Digest};
+use std::path::{Path, PathBuf};
+use time::OffsetDateTime;
+use walkdir::WalkDir;
+
+/// Maximum file size to process (10MB).
+const MAX_FILE_SIZE: u64 = 10 * 1024 * 1024;
+
+/// Supported file extensions.
+const SUPPORTED_EXTENSIONS: &[&str] = &["md", "markdown", "txt"];
+
+/// A document section extracted from a markdown file.
+#[derive(Clone, Debug)]
+pub struct DocSection {
+    /// The heading path (e.g., "kubectl.md > Common Issues > CrashLoopBackOff")
+    pub breadcrumb: String,
+    /// The section content (heading + body)
+    pub content: String,
+    /// Source URI (absolute path or URL)
+    pub source_uri: String,
+    /// SHA256 hash of the original file
+    pub file_hash: String,
+    /// Whether this is a continuation of an over-long section
+    pub continuation: bool,
+    /// Section index within the file
+    pub section_index: usize,
+}
+
+/// Document corpus source that reads markdown files from a directory.
+pub struct DocCorpusSource {
+    root_path: PathBuf,
+    max_tokens: usize,
+}
+
+impl DocCorpusSource {
+    /// Create a new document corpus source.
+    ///
+    /// # Arguments
+    /// * `root_path` - Root directory to scan for markdown files
+    /// * `max_tokens` - Maximum tokens per section (for splitting over-long sections)
+    pub fn new(root_path: impl AsRef<Path>, max_tokens: usize) -> Self {
+        DocCorpusSource {
+            root_path: root_path.as_ref().to_path_buf(),
+            max_tokens,
+        }
+    }
+
+    /// Scan the directory and return all sections.
+    pub fn scan(&self) -> Result<Vec<DocSection>> {
+        let mut sections = Vec::new();
+
+        for entry in WalkDir::new(&self.root_path)
+            .follow_links(false) // Don't follow symlinks to avoid loops
+            .into_iter()
+            .filter_entry(|e| !is_hidden(e))
+        {
+            let entry = entry?;
+            let path = entry.path();
+
+            // Skip directories
+            if path.is_dir() {
+                continue;
+            }
+
+            // Check extension
+            let ext = path.extension()
+                .and_then(|e| e.to_str())
+                .map(|e| e.to_lowercase());
+            
+            if !ext.map(|e| SUPPORTED_EXTENSIONS.contains(&e.as_str())).unwrap_or(false) {
+                continue;
+            }
+
+            // Check file size
+            let metadata = std::fs::metadata(path)?;
+            if metadata.len() > MAX_FILE_SIZE {
+                continue;
+            }
+
+            // Read and process the file
+            let content = std::fs::read_to_string(path)?;
+            let file_hash = compute_sha256(&content);
+            let source_uri = path.canonicalize()?.to_string_lossy().to_string();
+            let filename = path.file_name()
+                .and_then(|n| n.to_str())
+                .unwrap_or("unknown");
+
+            let file_sections = split_on_headings(&content, filename, &source_uri, &file_hash, self.max_tokens);
+            sections.extend(file_sections);
+        }
+
+        Ok(sections)
+    }
+
+    /// Dry-run: scan and print the plan without making any network calls.
+    pub fn dry_run(&self) -> Result<DryRunReport> {
+        let sections = self.scan()?;
+        
+        let mut files = std::collections::HashMap::new();
+        for section in &sections {
+            let entry = files.entry(section.source_uri.clone()).or_insert_with(|| FileReport {
+                path: section.source_uri.clone(),
+                sections: Vec::new(),
+                total_tokens: 0,
+            });
+            let tokens = section.content.len() / 4; // Approximate
+            entry.sections.push(SectionReport {
+                breadcrumb: section.breadcrumb.clone(),
+                tokens,
+                continuation: section.continuation,
+            });
+            entry.total_tokens += tokens;
+        }
+
+        Ok(DryRunReport {
+            total_files: files.len(),
+            total_sections: sections.len(),
+            files: files.into_values().collect(),
+        })
+    }
+}
+
+/// Check if a directory entry is hidden (starts with .)
+fn is_hidden(entry: &walkdir::DirEntry) -> bool {
+    entry.file_name()
+        .to_str()
+        .map(|s| s.starts_with('.'))
+        .unwrap_or(false)
+}
+
+/// Compute SHA256 hash of content.
+fn compute_sha256(content: &str) -> String {
+    let mut hasher = Sha256::new();
+    hasher.update(content.as_bytes());
+    format!("{:x}", hasher.finalize())
+}
+
+/// Split markdown content on ATX headings.
+fn split_on_headings(
+    content: &str,
+    filename: &str,
+    source_uri: &str,
+    file_hash: &str,
+    max_tokens: usize,
+) -> Vec<DocSection> {
+    let mut sections = Vec::new();
+    let mut heading_stack: Vec<(usize, String)> = Vec::new(); // (level, text)
+    let mut current_content = String::new();
+    let mut section_index = 0;
+
+    // Helper to build breadcrumb from stack
+    let build_breadcrumb = |stack: &[(usize, String)]| -> String {
+        let mut parts = vec![filename.to_string()];
+        parts.extend(stack.iter().map(|(_, text)| text.clone()));
+        parts.join(" > ")
+    };
+
+    // Helper to flush current section
+    let flush_section = |content: &mut String, stack: &[(usize, String)], sections: &mut Vec<DocSection>, index: &mut usize| {
+        let trimmed = content.trim();
+        // Skip sections that are only a heading (no body content)
+        // A section with only one line that starts with # is just a heading
+        let lines: Vec<&str> = trimmed.lines().collect();
+        let has_body = lines.len() > 1 || (lines.len() == 1 && !lines[0].starts_with('#'));
+        if !trimmed.is_empty() && has_body {
+            let breadcrumb = build_breadcrumb(stack);
+            let tokens = trimmed.len() / 4;
+
+            // Check if section is over-long and needs splitting
+            if tokens > max_tokens {
+                let split_sections = split_overlong_section(trimmed, &breadcrumb, source_uri, file_hash, max_tokens, *index);
+                *index += split_sections.len();
+                sections.extend(split_sections);
+            } else {
+                // Prepend breadcrumb to content for embedding
+                let content_with_breadcrumb = format!("[{}]\n\n{}", breadcrumb, trimmed);
+                sections.push(DocSection {
+                    breadcrumb,
+                    content: content_with_breadcrumb,
+                    source_uri: source_uri.to_string(),
+                    file_hash: file_hash.to_string(),
+                    continuation: false,
+                    section_index: *index,
+                });
+                *index += 1;
+            }
+        }
+        content.clear();
+    };
+
+    for line in content.lines() {
+        // Check for ATX heading (# to ######)
+        if let Some(heading) = parse_atx_heading(line) {
+            // Flush previous section
+            flush_section(&mut current_content, &heading_stack, &mut sections, &mut section_index);
+
+            // Update heading stack
+            // Pop all headings at same or deeper level
+            while heading_stack.last().map(|(l, _)| *l >= heading.level).unwrap_or(false) {
+                heading_stack.pop();
+            }
+            heading_stack.push((heading.level, heading.text.clone()));
+
+            // Start new section with the heading
+            current_content.push_str(line);
+            current_content.push('\n');
+        } else {
+            // Regular content
+            current_content.push_str(line);
+            current_content.push('\n');
+        }
+    }
+
+    // Flush final section
+    flush_section(&mut current_content, &heading_stack, &mut sections, &mut section_index);
+
+    // If no sections were created (file has no headings), create one section for entire file
+    if sections.is_empty() && !content.trim().is_empty() {
+        let breadcrumb = filename.to_string();
+        let tokens = content.len() / 4;
+
+        if tokens > max_tokens {
+            sections = split_overlong_section(content.trim(), &breadcrumb, source_uri, file_hash, max_tokens, 0);
+        } else {
+            let content_with_breadcrumb = format!("[{}]\n\n{}", breadcrumb, content.trim());
+            sections.push(DocSection {
+                breadcrumb,
+                content: content_with_breadcrumb,
+                source_uri: source_uri.to_string(),
+                file_hash: file_hash.to_string(),
+                continuation: false,
+                section_index: 0,
+            });
+        }
+    }
+
+    sections
+}
+
+/// Parse an ATX heading line.
+struct AtxHeading {
+    level: usize,
+    text: String,
+}
+
+fn parse_atx_heading(line: &str) -> Option<AtxHeading> {
+    let trimmed = line.trim_start();
+    
+    // Count leading #s
+    let hashes = trimmed.chars().take_while(|&c| c == '#').count();
+    
+    // Must be 1-6 hashes followed by space or end of line
+    if hashes == 0 || hashes > 6 {
+        return None;
+    }
+
+    let rest = &trimmed[hashes..];
+    if !rest.is_empty() && !rest.starts_with(' ') && !rest.starts_with('\t') {
+        return None;
+    }
+
+    let text = rest.trim().trim_end_matches('#').trim().to_string();
+    
+    Some(AtxHeading { level: hashes, text })
+}
+
+/// Split an over-long section into smaller chunks.
+fn split_overlong_section(
+    content: &str,
+    breadcrumb: &str,
+    source_uri: &str,
+    file_hash: &str,
+    max_tokens: usize,
+    start_index: usize,
+) -> Vec<DocSection> {
+    let mut sections = Vec::new();
+    let max_chars = max_tokens * 4; // Approximate chars from tokens
+
+    // First, try to split on paragraph boundaries
+    let paragraphs: Vec<&str> = content.split("\n\n").collect();
+    
+    let mut current_content = String::new();
+    let mut chunk_index = 0;
+
+    for para in paragraphs {
+        let para_with_sep = if current_content.is_empty() {
+            para.to_string()
+        } else {
+            format!("\n\n{}", para)
+        };
+
+        if current_content.len() + para_with_sep.len() > max_chars && !current_content.is_empty() {
+            // Flush current chunk
+            let is_continuation = chunk_index > 0;
+            let content_with_breadcrumb = if is_continuation {
+                format!("[{} (continued)]\n\n{}", breadcrumb, current_content.trim())
+            } else {
+                format!("[{}]\n\n{}", breadcrumb, current_content.trim())
+            };
+
+            sections.push(DocSection {
+                breadcrumb: breadcrumb.to_string(),
+                content: content_with_breadcrumb,
+                source_uri: source_uri.to_string(),
+                file_hash: file_hash.to_string(),
+                continuation: is_continuation,
+                section_index: start_index + chunk_index,
+            });
+
+            current_content = para.to_string();
+            chunk_index += 1;
+        } else {
+            if current_content.is_empty() {
+                current_content = para.to_string();
+            } else {
+                current_content.push_str(&para_with_sep);
+            }
+        }
+    }
+
+    // Flush remaining content
+    if !current_content.trim().is_empty() {
+        let is_continuation = chunk_index > 0;
+        let content_with_breadcrumb = if is_continuation {
+            format!("[{} (continued)]\n\n{}", breadcrumb, current_content.trim())
+        } else {
+            format!("[{}]\n\n{}", breadcrumb, current_content.trim())
+        };
+
+        sections.push(DocSection {
+            breadcrumb: breadcrumb.to_string(),
+            content: content_with_breadcrumb,
+            source_uri: source_uri.to_string(),
+            file_hash: file_hash.to_string(),
+            continuation: is_continuation,
+            section_index: start_index + chunk_index,
+        });
+    }
+
+    // If a single paragraph is still too long, hard split
+    if sections.len() == 1 && sections[0].content.len() > max_chars {
+        sections = hard_split(&sections[0].content, breadcrumb, source_uri, file_hash, max_chars, start_index);
+    }
+
+    sections
+}
+
+/// Hard split content at character boundaries when paragraph splitting isn't enough.
+fn hard_split(
+    content: &str,
+    breadcrumb: &str,
+    source_uri: &str,
+    file_hash: &str,
+    max_chars: usize,
+    start_index: usize,
+) -> Vec<DocSection> {
+    let mut sections = Vec::new();
+    let chars: Vec<char> = content.chars().collect();
+    let mut chunk_index = 0;
+    let mut start = 0;
+
+    while start < chars.len() {
+        let end = (start + max_chars).min(chars.len());
+        let chunk: String = chars[start..end].iter().collect();
+        
+        let is_continuation = chunk_index > 0;
+        let content_with_breadcrumb = if is_continuation {
+            format!("[{} (continued)]\n\n{}", breadcrumb, chunk.trim())
+        } else {
+            chunk.to_string() // First chunk already has breadcrumb
+        };
+
+        sections.push(DocSection {
+            breadcrumb: breadcrumb.to_string(),
+            content: content_with_breadcrumb,
+            source_uri: source_uri.to_string(),
+            file_hash: file_hash.to_string(),
+            continuation: is_continuation,
+            section_index: start_index + chunk_index,
+        });
+
+        start = end;
+        chunk_index += 1;
+    }
+
+    sections
+}
+
+/// Dry-run report.
+#[derive(Debug)]
+pub struct DryRunReport {
+    pub total_files: usize,
+    pub total_sections: usize,
+    pub files: Vec<FileReport>,
+}
+
+/// Per-file report in dry-run.
+#[derive(Debug)]
+pub struct FileReport {
+    pub path: String,
+    pub sections: Vec<SectionReport>,
+    pub total_tokens: usize,
+}
+
+/// Per-section report in dry-run.
+#[derive(Debug)]
+pub struct SectionReport {
+    pub breadcrumb: String,
+    pub tokens: usize,
+    pub continuation: bool,
+}
+
+impl std::fmt::Display for DryRunReport {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        writeln!(f, "=== Document Corpus Dry Run ===")?;
+        writeln!(f, "Total files: {}", self.total_files)?;
+        writeln!(f, "Total sections: {}", self.total_sections)?;
+        writeln!(f)?;
+
+        for file in &self.files {
+            writeln!(f, "📄 {}", file.path)?;
+            writeln!(f, "   Total tokens: ~{}", file.total_tokens)?;
+            for section in &file.sections {
+                let cont = if section.continuation { " (cont)" } else { "" };
+                writeln!(f, "   └─ {} (~{} tokens){}", section.breadcrumb, section.tokens, cont)?;
+            }
+            writeln!(f)?;
+        }
+
+        Ok(())
+    }
+}
+
+impl RecordSource for DocCorpusSource {
+    fn records(self) -> Box<dyn Stream<Item = Result<Record, String>> + Unpin> {
+        // Scan sections synchronously (blocking is okay for file I/O)
+        let sections = match self.scan() {
+            Ok(s) => s,
+            Err(e) => {
+                return Box::new(stream::iter(vec![
+                    Err(format!("Failed to scan corpus: {}", e))
+                ]));
+            }
+        };
+
+        let records: Vec<Result<Record, String>> = sections
+            .into_iter()
+            .enumerate()
+            .map(|(i, section)| {
+                Ok(Record {
+                    role: Role::User, // Documentation is always "user" content
+                    text: section.content,
+                    timestamp: OffsetDateTime::now_utc(),
+                    provenance: Provenance {
+                        source_id: section.source_uri,
+                        offset: i as u64,
+                    },
+                })
+            })
+            .collect();
+
+        Box::new(stream::iter(records))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_parse_atx_heading() {
+        assert!(parse_atx_heading("# Heading").is_some());
+        assert!(parse_atx_heading("## Sub").is_some());
+        assert!(parse_atx_heading("###### Deep").is_some());
+        assert!(parse_atx_heading("####### Too deep").is_none());
+        assert!(parse_atx_heading("#NoSpace").is_none());
+        assert!(parse_atx_heading("Not a heading").is_none());
+        
+        let h = parse_atx_heading("## My Heading ##").unwrap();
+        assert_eq!(h.level, 2);
+        assert_eq!(h.text, "My Heading");
+    }
+
+    #[test]
+    fn test_split_on_headings_simple() {
+        let content = r#"# Title
+
+Some intro text.
+
+## Section 1
+
+Content 1.
+
+## Section 2
+
+Content 2.
+"#;
+        let sections = split_on_headings(content, "test.md", "/test.md", "abc123", 5000);
+        
+        assert_eq!(sections.len(), 3);
+        assert_eq!(sections[0].breadcrumb, "test.md > Title");
+        assert_eq!(sections[1].breadcrumb, "test.md > Title > Section 1");
+        assert_eq!(sections[2].breadcrumb, "test.md > Title > Section 2");
+    }
+
+    #[test]
+    fn test_split_on_headings_nested() {
+        let content = r#"# Top
+
+## A
+
+### A1
+
+Content A1.
+
+### A2
+
+Content A2.
+
+## B
+
+Content B.
+"#;
+        let sections = split_on_headings(content, "nested.md", "/nested.md", "xyz", 5000);
+        
+        // Sections without body content (only headings) are skipped
+        // So we get: A1, A2, B (3 sections, not 4)
+        assert_eq!(sections.len(), 3);
+        assert_eq!(sections[0].breadcrumb, "nested.md > Top > A > A1");
+        assert_eq!(sections[1].breadcrumb, "nested.md > Top > A > A2");
+        assert_eq!(sections[2].breadcrumb, "nested.md > Top > B");
+    }
+
+    #[test]
+    fn test_no_headings() {
+        let content = "Just some text without any headings.";
+        let sections = split_on_headings(content, "plain.md", "/plain.md", "hash", 5000);
+        
+        assert_eq!(sections.len(), 1);
+        assert_eq!(sections[0].breadcrumb, "plain.md");
+        assert!(!sections[0].continuation);
+    }
+
+    #[test]
+    fn test_overlong_section_splits() {
+        let content = format!("# Big Section\n\n{}", "word ".repeat(2000)); // ~10000 chars
+        let sections = split_on_headings(&content, "big.md", "/big.md", "hash", 500); // Force split
+        
+        assert!(sections.len() > 1);
+        assert!(!sections[0].continuation);
+        assert!(sections[1].continuation);
+    }
+}

crates/mem-ingest/src/lib.rs

@@ -1,5 +1,7 @@
 pub mod pi_session;
 pub mod claude_transcript;
+pub mod doc_corpus;
 
 pub use pi_session::PiSessionSource;
 pub use claude_transcript::ClaudeTranscriptSource;
+pub use doc_corpus::{DocCorpusSource, DocSection, DryRunReport};