1use crate::{parsing::Span, JobHandle};
2use ai::embedding::EmbeddingProvider;
3use gpui::executor::Background;
4use parking_lot::Mutex;
5use smol::channel;
6use std::{mem, ops::Range, path::Path, sync::Arc, time::SystemTime};
7
8#[derive(Clone)]
9pub struct FileToEmbed {
10 pub worktree_id: i64,
11 pub path: Arc<Path>,
12 pub mtime: SystemTime,
13 pub spans: Vec<Span>,
14 pub job_handle: JobHandle,
15}
16
17impl std::fmt::Debug for FileToEmbed {
18 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
19 f.debug_struct("FileToEmbed")
20 .field("worktree_id", &self.worktree_id)
21 .field("path", &self.path)
22 .field("mtime", &self.mtime)
23 .field("spans", &self.spans)
24 .finish_non_exhaustive()
25 }
26}
27
28impl PartialEq for FileToEmbed {
29 fn eq(&self, other: &Self) -> bool {
30 self.worktree_id == other.worktree_id
31 && self.path == other.path
32 && self.mtime == other.mtime
33 && self.spans == other.spans
34 }
35}
36
37pub struct EmbeddingQueue {
38 embedding_provider: Arc<dyn EmbeddingProvider>,
39 pending_batch: Vec<FileFragmentToEmbed>,
40 executor: Arc<Background>,
41 pending_batch_token_count: usize,
42 finished_files_tx: channel::Sender<FileToEmbed>,
43 finished_files_rx: channel::Receiver<FileToEmbed>,
44}
45
46#[derive(Clone)]
47pub struct FileFragmentToEmbed {
48 file: Arc<Mutex<FileToEmbed>>,
49 span_range: Range<usize>,
50}
51
52impl EmbeddingQueue {
53 pub fn new(embedding_provider: Arc<dyn EmbeddingProvider>, executor: Arc<Background>) -> Self {
54 let (finished_files_tx, finished_files_rx) = channel::unbounded();
55 Self {
56 embedding_provider,
57 executor,
58 pending_batch: Vec::new(),
59 pending_batch_token_count: 0,
60 finished_files_tx,
61 finished_files_rx,
62 }
63 }
64
65 pub fn push(&mut self, file: FileToEmbed) {
66 if file.spans.is_empty() {
67 self.finished_files_tx.try_send(file).unwrap();
68 return;
69 }
70
71 let file = Arc::new(Mutex::new(file));
72
73 self.pending_batch.push(FileFragmentToEmbed {
74 file: file.clone(),
75 span_range: 0..0,
76 });
77
78 let mut fragment_range = &mut self.pending_batch.last_mut().unwrap().span_range;
79 for (ix, span) in file.lock().spans.iter().enumerate() {
80 let span_token_count = if span.embedding.is_none() {
81 span.token_count
82 } else {
83 0
84 };
85
86 let next_token_count = self.pending_batch_token_count + span_token_count;
87 if next_token_count > self.embedding_provider.max_tokens_per_batch() {
88 let range_end = fragment_range.end;
89 self.flush();
90 self.pending_batch.push(FileFragmentToEmbed {
91 file: file.clone(),
92 span_range: range_end..range_end,
93 });
94 fragment_range = &mut self.pending_batch.last_mut().unwrap().span_range;
95 }
96
97 fragment_range.end = ix + 1;
98 self.pending_batch_token_count += span_token_count;
99 }
100 }
101
102 pub fn flush(&mut self) {
103 let batch = mem::take(&mut self.pending_batch);
104 self.pending_batch_token_count = 0;
105 if batch.is_empty() {
106 return;
107 }
108
109 let finished_files_tx = self.finished_files_tx.clone();
110 let embedding_provider = self.embedding_provider.clone();
111
112 self.executor
113 .spawn(async move {
114 let mut spans = Vec::new();
115 for fragment in &batch {
116 let file = fragment.file.lock();
117 spans.extend(
118 file.spans[fragment.span_range.clone()]
119 .iter()
120 .filter(|d| d.embedding.is_none())
121 .map(|d| d.content.clone()),
122 );
123 }
124
125 // If spans is 0, just send the fragment to the finished files if its the last one.
126 if spans.is_empty() {
127 for fragment in batch.clone() {
128 if let Some(file) = Arc::into_inner(fragment.file) {
129 finished_files_tx.try_send(file.into_inner()).unwrap();
130 }
131 }
132 return;
133 };
134
135 match embedding_provider.embed_batch(spans).await {
136 Ok(embeddings) => {
137 let mut embeddings = embeddings.into_iter();
138 for fragment in batch {
139 for span in &mut fragment.file.lock().spans[fragment.span_range.clone()]
140 .iter_mut()
141 .filter(|d| d.embedding.is_none())
142 {
143 if let Some(embedding) = embeddings.next() {
144 span.embedding = Some(embedding);
145 } else {
146 log::error!("number of embeddings != number of documents");
147 }
148 }
149
150 if let Some(file) = Arc::into_inner(fragment.file) {
151 finished_files_tx.try_send(file.into_inner()).unwrap();
152 }
153 }
154 }
155 Err(error) => {
156 log::error!("{:?}", error);
157 }
158 }
159 })
160 .detach();
161 }
162
163 pub fn finished_files(&self) -> channel::Receiver<FileToEmbed> {
164 self.finished_files_rx.clone()
165 }
166}