updated vector_store db to leverage EMBEDDINGS_DIR path

KCaverly created 2 years ago

Change summary

crates/util/src/paths.rs                      |  1 
crates/vector_store/src/db.rs                 |  7 ---
crates/vector_store/src/vector_store.rs       | 33 ++++++++++++++------
crates/vector_store/src/vector_store_tests.rs |  2 
4 files changed, 26 insertions(+), 17 deletions(-)

Detailed changes

crates/util/src/paths.rs

@@ -6,6 +6,7 @@ lazy_static::lazy_static! {
     pub static ref HOME: PathBuf = dirs::home_dir().expect("failed to determine home directory");
     pub static ref CONFIG_DIR: PathBuf = HOME.join(".config").join("zed");
     pub static ref CONVERSATIONS_DIR: PathBuf = HOME.join(".config/zed/conversations");
+    pub static ref EMBEDDINGS_DIR: PathBuf = HOME.join(".config/zed/embeddings");
     pub static ref LOGS_DIR: PathBuf = HOME.join("Library/Logs/Zed");
     pub static ref SUPPORT_DIR: PathBuf = HOME.join("Library/Application Support/Zed");
     pub static ref LANGUAGES_DIR: PathBuf = HOME.join("Library/Application Support/Zed/languages");

crates/vector_store/src/db.rs

@@ -15,11 +15,6 @@ use sha1::{Digest, Sha1};
 
 use crate::IndexedFile;
 
-// This is saving to a local database store within the users dev zed path
-// Where do we want this to sit?
-// Assuming near where the workspace DB sits.
-pub const VECTOR_DB_URL: &str = "embeddings_db";
-
 // Note this is not an appropriate document
 #[derive(Debug)]
 pub struct DocumentRecord {
@@ -109,7 +104,7 @@ pub struct VectorDatabase {
 }
 
 impl VectorDatabase {
-    pub fn new(path: &str) -> Result<Self> {
+    pub fn new(path: String) -> Result<Self> {
         let this = Self {
             db: rusqlite::Connection::open(path)?,
         };

crates/vector_store/src/vector_store.rs

@@ -6,16 +6,23 @@ mod modal;
 mod vector_store_tests;
 
 use anyhow::{anyhow, Result};
-use db::{FileSha1, VectorDatabase, VECTOR_DB_URL};
+use db::{FileSha1, VectorDatabase};
 use embedding::{EmbeddingProvider, OpenAIEmbeddings};
 use gpui::{AppContext, Entity, ModelContext, ModelHandle, Task, ViewContext};
 use language::{Language, LanguageRegistry};
 use modal::{SemanticSearch, SemanticSearchDelegate, Toggle};
 use project::{Fs, Project, WorktreeId};
 use smol::channel;
-use std::{cmp::Ordering, collections::HashMap, path::PathBuf, sync::Arc};
+use std::{
+    cmp::Ordering,
+    collections::HashMap,
+    path::{Path, PathBuf},
+    sync::Arc,
+};
 use tree_sitter::{Parser, QueryCursor};
-use util::{http::HttpClient, ResultExt, TryFutureExt};
+use util::{
+    channel::RELEASE_CHANNEL_NAME, http::HttpClient, paths::EMBEDDINGS_DIR, ResultExt, TryFutureExt,
+};
 use workspace::{Workspace, WorkspaceCreated};
 
 #[derive(Debug)]
@@ -31,11 +38,14 @@ pub fn init(
     language_registry: Arc<LanguageRegistry>,
     cx: &mut AppContext,
 ) {
+    let db_file_path = EMBEDDINGS_DIR
+        .join(Path::new(RELEASE_CHANNEL_NAME.as_str()))
+        .join("embeddings_db");
+
     let vector_store = cx.add_model(|_| {
         VectorStore::new(
             fs,
-            VECTOR_DB_URL.to_string(),
-            // Arc::new(DummyEmbeddings {}),
+            db_file_path,
             Arc::new(OpenAIEmbeddings {
                 client: http_client,
             }),
@@ -87,7 +97,7 @@ pub struct IndexedFile {
 
 pub struct VectorStore {
     fs: Arc<dyn Fs>,
-    database_url: Arc<str>,
+    database_url: Arc<PathBuf>,
     embedding_provider: Arc<dyn EmbeddingProvider>,
     language_registry: Arc<LanguageRegistry>,
     worktree_db_ids: Vec<(WorktreeId, i64)>,
@@ -104,13 +114,13 @@ pub struct SearchResult {
 impl VectorStore {
     fn new(
         fs: Arc<dyn Fs>,
-        database_url: String,
+        database_url: PathBuf,
         embedding_provider: Arc<dyn EmbeddingProvider>,
         language_registry: Arc<LanguageRegistry>,
     ) -> Self {
         Self {
             fs,
-            database_url: database_url.into(),
+            database_url: Arc::new(database_url),
             embedding_provider,
             language_registry,
             worktree_db_ids: Vec::new(),
@@ -209,7 +219,10 @@ impl VectorStore {
                 .timer(std::time::Duration::from_secs(3))
                 .await;
 
-            let db = VectorDatabase::new(&database_url)?;
+            if let Some(db_directory) = database_url.parent() {
+                fs.create_dir(db_directory).await.log_err();
+            }
+            let db = VectorDatabase::new(database_url.to_string_lossy().into())?;
 
             let worktrees = project.read_with(&cx, |project, cx| {
                 project
@@ -372,7 +385,7 @@ impl VectorStore {
             let documents = cx
                 .background()
                 .spawn(async move {
-                    let database = VectorDatabase::new(database_url.as_ref())?;
+                    let database = VectorDatabase::new(database_url.to_string_lossy().into())?;
 
                     let phrase_embedding = embedding_provider
                         .embed_batch(vec![&phrase])

crates/vector_store/src/vector_store_tests.rs

@@ -63,7 +63,7 @@ async fn test_vector_store(cx: &mut TestAppContext) {
     let store = cx.add_model(|_| {
         VectorStore::new(
             fs.clone(),
-            db_path.to_string_lossy().to_string(),
+            db_path,
             Arc::new(FakeEmbeddingProvider),
             languages,
         )