tenant.openagents/omega

1use crate::editable_context::EditHistoryContextEntry;
2use anyhow::{Context as _, Result, bail};
3use gpui::{AppContext as _, AsyncApp, Entity};
4use language::{Buffer, Point, ToPoint as _};
5use project::Project;
6use std::{
7    cmp::Ordering,
8    collections::{HashMap, HashSet},
9    fs,
10    ops::Range,
11    path::{Path, PathBuf},
12    time::Instant,
13};
14use text::Anchor;
15use util::command::new_command;
16
17const BM25_CONTEXT_QUERY_LINE_COUNT: u32 = 20;
18const BM25_CONTEXT_EDIT_HISTORY_QUERY_ENTRY_COUNT: usize = 8;
19const BM25_CONTEXT_CHUNK_LINE_COUNT: usize = 40;
20const BM25_CONTEXT_CHUNK_OVERLAP_LINE_COUNT: usize = 10;
21const BM25_CONTEXT_CHUNK_COUNT: usize = 12;
22const BM25_CONTEXT_MAX_CHUNKS_PER_FILE: usize = 3;
23const BM25_CONTEXT_MAX_FILE_BYTES: u64 = 1_000_000;
24const BM25_K1: f64 = 1.2;
25const BM25_B: f64 = 0.75;
26
27pub(super) struct Bm25ContextCandidate {
28    pub path: PathBuf,
29    pub row_range: Range<u32>,
30    pub order: usize,
31}
32
33pub async fn collect_bm25_context(
34    project: Entity<Project>,
35    active_buffer: Entity<Buffer>,
36    cursor_position: Anchor,
37    edit_history: &[EditHistoryContextEntry],
38    next_order: usize,
39    cx: &mut AsyncApp,
40) -> Vec<Bm25ContextCandidate> {
41    let Some(query) = build_query(&project, &active_buffer, cursor_position, edit_history, cx)
42    else {
43        return Vec::new();
44    };
45
46    let result = cx
47        .background_spawn(async move { collect_bm25_context_from_disk(query, next_order).await })
48        .await;
49
50    match result {
51        Ok(context) => context,
52        Err(error) => {
53            log::debug!("failed to collect BM25 context: {error:#}");
54            Vec::new()
55        }
56    }
57}
58
59struct Bm25ContextQuery {
60    worktree_abs_path: PathBuf,
61    worktree_root_name: String,
62    active_path: String,
63    cursor_excerpt: String,
64    edit_history_excerpts: Vec<String>,
65}
66
67fn build_query(
68    project: &Entity<Project>,
69    active_buffer: &Entity<Buffer>,
70    cursor_position: Anchor,
71    edit_history: &[EditHistoryContextEntry],
72    cx: &mut AsyncApp,
73) -> Option<Bm25ContextQuery> {
74    let (worktree_abs_path, worktree_root_name, active_path, cursor_excerpt) = cx.update(|cx| {
75        let buffer = active_buffer.read(cx);
76        let file = buffer.file()?;
77        let project = project.read(cx);
78        if !project.is_local() {
79            return None;
80        }
81        let worktree = project.worktree_for_id(file.worktree_id(cx), cx)?;
82        let worktree = worktree.read(cx);
83        if !worktree.is_local() {
84            return None;
85        }
86
87        let snapshot = buffer.snapshot();
88        let range = expanded_anchor_range(&snapshot, cursor_position..cursor_position);
89        let cursor_excerpt = snapshot.text_for_range(range).collect::<String>();
90
91        Some((
92            worktree.abs_path(),
93            worktree.root_name().as_unix_str().to_string(),
94            file.path().as_unix_str().to_string(),
95            cursor_excerpt,
96        ))
97    })?;
98
99    let edit_history_excerpts = edit_history
100        .iter()
101        .take(BM25_CONTEXT_EDIT_HISTORY_QUERY_ENTRY_COUNT)
102        .map(|entry| {
103            entry.buffer.read_with(cx, |buffer, _cx| {
104                let snapshot = buffer.snapshot();
105                let range = expanded_anchor_range(&snapshot, entry.edited_range.clone());
106                snapshot.text_for_range(range).collect::<String>()
107            })
108        })
109        .collect();
110
111    Some(Bm25ContextQuery {
112        worktree_abs_path: worktree_abs_path.to_path_buf(),
113        worktree_root_name,
114        active_path,
115        cursor_excerpt,
116        edit_history_excerpts,
117    })
118}
119
120fn expanded_anchor_range(
121    snapshot: &language::BufferSnapshot,
122    range: Range<Anchor>,
123) -> Range<Anchor> {
124    let start = range.start.to_point(snapshot);
125    let end = range.end.to_point(snapshot);
126    let start_row = start.row.saturating_sub(BM25_CONTEXT_QUERY_LINE_COUNT);
127    let end_row = end
128        .row
129        .saturating_add(BM25_CONTEXT_QUERY_LINE_COUNT)
130        .min(snapshot.max_point().row);
131    let start = snapshot.anchor_before(Point::new(start_row, 0));
132    let end = snapshot.anchor_after(Point::new(end_row, snapshot.line_len(end_row)));
133    start..end
134}
135
136async fn collect_bm25_context_from_disk(
137    query: Bm25ContextQuery,
138    next_order: usize,
139) -> Result<Vec<Bm25ContextCandidate>> {
140    let query_terms = query_terms(&query);
141    if query_terms.is_empty() {
142        return Ok(Vec::new());
143    }
144
145    let started_at = Instant::now();
146    let index = Bm25Index::build(&query.worktree_abs_path).await?;
147    let elapsed = started_at.elapsed();
148    log::debug!(
149        "built BM25 context index: candidate_files:{}, indexed_files:{}, indexed_bytes:{}, chunks:{}, terms:{}, latency:{elapsed:?}",
150        index.stats.candidate_file_count,
151        index.stats.indexed_file_count,
152        index.stats.indexed_bytes,
153        index.stats.document_count,
154        index.stats.term_count,
155    );
156
157    let candidates = index.search(&query_terms, &query.worktree_root_name, next_order);
158    log::debug!("selected {} BM25 context chunks", candidates.len());
159    Ok(candidates)
160}
161
162fn query_terms(query: &Bm25ContextQuery) -> HashMap<String, f64> {
163    let mut terms = HashMap::new();
164    add_query_terms(&mut terms, &query.active_path, 3.0);
165    add_query_terms(&mut terms, &query.cursor_excerpt, 1.0);
166    for excerpt in &query.edit_history_excerpts {
167        add_query_terms(&mut terms, excerpt, 2.0);
168    }
169    terms
170}
171
172fn add_query_terms(terms: &mut HashMap<String, f64>, text: &str, weight: f64) {
173    for token in tokenize(text) {
174        *terms.entry(token).or_default() += weight;
175    }
176}
177
178struct Bm25Index {
179    documents: Vec<Document>,
180    document_frequencies: HashMap<String, usize>,
181    average_document_len: f64,
182    stats: Bm25IndexStats,
183}
184
185#[derive(Default)]
186struct Bm25IndexStats {
187    candidate_file_count: usize,
188    indexed_file_count: usize,
189    indexed_bytes: u64,
190    document_count: usize,
191    term_count: usize,
192}
193
194struct Document {
195    relative_path: PathBuf,
196    row_range: Range<u32>,
197    term_frequencies: HashMap<String, usize>,
198    len: usize,
199}
200
201struct ScoredDocument {
202    document_index: usize,
203    score: f64,
204}
205
206struct DocumentsForFile {
207    documents: Vec<Document>,
208    byte_len: u64,
209}
210
211impl Bm25Index {
212    async fn build(worktree_abs_path: &Path) -> Result<Self> {
213        let relative_paths = git_ls_files(worktree_abs_path).await?;
214        let mut stats = Bm25IndexStats {
215            candidate_file_count: relative_paths.len(),
216            ..Default::default()
217        };
218        let mut documents = Vec::new();
219        for relative_path in relative_paths {
220            let Some(documents_for_file) = documents_for_file(worktree_abs_path, relative_path)
221            else {
222                continue;
223            };
224
225            stats.indexed_file_count += 1;
226            stats.indexed_bytes += documents_for_file.byte_len;
227            documents.extend(documents_for_file.documents);
228        }
229
230        let mut document_frequencies = HashMap::new();
231        let mut total_document_len = 0;
232        for document in &documents {
233            total_document_len += document.len;
234            let mut seen_terms = HashSet::new();
235            for term in document.term_frequencies.keys() {
236                if seen_terms.insert(term) {
237                    *document_frequencies.entry(term.clone()).or_default() += 1;
238                }
239            }
240        }
241
242        let average_document_len = if documents.is_empty() {
243            0.0
244        } else {
245            total_document_len as f64 / documents.len() as f64
246        };
247        stats.document_count = documents.len();
248        stats.term_count = document_frequencies.len();
249
250        Ok(Self {
251            documents,
252            document_frequencies,
253            average_document_len,
254            stats,
255        })
256    }
257
258    fn search(
259        &self,
260        query_terms: &HashMap<String, f64>,
261        worktree_root_name: &str,
262        next_order: usize,
263    ) -> Vec<Bm25ContextCandidate> {
264        if self.documents.is_empty() || self.average_document_len == 0.0 {
265            return Vec::new();
266        }
267
268        let mut scored_documents = self
269            .documents
270            .iter()
271            .enumerate()
272            .filter_map(|(document_index, document)| {
273                let score = self.score_document(document, query_terms);
274                (score > 0.0).then_some(ScoredDocument {
275                    document_index,
276                    score,
277                })
278            })
279            .collect::<Vec<_>>();
280
281        scored_documents.sort_by(|left, right| {
282            right
283                .score
284                .partial_cmp(&left.score)
285                .unwrap_or(Ordering::Equal)
286                .then_with(|| {
287                    self.documents[left.document_index]
288                        .relative_path
289                        .cmp(&self.documents[right.document_index].relative_path)
290                })
291                .then_with(|| {
292                    self.documents[left.document_index]
293                        .row_range
294                        .start
295                        .cmp(&self.documents[right.document_index].row_range.start)
296                })
297        });
298
299        let mut selected_documents = Vec::new();
300        let mut chunks_per_file = HashMap::<PathBuf, usize>::new();
301        for scored_document in scored_documents {
302            let document = &self.documents[scored_document.document_index];
303            let chunk_count = chunks_per_file
304                .entry(document.relative_path.clone())
305                .or_default();
306            if *chunk_count >= BM25_CONTEXT_MAX_CHUNKS_PER_FILE {
307                continue;
308            }
309
310            *chunk_count += 1;
311            selected_documents.push(Bm25ContextCandidate {
312                path: Path::new(&format!(
313                    "{}/{}",
314                    worktree_root_name,
315                    document.relative_path.to_string_lossy()
316                ))
317                .into(),
318                row_range: document.row_range.clone(),
319                order: next_order + selected_documents.len(),
320            });
321
322            if selected_documents.len() >= BM25_CONTEXT_CHUNK_COUNT {
323                break;
324            }
325        }
326
327        selected_documents
328    }
329
330    fn score_document(&self, document: &Document, query_terms: &HashMap<String, f64>) -> f64 {
331        let document_count = self.documents.len() as f64;
332        let document_len = document.len as f64;
333        let mut score = 0.0;
334
335        for (term, query_weight) in query_terms {
336            let Some(term_frequency) = document.term_frequencies.get(term) else {
337                continue;
338            };
339            let document_frequency = self
340                .document_frequencies
341                .get(term)
342                .copied()
343                .unwrap_or_default() as f64;
344            if document_frequency == 0.0 {
345                continue;
346            }
347
348            let inverse_document_frequency =
349                ((document_count - document_frequency + 0.5) / (document_frequency + 0.5) + 1.0)
350                    .ln();
351            let term_frequency = *term_frequency as f64;
352            let denominator = term_frequency
353                + BM25_K1
354                    * (1.0 - BM25_B + BM25_B * document_len / self.average_document_len.max(1.0));
355            score += query_weight * inverse_document_frequency * term_frequency * (BM25_K1 + 1.0)
356                / denominator;
357        }
358
359        score
360    }
361}
362
363async fn git_ls_files(worktree_abs_path: &Path) -> Result<Vec<PathBuf>> {
364    let output = new_command("git")
365        .arg("ls-files")
366        .arg("-z")
367        .current_dir(worktree_abs_path)
368        .output()
369        .await
370        .with_context(|| {
371            format!(
372                "failed to run git ls-files in {}",
373                worktree_abs_path.display()
374            )
375        })?;
376
377    if !output.status.success() {
378        let stderr = String::from_utf8_lossy(&output.stderr);
379        bail!(
380            "git ls-files failed in {} with status {}: {}",
381            worktree_abs_path.display(),
382            output.status,
383            stderr.trim()
384        );
385    }
386
387    let output =
388        String::from_utf8(output.stdout).context("git ls-files output was not valid UTF-8")?;
389    Ok(output
390        .split('\0')
391        .filter(|path| !path.is_empty())
392        .map(PathBuf::from)
393        .collect())
394}
395
396fn documents_for_file(
397    worktree_abs_path: &Path,
398    relative_path: PathBuf,
399) -> Option<DocumentsForFile> {
400    let absolute_path = worktree_abs_path.join(&relative_path);
401    let metadata = fs::symlink_metadata(&absolute_path).ok()?;
402    if metadata.file_type().is_symlink()
403        || !metadata.is_file()
404        || metadata.len() > BM25_CONTEXT_MAX_FILE_BYTES
405    {
406        return None;
407    }
408
409    let text = fs::read_to_string(&absolute_path).ok()?;
410    if text.is_empty() {
411        return None;
412    }
413
414    let byte_len = metadata.len();
415    let lines = lines(&text);
416    let path_tokens = tokenize(&relative_path.to_string_lossy());
417
418    let documents = chunk_line_ranges(
419        &lines,
420        BM25_CONTEXT_CHUNK_LINE_COUNT,
421        BM25_CONTEXT_CHUNK_OVERLAP_LINE_COUNT,
422    )
423    .into_iter()
424    .filter_map(|row_range| {
425        let chunk_text = text_for_line_range(&text, row_range.clone());
426        let mut term_frequencies = HashMap::new();
427        add_term_frequencies(&mut term_frequencies, tokenize(&chunk_text), 1);
428        add_term_frequencies(&mut term_frequencies, path_tokens.clone(), 2);
429        let len = term_frequencies.values().sum();
430        if len == 0 {
431            return None;
432        }
433
434        Some(Document {
435            relative_path: relative_path.clone(),
436            row_range: row_range.start as u32..row_range.end as u32,
437            term_frequencies,
438            len,
439        })
440    })
441    .collect::<Vec<_>>();
442
443    (!documents.is_empty()).then_some(DocumentsForFile {
444        documents,
445        byte_len,
446    })
447}
448
449fn add_term_frequencies(
450    term_frequencies: &mut HashMap<String, usize>,
451    tokens: Vec<String>,
452    weight: usize,
453) {
454    for token in tokens {
455        *term_frequencies.entry(token).or_default() += weight;
456    }
457}
458
459fn chunk_line_ranges(
460    lines: &[&str],
461    target_line_count: usize,
462    overlap_line_count: usize,
463) -> Vec<Range<usize>> {
464    if lines.is_empty() || target_line_count == 0 {
465        return Vec::new();
466    }
467
468    let mut ranges = Vec::new();
469    let mut start = 0;
470    while start < lines.len() {
471        let ideal_end = start.saturating_add(target_line_count).min(lines.len());
472        let mut end = ideal_end;
473        if ideal_end < lines.len()
474            && let Some(boundary) =
475                empty_line_boundary_near(lines, start, ideal_end, overlap_line_count)
476        {
477            end = boundary;
478        }
479        if end <= start {
480            end = ideal_end;
481        }
482        if end <= start {
483            break;
484        }
485
486        ranges.push(start..end);
487        if end == lines.len() {
488            break;
489        }
490
491        let next_start = end.saturating_sub(overlap_line_count);
492        start = if next_start <= start { end } else { next_start };
493    }
494
495    ranges
496}
497
498fn empty_line_boundary_near(
499    lines: &[&str],
500    start: usize,
501    ideal_end: usize,
502    overlap_line_count: usize,
503) -> Option<usize> {
504    let search_start = ideal_end.saturating_sub(overlap_line_count).max(start + 1);
505    let search_end = ideal_end
506        .saturating_add(overlap_line_count)
507        .min(lines.len());
508
509    (search_start..search_end)
510        .filter(|row| lines[*row].trim().is_empty())
511        .min_by_key(|row| row.abs_diff(ideal_end))
512        .map(|row| row + 1)
513}
514
515fn lines(text: &str) -> Vec<&str> {
516    text.split_inclusive('\n').collect()
517}
518
519fn text_for_line_range(text: &str, range: Range<usize>) -> String {
520    lines(text)
521        .into_iter()
522        .skip(range.start)
523        .take(range.end.saturating_sub(range.start))
524        .collect()
525}
526
527fn tokenize(text: &str) -> Vec<String> {
528    let mut tokens = Vec::new();
529    let mut segment = String::new();
530
531    for character in text.chars() {
532        if character.is_alphanumeric() || character == '_' || character == '-' {
533            segment.push(character);
534        } else {
535            push_segment_tokens(&segment, &mut tokens);
536            segment.clear();
537        }
538    }
539    push_segment_tokens(&segment, &mut tokens);
540
541    tokens
542}
543
544fn push_segment_tokens(segment: &str, tokens: &mut Vec<String>) {
545    if segment.is_empty() {
546        return;
547    }
548
549    let mut segment_tokens = Vec::new();
550    push_token(segment, &mut segment_tokens);
551    for part in segment.split(['_', '-']).filter(|part| !part.is_empty()) {
552        push_token(part, &mut segment_tokens);
553        for camel_part in camel_case_parts(part) {
554            push_token(camel_part, &mut segment_tokens);
555        }
556    }
557
558    let mut unique_segment_tokens = Vec::new();
559    for token in segment_tokens {
560        if !unique_segment_tokens.contains(&token) {
561            unique_segment_tokens.push(token);
562        }
563    }
564    tokens.extend(unique_segment_tokens);
565}
566
567fn camel_case_parts(text: &str) -> Vec<&str> {
568    let mut parts = Vec::new();
569    let mut start = 0;
570    let mut previous = None;
571
572    for (index, character) in text.char_indices() {
573        if index > 0
574            && character.is_uppercase()
575            && previous
576                .is_some_and(|previous: char| previous.is_lowercase() || previous.is_numeric())
577        {
578            parts.push(&text[start..index]);
579            start = index;
580        }
581        previous = Some(character);
582    }
583
584    if start < text.len() {
585        parts.push(&text[start..]);
586    }
587
588    parts
589}
590
591fn push_token(token: &str, tokens: &mut Vec<String>) {
592    let token = token.to_lowercase();
593    if token.len() <= 1
594        || token.len() > 128
595        || !token.chars().any(|character| character.is_alphabetic())
596    {
597        return;
598    }
599    tokens.push(token);
600}
601
602#[cfg(test)]
603mod tests {
604    use super::*;
605
606    #[test]
607    fn test_tokenize_splits_code_identifiers() {
608        let tokens =
609            tokenize("PrivateNetworkRequestPolicy foo_bar config/reg_default_16M_retrieval.json");
610
611        assert!(tokens.contains(&"privatenetworkrequestpolicy".to_string()));
612        assert!(tokens.contains(&"private".to_string()));
613        assert!(tokens.contains(&"network".to_string()));
614        assert!(tokens.contains(&"request".to_string()));
615        assert!(tokens.contains(&"policy".to_string()));
616        assert!(tokens.contains(&"foo_bar".to_string()));
617        assert!(tokens.contains(&"foo".to_string()));
618        assert!(tokens.contains(&"bar".to_string()));
619        assert!(tokens.contains(&"reg_default_16m_retrieval".to_string()));
620        assert!(tokens.contains(&"retrieval".to_string()));
621    }
622
623    #[test]
624    fn test_chunk_line_ranges_prefers_empty_line_boundaries_with_overlap() {
625        let text = "a\nb\n\nc\nd\ne\nf\n\ng\nh\ni\nj\n";
626        let lines = lines(text);
627        let ranges = chunk_line_ranges(&lines, 3, 1);
628
629        assert_eq!(ranges[0], 0..3);
630        assert!(ranges[1].start < ranges[0].end);
631    }
632
633    #[test]
634    fn test_bm25_ranks_matching_chunk() {
635        let documents = vec![
636            Document {
637                relative_path: PathBuf::from("src/unrelated.rs"),
638                row_range: 0..1,
639                term_frequencies: {
640                    let mut terms = HashMap::new();
641                    add_term_frequencies(&mut terms, tokenize("fn unrelated"), 1);
642                    terms
643                },
644                len: 2,
645            },
646            Document {
647                relative_path: PathBuf::from("src/network.rs"),
648                row_range: 0..1,
649                term_frequencies: {
650                    let mut terms = HashMap::new();
651                    add_term_frequencies(
652                        &mut terms,
653                        tokenize("fn update_private_network_request_policy"),
654                        1,
655                    );
656                    terms
657                },
658                len: 6,
659            },
660        ];
661        let mut document_frequencies = HashMap::new();
662        for document in &documents {
663            for term in document.term_frequencies.keys() {
664                *document_frequencies.entry(term.clone()).or_default() += 1;
665            }
666        }
667        let index = Bm25Index {
668            documents,
669            document_frequencies,
670            average_document_len: 4.0,
671            stats: Bm25IndexStats::default(),
672        };
673        let mut query = HashMap::new();
674        add_query_terms(&mut query, "PrivateNetworkRequestPolicy", 1.0);
675
676        let candidates = index.search(&query, "repo", 0);
677
678        assert_eq!(candidates[0].path, Path::new("repo/src/network.rs"));
679        assert_eq!(candidates[0].row_range, 0..1);
680        assert_eq!(candidates[0].order, 0);
681    }
682}
683