1mod db;
2mod embedding;
3mod modal;
4
5#[cfg(test)]
6mod vector_store_tests;
7
8use anyhow::{anyhow, Result};
9use db::{FileSha1, VectorDatabase};
10use embedding::{EmbeddingProvider, OpenAIEmbeddings};
11use gpui::{AppContext, Entity, ModelContext, ModelHandle, Task, ViewContext};
12use language::{Language, LanguageRegistry};
13use modal::{SemanticSearch, SemanticSearchDelegate, Toggle};
14use project::{Fs, Project, WorktreeId};
15use smol::channel;
16use std::{
17 cmp::Ordering,
18 collections::HashMap,
19 path::{Path, PathBuf},
20 sync::Arc,
21};
22use tree_sitter::{Parser, QueryCursor};
23use util::{
24 channel::RELEASE_CHANNEL_NAME, http::HttpClient, paths::EMBEDDINGS_DIR, ResultExt, TryFutureExt,
25};
26use workspace::{Workspace, WorkspaceCreated};
27
28#[derive(Debug)]
29pub struct Document {
30 pub offset: usize,
31 pub name: String,
32 pub embedding: Vec<f32>,
33}
34
35pub fn init(
36 fs: Arc<dyn Fs>,
37 http_client: Arc<dyn HttpClient>,
38 language_registry: Arc<LanguageRegistry>,
39 cx: &mut AppContext,
40) {
41 let db_file_path = EMBEDDINGS_DIR
42 .join(Path::new(RELEASE_CHANNEL_NAME.as_str()))
43 .join("embeddings_db");
44
45 let vector_store = cx.add_model(|_| {
46 VectorStore::new(
47 fs,
48 db_file_path,
49 Arc::new(OpenAIEmbeddings {
50 client: http_client,
51 }),
52 language_registry,
53 )
54 });
55
56 cx.subscribe_global::<WorkspaceCreated, _>({
57 let vector_store = vector_store.clone();
58 move |event, cx| {
59 let workspace = &event.0;
60 if let Some(workspace) = workspace.upgrade(cx) {
61 let project = workspace.read(cx).project().clone();
62 if project.read(cx).is_local() {
63 vector_store.update(cx, |store, cx| {
64 store.add_project(project, cx).detach();
65 });
66 }
67 }
68 }
69 })
70 .detach();
71
72 cx.add_action({
73 move |workspace: &mut Workspace, _: &Toggle, cx: &mut ViewContext<Workspace>| {
74 let vector_store = vector_store.clone();
75 workspace.toggle_modal(cx, |workspace, cx| {
76 let project = workspace.project().clone();
77 let workspace = cx.weak_handle();
78 cx.add_view(|cx| {
79 SemanticSearch::new(
80 SemanticSearchDelegate::new(workspace, project, vector_store),
81 cx,
82 )
83 })
84 })
85 }
86 });
87
88 SemanticSearch::init(cx);
89}
90
91#[derive(Debug)]
92pub struct IndexedFile {
93 path: PathBuf,
94 sha1: FileSha1,
95 documents: Vec<Document>,
96}
97
98pub struct VectorStore {
99 fs: Arc<dyn Fs>,
100 database_url: Arc<PathBuf>,
101 embedding_provider: Arc<dyn EmbeddingProvider>,
102 language_registry: Arc<LanguageRegistry>,
103 worktree_db_ids: Vec<(WorktreeId, i64)>,
104}
105
106#[derive(Debug, Clone)]
107pub struct SearchResult {
108 pub worktree_id: WorktreeId,
109 pub name: String,
110 pub offset: usize,
111 pub file_path: PathBuf,
112}
113
114impl VectorStore {
115 fn new(
116 fs: Arc<dyn Fs>,
117 database_url: PathBuf,
118 embedding_provider: Arc<dyn EmbeddingProvider>,
119 language_registry: Arc<LanguageRegistry>,
120 ) -> Self {
121 Self {
122 fs,
123 database_url: Arc::new(database_url),
124 embedding_provider,
125 language_registry,
126 worktree_db_ids: Vec::new(),
127 }
128 }
129
130 async fn index_file(
131 cursor: &mut QueryCursor,
132 parser: &mut Parser,
133 embedding_provider: &dyn EmbeddingProvider,
134 language: Arc<Language>,
135 file_path: PathBuf,
136 content: String,
137 ) -> Result<IndexedFile> {
138 let grammar = language.grammar().ok_or_else(|| anyhow!("no grammar"))?;
139 let outline_config = grammar
140 .outline_config
141 .as_ref()
142 .ok_or_else(|| anyhow!("no outline query"))?;
143
144 parser.set_language(grammar.ts_language).unwrap();
145 let tree = parser
146 .parse(&content, None)
147 .ok_or_else(|| anyhow!("parsing failed"))?;
148
149 let mut documents = Vec::new();
150 let mut context_spans = Vec::new();
151 for mat in cursor.matches(&outline_config.query, tree.root_node(), content.as_bytes()) {
152 let mut item_range = None;
153 let mut name_range = None;
154 for capture in mat.captures {
155 if capture.index == outline_config.item_capture_ix {
156 item_range = Some(capture.node.byte_range());
157 } else if capture.index == outline_config.name_capture_ix {
158 name_range = Some(capture.node.byte_range());
159 }
160 }
161
162 if let Some((item_range, name_range)) = item_range.zip(name_range) {
163 if let Some((item, name)) =
164 content.get(item_range.clone()).zip(content.get(name_range))
165 {
166 context_spans.push(item);
167 documents.push(Document {
168 name: name.to_string(),
169 offset: item_range.start,
170 embedding: Vec::new(),
171 });
172 }
173 }
174 }
175
176 if !documents.is_empty() {
177 let embeddings = embedding_provider.embed_batch(context_spans).await?;
178 for (document, embedding) in documents.iter_mut().zip(embeddings) {
179 document.embedding = embedding;
180 }
181 }
182
183 let sha1 = FileSha1::from_str(content);
184
185 return Ok(IndexedFile {
186 path: file_path,
187 sha1,
188 documents,
189 });
190 }
191
192 fn add_project(
193 &mut self,
194 project: ModelHandle<Project>,
195 cx: &mut ModelContext<Self>,
196 ) -> Task<Result<()>> {
197 let worktree_scans_complete = project
198 .read(cx)
199 .worktrees(cx)
200 .map(|worktree| {
201 let scan_complete = worktree.read(cx).as_local().unwrap().scan_complete();
202 async move {
203 scan_complete.await;
204 log::info!("worktree scan completed");
205 }
206 })
207 .collect::<Vec<_>>();
208
209 let fs = self.fs.clone();
210 let language_registry = self.language_registry.clone();
211 let embedding_provider = self.embedding_provider.clone();
212 let database_url = self.database_url.clone();
213
214 cx.spawn(|this, mut cx| async move {
215 futures::future::join_all(worktree_scans_complete).await;
216
217 if let Some(db_directory) = database_url.parent() {
218 fs.create_dir(db_directory).await.log_err();
219 }
220 let db = VectorDatabase::new(database_url.to_string_lossy().into())?;
221
222 let worktrees = project.read_with(&cx, |project, cx| {
223 project
224 .worktrees(cx)
225 .map(|worktree| worktree.read(cx).snapshot())
226 .collect::<Vec<_>>()
227 });
228
229 // Here we query the worktree ids, and yet we dont have them elsewhere
230 // We likely want to clean up these datastructures
231 let (db, worktree_hashes, worktree_db_ids) = cx
232 .background()
233 .spawn({
234 let worktrees = worktrees.clone();
235 async move {
236 let mut worktree_db_ids: HashMap<WorktreeId, i64> = HashMap::new();
237 let mut hashes: HashMap<WorktreeId, HashMap<PathBuf, FileSha1>> =
238 HashMap::new();
239 for worktree in worktrees {
240 let worktree_db_id =
241 db.find_or_create_worktree(worktree.abs_path().as_ref())?;
242 worktree_db_ids.insert(worktree.id(), worktree_db_id);
243 hashes.insert(worktree.id(), db.get_file_hashes(worktree_db_id)?);
244 }
245 anyhow::Ok((db, hashes, worktree_db_ids))
246 }
247 })
248 .await?;
249
250 let (paths_tx, paths_rx) =
251 channel::unbounded::<(i64, PathBuf, String, Arc<Language>)>();
252 let (indexed_files_tx, indexed_files_rx) = channel::unbounded::<(i64, IndexedFile)>();
253 cx.background()
254 .spawn({
255 let fs = fs.clone();
256 let worktree_db_ids = worktree_db_ids.clone();
257 async move {
258 for worktree in worktrees.into_iter() {
259 let file_hashes = &worktree_hashes[&worktree.id()];
260 for file in worktree.files(false, 0) {
261 let absolute_path = worktree.absolutize(&file.path);
262
263 if let Ok(language) = language_registry
264 .language_for_file(&absolute_path, None)
265 .await
266 {
267 if language.name().as_ref() != "Rust" {
268 continue;
269 }
270
271 if let Some(content) = fs.load(&absolute_path).await.log_err() {
272 log::info!("loaded file: {absolute_path:?}");
273
274 let path_buf = file.path.to_path_buf();
275 let already_stored = file_hashes
276 .get(&path_buf)
277 .map_or(false, |existing_hash| {
278 existing_hash.equals(&content)
279 });
280
281 if !already_stored {
282 log::info!(
283 "File Changed (Sending to Parse): {:?}",
284 &path_buf
285 );
286 paths_tx
287 .try_send((
288 worktree_db_ids[&worktree.id()],
289 path_buf,
290 content,
291 language,
292 ))
293 .unwrap();
294 }
295 }
296 }
297 }
298 }
299 }
300 })
301 .detach();
302
303 let db_write_task = cx.background().spawn(
304 async move {
305 while let Ok((worktree_id, indexed_file)) = indexed_files_rx.recv().await {
306 db.insert_file(worktree_id, indexed_file).log_err();
307 }
308
309 anyhow::Ok(())
310 }
311 .log_err(),
312 );
313
314 cx.background()
315 .scoped(|scope| {
316 for _ in 0..cx.background().num_cpus() {
317 scope.spawn(async {
318 let mut parser = Parser::new();
319 let mut cursor = QueryCursor::new();
320 while let Ok((worktree_id, file_path, content, language)) =
321 paths_rx.recv().await
322 {
323 if let Some(indexed_file) = Self::index_file(
324 &mut cursor,
325 &mut parser,
326 embedding_provider.as_ref(),
327 language,
328 file_path,
329 content,
330 )
331 .await
332 .log_err()
333 {
334 indexed_files_tx
335 .try_send((worktree_id, indexed_file))
336 .unwrap();
337 }
338 }
339 });
340 }
341 })
342 .await;
343 drop(indexed_files_tx);
344
345 db_write_task.await;
346
347 this.update(&mut cx, |this, _| {
348 this.worktree_db_ids.extend(worktree_db_ids);
349 });
350
351 anyhow::Ok(())
352 })
353 }
354
355 pub fn search(
356 &mut self,
357 project: &ModelHandle<Project>,
358 phrase: String,
359 limit: usize,
360 cx: &mut ModelContext<Self>,
361 ) -> Task<Result<Vec<SearchResult>>> {
362 let project = project.read(cx);
363 let worktree_db_ids = project
364 .worktrees(cx)
365 .filter_map(|worktree| {
366 let worktree_id = worktree.read(cx).id();
367 self.worktree_db_ids.iter().find_map(|(id, db_id)| {
368 if *id == worktree_id {
369 Some(*db_id)
370 } else {
371 None
372 }
373 })
374 })
375 .collect::<Vec<_>>();
376
377 let embedding_provider = self.embedding_provider.clone();
378 let database_url = self.database_url.clone();
379 cx.spawn(|this, cx| async move {
380 let documents = cx
381 .background()
382 .spawn(async move {
383 let database = VectorDatabase::new(database_url.to_string_lossy().into())?;
384
385 let phrase_embedding = embedding_provider
386 .embed_batch(vec![&phrase])
387 .await?
388 .into_iter()
389 .next()
390 .unwrap();
391
392 let mut results = Vec::<(i64, f32)>::with_capacity(limit + 1);
393 database.for_each_document(&worktree_db_ids, |id, embedding| {
394 let similarity = dot(&embedding.0, &phrase_embedding);
395 let ix = match results.binary_search_by(|(_, s)| {
396 similarity.partial_cmp(&s).unwrap_or(Ordering::Equal)
397 }) {
398 Ok(ix) => ix,
399 Err(ix) => ix,
400 };
401 results.insert(ix, (id, similarity));
402 results.truncate(limit);
403 })?;
404
405 let ids = results.into_iter().map(|(id, _)| id).collect::<Vec<_>>();
406 database.get_documents_by_ids(&ids)
407 })
408 .await?;
409
410 let results = this.read_with(&cx, |this, _| {
411 documents
412 .into_iter()
413 .filter_map(|(worktree_db_id, file_path, offset, name)| {
414 let worktree_id = this.worktree_db_ids.iter().find_map(|(id, db_id)| {
415 if *db_id == worktree_db_id {
416 Some(*id)
417 } else {
418 None
419 }
420 })?;
421 Some(SearchResult {
422 worktree_id,
423 name,
424 offset,
425 file_path,
426 })
427 })
428 .collect()
429 });
430
431 anyhow::Ok(results)
432 })
433 }
434}
435
436impl Entity for VectorStore {
437 type Event = ();
438}
439
440fn dot(vec_a: &[f32], vec_b: &[f32]) -> f32 {
441 let len = vec_a.len();
442 assert_eq!(len, vec_b.len());
443
444 let mut result = 0.0;
445 unsafe {
446 matrixmultiply::sgemm(
447 1,
448 len,
449 1,
450 1.0,
451 vec_a.as_ptr(),
452 len as isize,
453 1,
454 vec_b.as_ptr(),
455 1,
456 len as isize,
457 0.0,
458 &mut result as *mut f32,
459 1,
460 1,
461 );
462 }
463 result
464}