vecdb is reloading only if caps changes. It checks for it every 30s

smallcloudai · Jan 11, 2024 · 277b8f4 · 277b8f4
1 parent 4365813
commit 277b8f4
Showing 1 changed file with 18 additions and 28 deletions.
diff --git a/src/vecdb/vecdb.rs b/src/vecdb/vecdb.rs
@@ -30,10 +30,6 @@ pub struct VecDb {
 }
 
 
-const VECDB_BACKGROUND_RELOAD_ON_SUCCESS: u64 = 1200;
-const VECDB_BACKGROUND_RELOAD_ON_FAIL: u64 = 30;
-
-
 pub async fn create_vecdb(
     default_embeddings_model: String,
     endpoint_embeddings_template: String,
@@ -69,38 +65,19 @@ pub async fn vecdb_background_reload(
     global_context: Arc<ARwLock<GlobalContext>>,
 ) {
     let mut background_tasks = BackgroundTasksHolder::new(vec![]);
-
-    let mut vecdb_launched = false;
-    let mut first_loop = true;
     loop {
-        if !first_loop {
-            if vecdb_launched {
-                tokio::time::sleep(tokio::time::Duration::from_secs(VECDB_BACKGROUND_RELOAD_ON_SUCCESS)).await;
-            } else {
-                tokio::time::sleep(tokio::time::Duration::from_secs(VECDB_BACKGROUND_RELOAD_ON_FAIL)).await;
-            }
-        }
-
-        background_tasks.abort().await;
-        background_tasks = BackgroundTasksHolder::new(vec![]);
-
-        vecdb_launched = false;
-        first_loop = false;
-        info!("attempting to launch vecdb");
+        tokio::time::sleep(tokio::time::Duration::from_secs(30)).await;
 
         let mut gcx_locked = global_context.write().await;
         let caps_mb = gcx_locked.caps.clone();
         if caps_mb.is_none() {
-            info!("vecd launch failed: no caps");
             continue;
         }
 
         let cache_dir = &gcx_locked.cache_dir;
         let cmdline = &gcx_locked.cmdline;
 
         if !cmdline.vecdb {
-            info!("VecDB launch is disabled by cmdline");
-            vecdb_launched = true;
             continue;
         }
 
@@ -121,10 +98,26 @@ pub async fn vecdb_background_reload(
         };
 
         if default_embeddings_model.is_empty() || endpoint_embeddings_template.is_empty() {
-            info!("vecd launch failed: default_embeddings_model.is_empty() || endpoint_embeddings_template.is_empty()");
+            error!("vecd launch failed: default_embeddings_model.is_empty() || endpoint_embeddings_template.is_empty()");
             continue;
         }
 
+        match *gcx_locked.vec_db.lock().await {
+            None => {}
+            Some(ref db) => {
+                if db.model_name == default_embeddings_model &&
+                    db.endpoint_template == endpoint_embeddings_template &&
+                    db.endpoint_embeddings_style == endpoint_embeddings_style {
+                    continue;
+                }
+            }
+        }
+
+        info!("attempting to launch vecdb");
+
+        background_tasks.abort().await;
+        background_tasks = BackgroundTasksHolder::new(vec![]);
+
         let vecdb_mb = create_vecdb(
             default_embeddings_model.clone(),
             endpoint_embeddings_template,
@@ -136,12 +129,10 @@ pub async fn vecdb_background_reload(
         ).await;
 
         if vecdb_mb.is_none() {
-            info!("vecd launch failed: vecdb_mb.is_none()");
             continue;
         }
         gcx_locked.vec_db = Arc::new(AMutex::new(vecdb_mb));
         info!("VECDB is launched successfully");
-        vecdb_launched = true;
 
         background_tasks.extend(match *gcx_locked.vec_db.lock().await {
             Some(ref db) => db.start_background_tasks().await,
@@ -156,7 +147,6 @@ pub async fn vecdb_background_reload(
                 }
             }
         }
-
     }
 }