ep: Add BM25 context retriever (#58526)

Release Notes:

- N/A

crates/edit_prediction_cli/src/retrieve_context.rs

@@ -27,6 +27,7 @@ pub enum ContextRetrievalType {
     EditHistory,
     EditHistoryFile,
     GitLog,
+    Bm25,
     OracleFile,
     #[default]
     All,
@@ -42,6 +43,7 @@ impl std::fmt::Display for ContextRetrievalType {
             ContextRetrievalType::EditHistory => write!(f, "edit-history"),
             ContextRetrievalType::EditHistoryFile => write!(f, "edit-history-file"),
             ContextRetrievalType::GitLog => write!(f, "git-log"),
+            ContextRetrievalType::Bm25 => write!(f, "bm25"),
             ContextRetrievalType::OracleFile => write!(f, "oracle-file"),
             ContextRetrievalType::All => write!(f, "all"),
             ContextRetrievalType::None => write!(f, "none"),
@@ -58,6 +60,7 @@ impl ContextRetrievalType {
             ContextRetrievalType::EditHistory => vec![ContextSource::EditHistory],
             ContextRetrievalType::EditHistoryFile => vec![ContextSource::EditHistoryFile],
             ContextRetrievalType::GitLog => vec![ContextSource::GitLog],
+            ContextRetrievalType::Bm25 => vec![ContextSource::Bm25],
             ContextRetrievalType::OracleFile => vec![ContextSource::OracleFile],
             ContextRetrievalType::All => {
                 let mut sources = vec![ContextSource::Lsp];
@@ -105,6 +108,8 @@ fn editable_context_sources() -> Vec<ContextSource> {
         ContextSource::EditHistory,
         ContextSource::EditHistoryFile,
         ContextSource::GitLog,
+        ContextSource::Bm25,
+        ContextSource::OracleFile,
     ]
 }
 

crates/edit_prediction_context/src/bm25_context.rs

@@ -0,0 +1,682 @@
+use crate::editable_context::EditHistoryContextEntry;
+use anyhow::{Context as _, Result, bail};
+use gpui::{AppContext as _, AsyncApp, Entity};
+use language::{Buffer, Point, ToPoint as _};
+use project::Project;
+use std::{
+    cmp::Ordering,
+    collections::{HashMap, HashSet},
+    fs,
+    ops::Range,
+    path::{Path, PathBuf},
+    time::Instant,
+};
+use text::Anchor;
+use util::command::new_command;
+
+const BM25_CONTEXT_QUERY_LINE_COUNT: u32 = 20;
+const BM25_CONTEXT_EDIT_HISTORY_QUERY_ENTRY_COUNT: usize = 8;
+const BM25_CONTEXT_CHUNK_LINE_COUNT: usize = 40;
+const BM25_CONTEXT_CHUNK_OVERLAP_LINE_COUNT: usize = 10;
+const BM25_CONTEXT_CHUNK_COUNT: usize = 12;
+const BM25_CONTEXT_MAX_CHUNKS_PER_FILE: usize = 3;
+const BM25_CONTEXT_MAX_FILE_BYTES: u64 = 1_000_000;
+const BM25_K1: f64 = 1.2;
+const BM25_B: f64 = 0.75;
+
+pub(super) struct Bm25ContextCandidate {
+    pub path: PathBuf,
+    pub row_range: Range<u32>,
+    pub order: usize,
+}
+
+pub async fn collect_bm25_context(
+    project: Entity<Project>,
+    active_buffer: Entity<Buffer>,
+    cursor_position: Anchor,
+    edit_history: &[EditHistoryContextEntry],
+    next_order: usize,
+    cx: &mut AsyncApp,
+) -> Vec<Bm25ContextCandidate> {
+    let Some(query) = build_query(&project, &active_buffer, cursor_position, edit_history, cx)
+    else {
+        return Vec::new();
+    };
+
+    let result = cx
+        .background_spawn(async move { collect_bm25_context_from_disk(query, next_order).await })
+        .await;
+
+    match result {
+        Ok(context) => context,
+        Err(error) => {
+            log::debug!("failed to collect BM25 context: {error:#}");
+            Vec::new()
+        }
+    }
+}
+
+struct Bm25ContextQuery {
+    worktree_abs_path: PathBuf,
+    worktree_root_name: String,
+    active_path: String,
+    cursor_excerpt: String,
+    edit_history_excerpts: Vec<String>,
+}
+
+fn build_query(
+    project: &Entity<Project>,
+    active_buffer: &Entity<Buffer>,
+    cursor_position: Anchor,
+    edit_history: &[EditHistoryContextEntry],
+    cx: &mut AsyncApp,
+) -> Option<Bm25ContextQuery> {
+    let (worktree_abs_path, worktree_root_name, active_path, cursor_excerpt) = cx.update(|cx| {
+        let buffer = active_buffer.read(cx);
+        let file = buffer.file()?;
+        let project = project.read(cx);
+        if !project.is_local() {
+            return None;
+        }
+        let worktree = project.worktree_for_id(file.worktree_id(cx), cx)?;
+        let worktree = worktree.read(cx);
+        if !worktree.is_local() {
+            return None;
+        }
+
+        let snapshot = buffer.snapshot();
+        let range = expanded_anchor_range(&snapshot, cursor_position..cursor_position);
+        let cursor_excerpt = snapshot.text_for_range(range).collect::<String>();
+
+        Some((
+            worktree.abs_path(),
+            worktree.root_name().as_unix_str().to_string(),
+            file.path().as_unix_str().to_string(),
+            cursor_excerpt,
+        ))
+    })?;
+
+    let edit_history_excerpts = edit_history
+        .iter()
+        .take(BM25_CONTEXT_EDIT_HISTORY_QUERY_ENTRY_COUNT)
+        .map(|entry| {
+            entry.buffer.read_with(cx, |buffer, _cx| {
+                let snapshot = buffer.snapshot();
+                let range = expanded_anchor_range(&snapshot, entry.edited_range.clone());
+                snapshot.text_for_range(range).collect::<String>()
+            })
+        })
+        .collect();
+
+    Some(Bm25ContextQuery {
+        worktree_abs_path: worktree_abs_path.to_path_buf(),
+        worktree_root_name,
+        active_path,
+        cursor_excerpt,
+        edit_history_excerpts,
+    })
+}
+
+fn expanded_anchor_range(
+    snapshot: &language::BufferSnapshot,
+    range: Range<Anchor>,
+) -> Range<Anchor> {
+    let start = range.start.to_point(snapshot);
+    let end = range.end.to_point(snapshot);
+    let start_row = start.row.saturating_sub(BM25_CONTEXT_QUERY_LINE_COUNT);
+    let end_row = end
+        .row
+        .saturating_add(BM25_CONTEXT_QUERY_LINE_COUNT)
+        .min(snapshot.max_point().row);
+    let start = snapshot.anchor_before(Point::new(start_row, 0));
+    let end = snapshot.anchor_after(Point::new(end_row, snapshot.line_len(end_row)));
+    start..end
+}
+
+async fn collect_bm25_context_from_disk(
+    query: Bm25ContextQuery,
+    next_order: usize,
+) -> Result<Vec<Bm25ContextCandidate>> {
+    let query_terms = query_terms(&query);
+    if query_terms.is_empty() {
+        return Ok(Vec::new());
+    }
+
+    let started_at = Instant::now();
+    let index = Bm25Index::build(&query.worktree_abs_path).await?;
+    let elapsed = started_at.elapsed();
+    log::debug!(
+        "built BM25 context index: candidate_files:{}, indexed_files:{}, indexed_bytes:{}, chunks:{}, terms:{}, latency:{elapsed:?}",
+        index.stats.candidate_file_count,
+        index.stats.indexed_file_count,
+        index.stats.indexed_bytes,
+        index.stats.document_count,
+        index.stats.term_count,
+    );
+
+    let candidates = index.search(&query_terms, &query.worktree_root_name, next_order);
+    log::debug!("selected {} BM25 context chunks", candidates.len());
+    Ok(candidates)
+}
+
+fn query_terms(query: &Bm25ContextQuery) -> HashMap<String, f64> {
+    let mut terms = HashMap::new();
+    add_query_terms(&mut terms, &query.active_path, 3.0);
+    add_query_terms(&mut terms, &query.cursor_excerpt, 1.0);
+    for excerpt in &query.edit_history_excerpts {
+        add_query_terms(&mut terms, excerpt, 2.0);
+    }
+    terms
+}
+
+fn add_query_terms(terms: &mut HashMap<String, f64>, text: &str, weight: f64) {
+    for token in tokenize(text) {
+        *terms.entry(token).or_default() += weight;
+    }
+}
+
+struct Bm25Index {
+    documents: Vec<Document>,
+    document_frequencies: HashMap<String, usize>,
+    average_document_len: f64,
+    stats: Bm25IndexStats,
+}
+
+#[derive(Default)]
+struct Bm25IndexStats {
+    candidate_file_count: usize,
+    indexed_file_count: usize,
+    indexed_bytes: u64,
+    document_count: usize,
+    term_count: usize,
+}
+
+struct Document {
+    relative_path: PathBuf,
+    row_range: Range<u32>,
+    term_frequencies: HashMap<String, usize>,
+    len: usize,
+}
+
+struct ScoredDocument {
+    document_index: usize,
+    score: f64,
+}
+
+struct DocumentsForFile {
+    documents: Vec<Document>,
+    byte_len: u64,
+}
+
+impl Bm25Index {
+    async fn build(worktree_abs_path: &Path) -> Result<Self> {
+        let relative_paths = git_ls_files(worktree_abs_path).await?;
+        let mut stats = Bm25IndexStats {
+            candidate_file_count: relative_paths.len(),
+            ..Default::default()
+        };
+        let mut documents = Vec::new();
+        for relative_path in relative_paths {
+            let Some(documents_for_file) = documents_for_file(worktree_abs_path, relative_path)
+            else {
+                continue;
+            };
+
+            stats.indexed_file_count += 1;
+            stats.indexed_bytes += documents_for_file.byte_len;
+            documents.extend(documents_for_file.documents);
+        }
+
+        let mut document_frequencies = HashMap::new();
+        let mut total_document_len = 0;
+        for document in &documents {
+            total_document_len += document.len;
+            let mut seen_terms = HashSet::new();
+            for term in document.term_frequencies.keys() {
+                if seen_terms.insert(term) {
+                    *document_frequencies.entry(term.clone()).or_default() += 1;
+                }
+            }
+        }
+
+        let average_document_len = if documents.is_empty() {
+            0.0
+        } else {
+            total_document_len as f64 / documents.len() as f64
+        };
+        stats.document_count = documents.len();
+        stats.term_count = document_frequencies.len();
+
+        Ok(Self {
+            documents,
+            document_frequencies,
+            average_document_len,
+            stats,
+        })
+    }
+
+    fn search(
+        &self,
+        query_terms: &HashMap<String, f64>,
+        worktree_root_name: &str,
+        next_order: usize,
+    ) -> Vec<Bm25ContextCandidate> {
+        if self.documents.is_empty() || self.average_document_len == 0.0 {
+            return Vec::new();
+        }
+
+        let mut scored_documents = self
+            .documents
+            .iter()
+            .enumerate()
+            .filter_map(|(document_index, document)| {
+                let score = self.score_document(document, query_terms);
+                (score > 0.0).then_some(ScoredDocument {
+                    document_index,
+                    score,
+                })
+            })
+            .collect::<Vec<_>>();
+
+        scored_documents.sort_by(|left, right| {
+            right
+                .score
+                .partial_cmp(&left.score)
+                .unwrap_or(Ordering::Equal)
+                .then_with(|| {
+                    self.documents[left.document_index]
+                        .relative_path
+                        .cmp(&self.documents[right.document_index].relative_path)
+                })
+                .then_with(|| {
+                    self.documents[left.document_index]
+                        .row_range
+                        .start
+                        .cmp(&self.documents[right.document_index].row_range.start)
+                })
+        });
+
+        let mut selected_documents = Vec::new();
+        let mut chunks_per_file = HashMap::<PathBuf, usize>::new();
+        for scored_document in scored_documents {
+            let document = &self.documents[scored_document.document_index];
+            let chunk_count = chunks_per_file
+                .entry(document.relative_path.clone())
+                .or_default();
+            if *chunk_count >= BM25_CONTEXT_MAX_CHUNKS_PER_FILE {
+                continue;
+            }
+
+            *chunk_count += 1;
+            selected_documents.push(Bm25ContextCandidate {
+                path: Path::new(&format!(
+                    "{}/{}",
+                    worktree_root_name,
+                    document.relative_path.to_string_lossy()
+                ))
+                .into(),
+                row_range: document.row_range.clone(),
+                order: next_order + selected_documents.len(),
+            });
+
+            if selected_documents.len() >= BM25_CONTEXT_CHUNK_COUNT {
+                break;
+            }
+        }
+
+        selected_documents
+    }
+
+    fn score_document(&self, document: &Document, query_terms: &HashMap<String, f64>) -> f64 {
+        let document_count = self.documents.len() as f64;
+        let document_len = document.len as f64;
+        let mut score = 0.0;
+
+        for (term, query_weight) in query_terms {
+            let Some(term_frequency) = document.term_frequencies.get(term) else {
+                continue;
+            };
+            let document_frequency = self
+                .document_frequencies
+                .get(term)
+                .copied()
+                .unwrap_or_default() as f64;
+            if document_frequency == 0.0 {
+                continue;
+            }
+
+            let inverse_document_frequency =
+                ((document_count - document_frequency + 0.5) / (document_frequency + 0.5) + 1.0)
+                    .ln();
+            let term_frequency = *term_frequency as f64;
+            let denominator = term_frequency
+                + BM25_K1
+                    * (1.0 - BM25_B + BM25_B * document_len / self.average_document_len.max(1.0));
+            score += query_weight * inverse_document_frequency * term_frequency * (BM25_K1 + 1.0)
+                / denominator;
+        }
+
+        score
+    }
+}
+
+async fn git_ls_files(worktree_abs_path: &Path) -> Result<Vec<PathBuf>> {
+    let output = new_command("git")
+        .arg("ls-files")
+        .arg("-z")
+        .current_dir(worktree_abs_path)
+        .output()
+        .await
+        .with_context(|| {
+            format!(
+                "failed to run git ls-files in {}",
+                worktree_abs_path.display()
+            )
+        })?;
+
+    if !output.status.success() {
+        let stderr = String::from_utf8_lossy(&output.stderr);
+        bail!(
+            "git ls-files failed in {} with status {}: {}",
+            worktree_abs_path.display(),
+            output.status,
+            stderr.trim()
+        );
+    }
+
+    let output =
+        String::from_utf8(output.stdout).context("git ls-files output was not valid UTF-8")?;
+    Ok(output
+        .split('\0')
+        .filter(|path| !path.is_empty())
+        .map(PathBuf::from)
+        .collect())
+}
+
+fn documents_for_file(
+    worktree_abs_path: &Path,
+    relative_path: PathBuf,
+) -> Option<DocumentsForFile> {
+    let absolute_path = worktree_abs_path.join(&relative_path);
+    let metadata = fs::symlink_metadata(&absolute_path).ok()?;
+    if metadata.file_type().is_symlink()
+        || !metadata.is_file()
+        || metadata.len() > BM25_CONTEXT_MAX_FILE_BYTES
+    {
+        return None;
+    }
+
+    let text = fs::read_to_string(&absolute_path).ok()?;
+    if text.is_empty() {
+        return None;
+    }
+
+    let byte_len = metadata.len();
+    let lines = lines(&text);
+    let path_tokens = tokenize(&relative_path.to_string_lossy());
+
+    let documents = chunk_line_ranges(
+        &lines,
+        BM25_CONTEXT_CHUNK_LINE_COUNT,
+        BM25_CONTEXT_CHUNK_OVERLAP_LINE_COUNT,
+    )
+    .into_iter()
+    .filter_map(|row_range| {
+        let chunk_text = text_for_line_range(&text, row_range.clone());
+        let mut term_frequencies = HashMap::new();
+        add_term_frequencies(&mut term_frequencies, tokenize(&chunk_text), 1);
+        add_term_frequencies(&mut term_frequencies, path_tokens.clone(), 2);
+        let len = term_frequencies.values().sum();
+        if len == 0 {
+            return None;
+        }
+
+        Some(Document {
+            relative_path: relative_path.clone(),
+            row_range: row_range.start as u32..row_range.end as u32,
+            term_frequencies,
+            len,
+        })
+    })
+    .collect::<Vec<_>>();
+
+    (!documents.is_empty()).then_some(DocumentsForFile {
+        documents,
+        byte_len,
+    })
+}
+
+fn add_term_frequencies(
+    term_frequencies: &mut HashMap<String, usize>,
+    tokens: Vec<String>,
+    weight: usize,
+) {
+    for token in tokens {
+        *term_frequencies.entry(token).or_default() += weight;
+    }
+}
+
+fn chunk_line_ranges(
+    lines: &[&str],
+    target_line_count: usize,
+    overlap_line_count: usize,
+) -> Vec<Range<usize>> {
+    if lines.is_empty() || target_line_count == 0 {
+        return Vec::new();
+    }
+
+    let mut ranges = Vec::new();
+    let mut start = 0;
+    while start < lines.len() {
+        let ideal_end = start.saturating_add(target_line_count).min(lines.len());
+        let mut end = ideal_end;
+        if ideal_end < lines.len()
+            && let Some(boundary) =
+                empty_line_boundary_near(lines, start, ideal_end, overlap_line_count)
+        {
+            end = boundary;
+        }
+        if end <= start {
+            end = ideal_end;
+        }
+        if end <= start {
+            break;
+        }
+
+        ranges.push(start..end);
+        if end == lines.len() {
+            break;
+        }
+
+        let next_start = end.saturating_sub(overlap_line_count);
+        start = if next_start <= start { end } else { next_start };
+    }
+
+    ranges
+}
+
+fn empty_line_boundary_near(
+    lines: &[&str],
+    start: usize,
+    ideal_end: usize,
+    overlap_line_count: usize,
+) -> Option<usize> {
+    let search_start = ideal_end.saturating_sub(overlap_line_count).max(start + 1);
+    let search_end = ideal_end
+        .saturating_add(overlap_line_count)
+        .min(lines.len());
+
+    (search_start..search_end)
+        .filter(|row| lines[*row].trim().is_empty())
+        .min_by_key(|row| row.abs_diff(ideal_end))
+        .map(|row| row + 1)
+}
+
+fn lines(text: &str) -> Vec<&str> {
+    text.split_inclusive('\n').collect()
+}
+
+fn text_for_line_range(text: &str, range: Range<usize>) -> String {
+    lines(text)
+        .into_iter()
+        .skip(range.start)
+        .take(range.end.saturating_sub(range.start))
+        .collect()
+}
+
+fn tokenize(text: &str) -> Vec<String> {
+    let mut tokens = Vec::new();
+    let mut segment = String::new();
+
+    for character in text.chars() {
+        if character.is_alphanumeric() || character == '_' || character == '-' {
+            segment.push(character);
+        } else {
+            push_segment_tokens(&segment, &mut tokens);
+            segment.clear();
+        }
+    }
+    push_segment_tokens(&segment, &mut tokens);
+
+    tokens
+}
+
+fn push_segment_tokens(segment: &str, tokens: &mut Vec<String>) {
+    if segment.is_empty() {
+        return;
+    }
+
+    let mut segment_tokens = Vec::new();
+    push_token(segment, &mut segment_tokens);
+    for part in segment.split(['_', '-']).filter(|part| !part.is_empty()) {
+        push_token(part, &mut segment_tokens);
+        for camel_part in camel_case_parts(part) {
+            push_token(camel_part, &mut segment_tokens);
+        }
+    }
+
+    let mut unique_segment_tokens = Vec::new();
+    for token in segment_tokens {
+        if !unique_segment_tokens.contains(&token) {
+            unique_segment_tokens.push(token);
+        }
+    }
+    tokens.extend(unique_segment_tokens);
+}
+
+fn camel_case_parts(text: &str) -> Vec<&str> {
+    let mut parts = Vec::new();
+    let mut start = 0;
+    let mut previous = None;
+
+    for (index, character) in text.char_indices() {
+        if index > 0
+            && character.is_uppercase()
+            && previous
+                .is_some_and(|previous: char| previous.is_lowercase() || previous.is_numeric())
+        {
+            parts.push(&text[start..index]);
+            start = index;
+        }
+        previous = Some(character);
+    }
+
+    if start < text.len() {
+        parts.push(&text[start..]);
+    }
+
+    parts
+}
+
+fn push_token(token: &str, tokens: &mut Vec<String>) {
+    let token = token.to_lowercase();
+    if token.len() <= 1
+        || token.len() > 128
+        || !token.chars().any(|character| character.is_alphabetic())
+    {
+        return;
+    }
+    tokens.push(token);
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_tokenize_splits_code_identifiers() {
+        let tokens =
+            tokenize("PrivateNetworkRequestPolicy foo_bar config/reg_default_16M_retrieval.json");
+
+        assert!(tokens.contains(&"privatenetworkrequestpolicy".to_string()));
+        assert!(tokens.contains(&"private".to_string()));
+        assert!(tokens.contains(&"network".to_string()));
+        assert!(tokens.contains(&"request".to_string()));
+        assert!(tokens.contains(&"policy".to_string()));
+        assert!(tokens.contains(&"foo_bar".to_string()));
+        assert!(tokens.contains(&"foo".to_string()));
+        assert!(tokens.contains(&"bar".to_string()));
+        assert!(tokens.contains(&"reg_default_16m_retrieval".to_string()));
+        assert!(tokens.contains(&"retrieval".to_string()));
+    }
+
+    #[test]
+    fn test_chunk_line_ranges_prefers_empty_line_boundaries_with_overlap() {
+        let text = "a\nb\n\nc\nd\ne\nf\n\ng\nh\ni\nj\n";
+        let lines = lines(text);
+        let ranges = chunk_line_ranges(&lines, 3, 1);
+
+        assert_eq!(ranges[0], 0..3);
+        assert!(ranges[1].start < ranges[0].end);
+    }
+
+    #[test]
+    fn test_bm25_ranks_matching_chunk() {
+        let documents = vec![
+            Document {
+                relative_path: PathBuf::from("src/unrelated.rs"),
+                row_range: 0..1,
+                term_frequencies: {
+                    let mut terms = HashMap::new();
+                    add_term_frequencies(&mut terms, tokenize("fn unrelated"), 1);
+                    terms
+                },
+                len: 2,
+            },
+            Document {
+                relative_path: PathBuf::from("src/network.rs"),
+                row_range: 0..1,
+                term_frequencies: {
+                    let mut terms = HashMap::new();
+                    add_term_frequencies(
+                        &mut terms,
+                        tokenize("fn update_private_network_request_policy"),
+                        1,
+                    );
+                    terms
+                },
+                len: 6,
+            },
+        ];
+        let mut document_frequencies = HashMap::new();
+        for document in &documents {
+            for term in document.term_frequencies.keys() {
+                *document_frequencies.entry(term.clone()).or_default() += 1;
+            }
+        }
+        let index = Bm25Index {
+            documents,
+            document_frequencies,
+            average_document_len: 4.0,
+            stats: Bm25IndexStats::default(),
+        };
+        let mut query = HashMap::new();
+        add_query_terms(&mut query, "PrivateNetworkRequestPolicy", 1.0);
+
+        let candidates = index.search(&query, "repo", 0);
+
+        assert_eq!(candidates[0].path, Path::new("repo/src/network.rs"));
+        assert_eq!(candidates[0].row_range, 0..1);
+        assert_eq!(candidates[0].order, 0);
+    }
+}

crates/edit_prediction_context/src/edit_prediction_context.rs

@@ -20,6 +20,7 @@ use util::rel_path::RelPath;
 use util::{RangeExt as _, ResultExt};
 
 mod assemble_excerpts;
+mod bm25_context;
 #[cfg(test)]
 mod edit_prediction_context_tests;
 mod editable_context;

crates/edit_prediction_context/src/editable_context.rs

@@ -11,7 +11,10 @@ use text::Anchor;
 use util::{paths::PathStyle, rel_path::RelPath};
 use zeta_prompt::{ContextSource, RelatedExcerpt, RelatedFile};
 
-use crate::git_log_context::build_git_log_index;
+use crate::{
+    bm25_context::{Bm25ContextCandidate, collect_bm25_context},
+    git_log_context::build_git_log_index,
+};
 
 /// This module contains collectors for editable context:
 /// excerpts or full files that are likely to be edited.
@@ -69,8 +72,27 @@ pub async fn collect_editable_context(
         collect_edit_history_file_context(&mut ranges_by_buffer, &edit_history, cx);
     }
     if context_sources.contains(&ContextSource::GitLog) {
-        collect_git_log_context(&mut ranges_by_buffer, project.clone(), active_buffer, cx).await;
+        collect_git_log_context(
+            &mut ranges_by_buffer,
+            project.clone(),
+            active_buffer.clone(),
+            cx,
+        )
+        .await;
     }
+
+    if context_sources.contains(&ContextSource::Bm25) {
+        collect_bm25_context_ranges(
+            &mut ranges_by_buffer,
+            project.clone(),
+            active_buffer,
+            cursor_position,
+            &edit_history,
+            cx,
+        )
+        .await;
+    }
+
     if context_sources.contains(&ContextSource::OracleFile) {
         collect_oracle_file_context(&mut ranges_by_buffer, project.clone(), oracle_paths, cx).await;
     }
@@ -298,6 +320,88 @@ fn collect_edit_history_file_context(
     }
 }
 
+async fn collect_bm25_context_ranges(
+    ranges_by_buffer: &mut RangesByBuffer,
+    project: Entity<Project>,
+    active_buffer: Entity<Buffer>,
+    cursor_position: Anchor,
+    edit_history: &[EditHistoryContextEntry],
+    cx: &mut AsyncApp,
+) {
+    let next_order = next_context_order(ranges_by_buffer);
+    let candidates = collect_bm25_context(
+        project.clone(),
+        active_buffer,
+        cursor_position,
+        edit_history,
+        next_order,
+        cx,
+    )
+    .await;
+
+    for candidate in candidates {
+        collect_bm25_candidate_context(ranges_by_buffer, &project, candidate, cx).await;
+    }
+}
+
+async fn collect_bm25_candidate_context(
+    ranges_by_buffer: &mut RangesByBuffer,
+    project: &Entity<Project>,
+    candidate: Bm25ContextCandidate,
+    cx: &mut AsyncApp,
+) {
+    let buffer = match open_buffer_for_path(project, &candidate.path, cx).await {
+        Ok(Some(buffer)) => buffer,
+        Ok(None) => {
+            log::debug!(
+                "failed to find BM25 context path: {}",
+                candidate.path.display()
+            );
+            return;
+        }
+        Err(error) => {
+            log::debug!(
+                "failed to open BM25 context path {}: {error:#}",
+                candidate.path.display()
+            );
+            return;
+        }
+    };
+
+    let Some(range) = buffer.read_with(cx, |buffer, _cx| {
+        anchor_range_for_row_range(&buffer.snapshot(), candidate.row_range.clone())
+    }) else {
+        return;
+    };
+
+    push_context_range(
+        ranges_by_buffer,
+        buffer,
+        range,
+        candidate.order,
+        ContextSource::Bm25,
+    );
+}
+
+fn anchor_range_for_row_range(
+    snapshot: &BufferSnapshot,
+    row_range: Range<u32>,
+) -> Option<Range<Anchor>> {
+    if row_range.start >= row_range.end || row_range.start > snapshot.max_point().row {
+        return None;
+    }
+
+    let max_point = snapshot.max_point();
+    let start = snapshot.anchor_before(Point::new(row_range.start, 0));
+    let end_point = if row_range.end > max_point.row {
+        max_point
+    } else {
+        Point::new(row_range.end, 0)
+    };
+    let end = snapshot.anchor_after(end_point);
+    Some(start..end)
+}
+
 async fn collect_oracle_file_context(
     ranges_by_buffer: &mut RangesByBuffer,
     project: Entity<Project>,
@@ -611,6 +715,7 @@ fn context_source_order(context_source: ContextSource) -> usize {
         ContextSource::EditHistory => 3,
         ContextSource::EditHistoryFile => 4,
         ContextSource::GitLog => 5,
-        ContextSource::OracleFile => 6,
+        ContextSource::Bm25 => 6,
+        ContextSource::OracleFile => 7,
     }
 }

crates/zeta_prompt/src/zeta_prompt.rs

@@ -254,6 +254,7 @@ pub enum ContextSource {
     EditHistory,
     EditHistoryFile,
     GitLog,
+    Bm25,
     OracleFile,
 }