Merge pull request #1286 from dcSpark/feature/change-embedding-model

Allows to change the embedding model. Recreate embeddings if model ch…

Author: guillevalin

Cargo.lock

@@ -19,7 +19,7 @@ members = [
 resolver = "2"
 
 [workspace.package]
-version = "1.1.9"
+version = "1.1.10"
 edition = "2021"
 authors = ["Nico Arqueros <nico@shinkai.com>"]
 

Cargo.toml

@@ -535,9 +535,7 @@ mod tests {
         let temp_file = NamedTempFile::new().unwrap();
         let db_path = std::path::PathBuf::from(temp_file.path());
         let api_url = String::new();
-        let model_type = EmbeddingModelType::OllamaTextEmbeddingsInference(
-            OllamaTextEmbeddingsInference::SnowflakeArcticEmbedM
-        );
+        let model_type = EmbeddingModelType::default();
 
         SqliteManager::new(db_path, api_url, model_type).unwrap()
     }

shinkai-bin/shinkai-node/src/llm_provider/providers/gemini.rs

@@ -1002,7 +1002,7 @@ mod tests {
     //     let fs_db_path = format!("db_tests/{}", "vector_fs");
     //     let profile_list = vec![default_test_profile()];
     //     let supported_embedding_models = vec![EmbeddingModelType::OllamaTextEmbeddingsInference(
-    //         OllamaTextEmbeddingsInference::SnowflakeArcticEmbed_M,
+    //         OllamaTextEmbeddingsInference::EmbeddingGemma300M,
     //     )];
 
     //     VectorFS::new(

shinkai-bin/shinkai-node/src/network/agent_payments_manager/my_agent_offerings_manager.rs

@@ -443,7 +443,25 @@ impl Node {
             callback_manager.update_cron_manager(cron_manager.clone());
         }
 
+        // Perform embedding migration if needed (BEFORE updating DB with current model)
+        {
+            let current_model = {
+                let default_model_guard = self.default_embedding_model.lock().await;
+                default_model_guard.clone()
+            };
+            
+            if let Err(e) = self.db.migrate_embeddings_to_new_model(&self.embedding_generator, &current_model).await {
+                shinkai_log(
+                    ShinkaiLogOption::Node,
+                    ShinkaiLogLevel::Error,
+                    &format!("Embedding migration failed: {e:?}"),
+                );
+                // Note: We continue even if migration fails to allow the node to start
+            }
+        }
+        
         self.initialize_embedding_models().await?;
+        
         {
             // Starting the WebSocket server
             if let (Some(ws_manager), Some(ws_address)) = (&self.ws_manager, self.ws_address) {

shinkai-bin/shinkai-node/src/network/node.rs

@@ -1376,7 +1376,7 @@ mod tests {
         let db_path = PathBuf::from(temp_file.path());
         let api_url = String::new();
         let model_type =
-            EmbeddingModelType::OllamaTextEmbeddingsInference(OllamaTextEmbeddingsInference::SnowflakeArcticEmbedM);
+            EmbeddingModelType::default();
         println!("Creating test db at {:?}", db_path);
         SqliteManager::new(db_path, api_url, model_type).unwrap()
     }

shinkai-bin/shinkai-node/src/network/zip_export_import/zip_export_import.rs

@@ -322,7 +322,7 @@ mod tests {
         let db_path = PathBuf::from(temp_file.path());
         let api_url = String::new();
         let model_type =
-            EmbeddingModelType::OllamaTextEmbeddingsInference(OllamaTextEmbeddingsInference::SnowflakeArcticEmbedM);
+            EmbeddingModelType::default();
 
         SqliteManager::new(db_path, api_url, model_type).unwrap()
     }
@@ -388,7 +388,7 @@ mod tests {
             },
             true,
         );
-        initial_tool.set_embedding(vec![0.0; 384]);
+        initial_tool.set_embedding(vec![0.0; EmbeddingModelType::default().vector_dimensions().unwrap_or(768)]);
         initial_tool
     }
 

shinkai-bin/shinkai-node/src/tools/tool_implementation/native_tools/config_setup.rs

@@ -2,7 +2,7 @@ use std::env;
 use std::net::{IpAddr, SocketAddr};
 use std::str::FromStr;
 
-use shinkai_embedding::model_type::{EmbeddingModelType, OllamaTextEmbeddingsInference};
+use shinkai_embedding::model_type::EmbeddingModelType;
 use shinkai_message_primitives::schemas::llm_providers::serialized_llm_provider::{
     LLMProviderInterface, SerializedLLMProvider
 };
@@ -169,11 +169,7 @@ pub fn fetch_node_environment() -> NodeEnvironment {
     }
 
     // Fetch the default embedding model
-    let default_embedding_model: EmbeddingModelType = env::var("DEFAULT_EMBEDDING_MODEL")
-        .map(|s| EmbeddingModelType::from_string(&s).expect("Failed to parse DEFAULT_EMBEDDING_MODEL"))
-        .unwrap_or_else(|_| {
-            EmbeddingModelType::OllamaTextEmbeddingsInference(OllamaTextEmbeddingsInference::SnowflakeArcticEmbedM)
-        });
+    let default_embedding_model: EmbeddingModelType = EmbeddingModelType::default();
 
     // Fetch the supported embedding models
     let supported_embedding_models: Vec<EmbeddingModelType> = env::var("SUPPORTED_EMBEDDING_MODELS")
@@ -183,9 +179,7 @@ pub fn fetch_node_environment() -> NodeEnvironment {
                 .collect()
         })
         .unwrap_or_else(|_| {
-            vec![EmbeddingModelType::OllamaTextEmbeddingsInference(
-                OllamaTextEmbeddingsInference::SnowflakeArcticEmbedM,
-            )]
+            vec![EmbeddingModelType::default()]
         });
 
     // Fetch the API_V2_KEY environment variable

shinkai-bin/shinkai-node/src/utils/environment.rs

@@ -23,7 +23,7 @@ fn setup_test_db() -> SqliteManager {
     let db_path = PathBuf::from(temp_file.path());
     let api_url = String::new();
     let model_type =
-        EmbeddingModelType::OllamaTextEmbeddingsInference(OllamaTextEmbeddingsInference::SnowflakeArcticEmbedM);
+        EmbeddingModelType::default();
 
     SqliteManager::new(db_path, api_url, model_type).unwrap()
 }

shinkai-bin/shinkai-node/tests/it/db_identity_tests.rs

@@ -23,7 +23,7 @@ fn setup_test_db() -> SqliteManager {
     let db_path = PathBuf::from(temp_file.path());
     let api_url = String::new();
     let model_type =
-        EmbeddingModelType::OllamaTextEmbeddingsInference(OllamaTextEmbeddingsInference::SnowflakeArcticEmbedM);
+        EmbeddingModelType::default();
 
     SqliteManager::new(db_path, api_url, model_type).unwrap()
 }