From 53f11178ea0687b3593970b9615bbb1f527ac413 Mon Sep 17 00:00:00 2001 From: Nico <1622112+nicarq@users.noreply.github.com> Date: Tue, 27 Aug 2024 21:36:31 -0500 Subject: [PATCH] fix: doc chat & api key v2 (#532) * fix * update api_v2_key * make ollama use embedding batches * api v2 always active * Revert "make ollama use embedding batches" This reverts commit c51e83a422130e5b8dd7cb91e413c6994e9bfd58. * bump: shinkai node v0.7.31 * fix test --------- Co-authored-by: paulclindo --- Cargo.lock | 2 +- shinkai-bin/shinkai-node/Cargo.toml | 2 +- shinkai-bin/shinkai-node/src/db/db_inbox.rs | 2 +- shinkai-bin/shinkai-node/src/db/db_main.rs | 16 +++++ .../chains/generic_chain/generic_prompts.rs | 22 +++---- .../llm_provider/execution/prompts/prompts.rs | 61 ++++++++----------- .../execution/prompts/subprompts.rs | 2 + shinkai-bin/shinkai-node/src/network/node.rs | 39 +++++++++++- .../src/network/node_api_router.rs | 25 +++----- .../src/network/v1_api/api_v1_commands.rs | 29 +++++++++ .../src/network/v1_api/api_v1_handlers.rs | 11 ++-- .../src/network/v2_api/api_v2_commands.rs | 40 ++++++++---- shinkai-bin/shinkai-node/src/runner.rs | 1 + .../shinkai-node/src/utils/environment.rs | 7 ++- .../tests/it/a0_subscription_manager_tests.rs | 2 + .../tests/it/a1_http_subscription_tests.rs | 2 + .../tests/it/change_nodes_name_tests.rs | 2 + .../shinkai-node/tests/it/db_restore_tests.rs | 1 + .../it/llm_provider_integration_tests.rs | 1 + .../tests/it/node_integration_tests.rs | 2 + .../tests/it/node_retrying_tests.rs | 2 + .../shinkai-node/tests/it/tcp_proxy_tests.rs | 4 ++ .../tests/it/utils/test_boilerplate.rs | 3 + .../tests/it/workflow_integration_tests.rs | 2 + .../tests/it/z_shinkai_mirror_tests.rs | 1 + 25 files changed, 194 insertions(+), 87 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c968dc594..f1de4434e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -9538,7 +9538,7 @@ dependencies = [ [[package]] name = "shinkai_node" -version = "0.7.30" +version = "0.7.31" dependencies = [ "aes-gcm", "anyhow", diff --git a/shinkai-bin/shinkai-node/Cargo.toml b/shinkai-bin/shinkai-node/Cargo.toml index b2df4d68f..ae0ca4020 100644 --- a/shinkai-bin/shinkai-node/Cargo.toml +++ b/shinkai-bin/shinkai-node/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "shinkai_node" -version = "0.7.30" +version = "0.7.31" edition = "2021" authors.workspace = true # this causes `cargo run` in the workspace root to run this package diff --git a/shinkai-bin/shinkai-node/src/db/db_inbox.rs b/shinkai-bin/shinkai-node/src/db/db_inbox.rs index 9bdbc12b0..63d0c9638 100644 --- a/shinkai-bin/shinkai-node/src/db/db_inbox.rs +++ b/shinkai-bin/shinkai-node/src/db/db_inbox.rs @@ -475,7 +475,7 @@ impl ShinkaiDB { } shinkai_log( ShinkaiLogOption::Api, - ShinkaiLogLevel::Info, + ShinkaiLogLevel::Debug, &format!("Inboxes: {}", inboxes.join(", ")), ); Ok(inboxes) diff --git a/shinkai-bin/shinkai-node/src/db/db_main.rs b/shinkai-bin/shinkai-node/src/db/db_main.rs index 3f9d8a6fb..4d7afff98 100644 --- a/shinkai-bin/shinkai-node/src/db/db_main.rs +++ b/shinkai-bin/shinkai-node/src/db/db_main.rs @@ -185,6 +185,22 @@ impl ShinkaiDB { .ok_or(ShinkaiDBError::ShinkaiNameLacksProfile) } + /// Required for intra-communications between node UI and node + pub fn read_api_v2_key(&self) -> Result, ShinkaiDBError> { + let cf = self.get_cf_handle(Topic::NodeAndUsers)?; + match self.db.get_cf(cf, b"api_v2_key") { + Ok(Some(value)) => Ok(Some(String::from_utf8(value).map_err(|_| ShinkaiDBError::InvalidData)?)), + Ok(None) => Ok(None), + Err(_) => Err(ShinkaiDBError::FailedFetchingValue), + } + } + + /// Sets the api_v2_key value + pub fn set_api_v2_key(&self, key: &str) -> Result<(), Error> { + let cf = self.get_cf_handle(Topic::NodeAndUsers).unwrap(); + self.db.put_cf(cf, b"api_v2_key", key.as_bytes()) + } + /// Returns the first half of the blake3 hash of the folder name value pub fn user_profile_to_half_hash(profile: ShinkaiName) -> String { let full_hash = blake3::hash(profile.full_name.as_bytes()).to_hex().to_string(); diff --git a/shinkai-bin/shinkai-node/src/llm_provider/execution/chains/generic_chain/generic_prompts.rs b/shinkai-bin/shinkai-node/src/llm_provider/execution/chains/generic_chain/generic_prompts.rs index 8c7f59ccc..94a498a2f 100644 --- a/shinkai-bin/shinkai-node/src/llm_provider/execution/chains/generic_chain/generic_prompts.rs +++ b/shinkai-bin/shinkai-node/src/llm_provider/execution/chains/generic_chain/generic_prompts.rs @@ -57,17 +57,6 @@ impl JobPromptGenerator { } } - // Add the user question and the preference prompt for the answer - if !user_message.is_empty() { - let user_prompt = custom_user_prompt.unwrap_or_default(); - let content = if user_prompt.is_empty() { - user_message.clone() - } else { - format!("{}\n {}", user_message, user_prompt) - }; - prompt.add_content(content, SubPromptType::User, 100); - } - // Parses the retrieved nodes as individual sub-prompts, to support priority pruning // and also grouping i.e. instead of having 100 tiny messages, we have a message with the chunks grouped { @@ -82,6 +71,17 @@ impl JobPromptGenerator { } } + // Add the user question and the preference prompt for the answer + if !user_message.is_empty() { + let user_prompt = custom_user_prompt.unwrap_or_default(); + let content = if user_prompt.is_empty() { + user_message.clone() + } else { + format!("{}\n {}", user_message, user_prompt) + }; + prompt.add_content(content, SubPromptType::UserLastMessage, 100); + } + // If function_call exists, it means that the LLM requested a function call and we need to send the response back if let Some(function_call) = function_call { // We add the assistant request to the prompt diff --git a/shinkai-bin/shinkai-node/src/llm_provider/execution/prompts/prompts.rs b/shinkai-bin/shinkai-node/src/llm_provider/execution/prompts/prompts.rs index 6fb28f0bd..9240dcd81 100644 --- a/shinkai-bin/shinkai-node/src/llm_provider/execution/prompts/prompts.rs +++ b/shinkai-bin/shinkai-node/src/llm_provider/execution/prompts/prompts.rs @@ -286,6 +286,7 @@ impl Prompt { // Accumulator for ExtraContext content let mut extra_context_content = String::new(); let mut processing_extra_context = false; + let mut last_user_message: Option = None; for sub_prompt in &self.sub_prompts { match sub_prompt { @@ -340,25 +341,10 @@ impl Prompt { current_length += sub_prompt.count_tokens_with_pregenerated_completion_message(&new_message); tiktoken_messages.push(new_message); } + SubPrompt::Content(SubPromptType::UserLastMessage, content, _) => { + last_user_message = Some(content.clone()); + } _ => { - // If we were processing ExtraContext, add it as a single System message - if processing_extra_context { - let extra_context_message = LlmMessage { - role: Some(SubPromptType::User.to_string()), - content: Some(extra_context_content.trim().to_string()), - name: None, - function_call: None, - functions: None, - }; - current_length += - ModelCapabilitiesManager::num_tokens_from_llama3(&[extra_context_message.clone()]); - tiktoken_messages.push(extra_context_message); - - // Reset the accumulator - extra_context_content.clear(); - processing_extra_context = false; - } - // Process the current sub-prompt let new_message = sub_prompt.into_chat_completion_request_message(); current_length += sub_prompt.count_tokens_with_pregenerated_completion_message(&new_message); @@ -367,17 +353,25 @@ impl Prompt { } } - // If there are any remaining ExtraContext sub-prompts, add them as a single message - if processing_extra_context && !extra_context_content.is_empty() { - let extra_context_message = LlmMessage { + // Combine ExtraContext and UserLastMessage into one message + if !extra_context_content.is_empty() || last_user_message.is_some() { + let combined_content = format!( + "{}\n{}", + extra_context_content.trim(), + last_user_message.unwrap_or_default() + ) + .trim() + .to_string(); + + let combined_message = LlmMessage { role: Some(SubPromptType::User.to_string()), - content: Some(extra_context_content.trim().to_string()), + content: Some(combined_content), name: None, function_call: None, functions: None, }; - current_length += ModelCapabilitiesManager::num_tokens_from_llama3(&[extra_context_message.clone()]); - tiktoken_messages.push(extra_context_message); + current_length += ModelCapabilitiesManager::num_tokens_from_llama3(&[combined_message.clone()]); + tiktoken_messages.push(combined_message); } (tiktoken_messages, current_length) @@ -525,11 +519,6 @@ mod tests { let (messages, _token_length) = prompt.generate_chat_completion_messages(); - // match serde_json::to_string_pretty(&messages) { - // Ok(pretty_json) => eprintln!("messages JSON: {}", pretty_json), - // Err(e) => eprintln!("Failed to serialize tools_json: {:?}", e), - // }; - // Expected messages let expected_messages = vec![ LlmMessage { @@ -553,13 +542,6 @@ mod tests { function_call: None, functions: None, }, - LlmMessage { - role: Some("user".to_string()), - content: Some("Here is a list of relevant new content provided for you to potentially use while answering:\n- FAQ Shinkai Overview What’s Shinkai? (Summary) (Source: Shinkai - Ask Me Anything.docx, Section: ) 2024-05-05T00:33:00\n- Shinkai is a comprehensive super app designed to enhance how users interact with AI. It allows users to run AI locally, facilitating direct conversations with documents and managing files converted into AI embeddings for advanced semantic searches across user data. This local execution ensures privacy and efficiency, putting control directly in the user's hands. (Source: Shinkai - Ask Me Anything.docx, Section: 2) 2024-05-05T00:33:00".to_string()), - name: None, - function_call: None, - functions: None, - }, LlmMessage { role: Some("user".to_string()), content: Some("tell me more about Shinkai. Answer the question using this markdown and the extra context provided: \n # Answer \n here goes the answer\n".to_string()), @@ -602,6 +584,13 @@ mod tests { }, }]), }, + LlmMessage { + role: Some("user".to_string()), + content: Some("Here is a list of relevant new content provided for you to potentially use while answering:\n- FAQ Shinkai Overview What’s Shinkai? (Summary) (Source: Shinkai - Ask Me Anything.docx, Section: ) 2024-05-05T00:33:00\n- Shinkai is a comprehensive super app designed to enhance how users interact with AI. It allows users to run AI locally, facilitating direct conversations with documents and managing files converted into AI embeddings for advanced semantic searches across user data. This local execution ensures privacy and efficiency, putting control directly in the user's hands. (Source: Shinkai - Ask Me Anything.docx, Section: 2) 2024-05-05T00:33:00".to_string()), + name: None, + function_call: None, + functions: None, + }, ]; // Check if the generated messages match the expected messages diff --git a/shinkai-bin/shinkai-node/src/llm_provider/execution/prompts/subprompts.rs b/shinkai-bin/shinkai-node/src/llm_provider/execution/prompts/subprompts.rs index 257f095a3..b7d99a9d7 100644 --- a/shinkai-bin/shinkai-node/src/llm_provider/execution/prompts/subprompts.rs +++ b/shinkai-bin/shinkai-node/src/llm_provider/execution/prompts/subprompts.rs @@ -12,6 +12,7 @@ use super::prompts::Prompt; #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] pub enum SubPromptType { User, + UserLastMessage, System, Assistant, ExtraContext, @@ -23,6 +24,7 @@ impl fmt::Display for SubPromptType { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { let s = match self { SubPromptType::User => "user", + SubPromptType::UserLastMessage => "user", SubPromptType::System => "system", SubPromptType::Assistant => "assistant", SubPromptType::ExtraContext => "user", diff --git a/shinkai-bin/shinkai-node/src/network/node.rs b/shinkai-bin/shinkai-node/src/network/node.rs index ef7a93a49..335ceba7c 100644 --- a/shinkai-bin/shinkai-node/src/network/node.rs +++ b/shinkai-bin/shinkai-node/src/network/node.rs @@ -31,7 +31,8 @@ use chrono::Utc; use core::panic; use ed25519_dalek::{Signer, SigningKey, VerifyingKey}; use futures::{future::FutureExt, pin_mut, prelude::*, select}; -use rand::Rng; +use rand::rngs::OsRng; +use rand::{Rng, RngCore}; use shinkai_message_primitives::schemas::llm_providers::serialized_llm_provider::SerializedLLMProvider; use shinkai_message_primitives::schemas::shinkai_name::ShinkaiName; use shinkai_message_primitives::schemas::shinkai_network::NetworkMessageType; @@ -139,6 +140,8 @@ pub struct Node { pub default_embedding_model: Arc>, // Supported embedding models for profiles pub supported_embedding_models: Arc>>, + // API V2 Key + pub api_v2_key: String, } impl Node { @@ -163,6 +166,7 @@ impl Node { ws_address: Option, default_embedding_model: EmbeddingModelType, supported_embedding_models: Vec, + api_v2_key: Option, ) -> Arc> { // if is_valid_node_identity_name_and_no_subidentities is false panic match ShinkaiName::new(node_name.to_string().clone()) { @@ -347,6 +351,25 @@ impl Node { .await; let sheet_manager = sheet_manager_result.unwrap(); + // It reads the api_v2_key from env, if not from db and if not, then it generates a new one that gets saved in the db + let api_v2_key = if let Some(key) = api_v2_key { + db_arc + .set_api_v2_key(&key) + .expect("Failed to set api_v2_key in the database"); + key + } else { + match db_arc.read_api_v2_key() { + Ok(Some(key)) => key, + Ok(None) | Err(_) => { + let new_key = Node::generate_api_v2_key(); + db_arc + .set_api_v2_key(&new_key) + .expect("Failed to set api_v2_key in the database"); + new_key + } + } + }; + Arc::new(Mutex::new(Node { node_name: node_name.clone(), identity_secret_key: clone_signature_secret_key(&identity_secret_key), @@ -383,6 +406,7 @@ impl Node { tool_router: Some(Arc::new(Mutex::new(tool_router))), default_embedding_model, supported_embedding_models, + api_v2_key, })) } @@ -460,7 +484,12 @@ impl Node { let reinstall_tools = std::env::var("REINSTALL_TOOLS").unwrap_or_else(|_| "false".to_string()) == "true"; tokio::spawn(async move { - let current_version = tool_router.lock().await.get_current_lancedb_version().await.unwrap_or(None); + let current_version = tool_router + .lock() + .await + .get_current_lancedb_version() + .await + .unwrap_or(None); if reinstall_tools || current_version != Some(LATEST_ROUTER_DB_VERSION.to_string()) { if let Err(e) = tool_router.lock().await.force_reinstall_all(generator).await { eprintln!("ToolRouter force reinstall failed: {:?}", e); @@ -1457,6 +1486,12 @@ impl Node { Err(e) => eprintln!("Failed to read validation data: {}", e), } } + + fn generate_api_v2_key() -> String { + let mut key = [0u8; 32]; // 256-bit key + OsRng.fill_bytes(&mut key); + base64::encode(&key) + } } impl Drop for Node { diff --git a/shinkai-bin/shinkai-node/src/network/node_api_router.rs b/shinkai-bin/shinkai-node/src/network/node_api_router.rs index 86c526963..432fc187e 100644 --- a/shinkai-bin/shinkai-node/src/network/node_api_router.rs +++ b/shinkai-bin/shinkai-node/src/network/node_api_router.rs @@ -122,24 +122,17 @@ pub async fn run_api( ); println!("API server running on http://{}", address); - if env::var("API_V2_KEY").is_ok() { - let v2_routes = warp::path("v2").and( - v2_routes(node_commands_sender.clone(), node_name.clone()) - .recover(handle_rejection) - .with(log) - .with(cors.clone()), - ); - - // Combine all routes - let routes = v1_routes.or(v2_routes).with(log).with(cors); + let v2_routes = warp::path("v2").and( + v2_routes(node_commands_sender.clone(), node_name.clone()) + .recover(handle_rejection) + .with(log) + .with(cors.clone()), + ); - warp::serve(routes).run(address).await; - } else { - // Combine all routes - let routes = v1_routes.with(log).with(cors); + // Combine all routes + let routes = v1_routes.or(v2_routes).with(log).with(cors); - warp::serve(routes).run(address).await; - } + warp::serve(routes).run(address).await; Ok(()) } diff --git a/shinkai-bin/shinkai-node/src/network/v1_api/api_v1_commands.rs b/shinkai-bin/shinkai-node/src/network/v1_api/api_v1_commands.rs index 4e7d23209..3f579a8e4 100644 --- a/shinkai-bin/shinkai-node/src/network/v1_api/api_v1_commands.rs +++ b/shinkai-bin/shinkai-node/src/network/v1_api/api_v1_commands.rs @@ -892,6 +892,20 @@ impl Node { identity_type: standard_identity_type, permission_type, }; + + let api_v2_key = match db.read_api_v2_key() { + Ok(Some(api_key)) => api_key, + Ok(None) | Err(_) => { + let api_error = APIError { + code: StatusCode::UNAUTHORIZED.as_u16(), + error: "Unauthorized".to_string(), + message: "Invalid bearer token".to_string(), + }; + let _ = res.send(Err(api_error)).await; + return Ok(()); + } + }; + let mut subidentity_manager = identity_manager.lock().await; match subidentity_manager.add_profile_subidentity(subidentity).await { Ok(_) => { @@ -900,6 +914,7 @@ impl Node { node_name: node_name.get_node_name_string().clone(), encryption_public_key: encryption_public_key_to_string(encryption_public_key), identity_public_key: signature_public_key_to_string(identity_public_key), + api_v2_key }; let _ = res.send(Ok(success_response)).await.map_err(|_| ()); } @@ -984,6 +999,19 @@ impl Node { permission_type, }; + let api_v2_key = match db.read_api_v2_key() { + Ok(Some(api_key)) => api_key, + Ok(None) | Err(_) => { + let api_error = APIError { + code: StatusCode::UNAUTHORIZED.as_u16(), + error: "Unauthorized".to_string(), + message: "Invalid bearer token".to_string(), + }; + let _ = res.send(Err(api_error)).await; + return Ok(()); + } + }; + let mut identity_manager_mut = identity_manager.lock().await; match identity_manager_mut.add_device_subidentity(device_identity).await { Ok(_) => { @@ -1009,6 +1037,7 @@ impl Node { node_name: node_name.get_node_name_string().clone(), encryption_public_key: encryption_public_key_to_string(encryption_public_key), identity_public_key: signature_public_key_to_string(identity_public_key), + api_v2_key, }; let _ = res.send(Ok(success_response)).await.map_err(|_| ()); } diff --git a/shinkai-bin/shinkai-node/src/network/v1_api/api_v1_handlers.rs b/shinkai-bin/shinkai-node/src/network/v1_api/api_v1_handlers.rs index 29f44396f..1a62dc609 100644 --- a/shinkai-bin/shinkai-node/src/network/v1_api/api_v1_handlers.rs +++ b/shinkai-bin/shinkai-node/src/network/v1_api/api_v1_handlers.rs @@ -1251,6 +1251,7 @@ pub struct APIUseRegistrationCodeSuccessResponse { pub node_name: String, pub encryption_public_key: String, pub identity_public_key: String, + pub api_v2_key: String, } pub async fn use_registration_code_handler( @@ -1274,16 +1275,12 @@ pub async fn use_registration_code_handler( "message": success_response.message, "node_name": success_response.node_name, "encryption_public_key": success_response.encryption_public_key, - "identity_public_key": success_response.identity_public_key + "identity_public_key": success_response.identity_public_key, + "api_v2_key": success_response.api_v2_key }); let response = serde_json::json!({ "status": "success", - "data": data, - // TODO: remove the below repeated data once the Apps have updated - "message": success_response.message, - "node_name": success_response.node_name, - "encryption_public_key": success_response.encryption_public_key, - "identity_public_key": success_response.identity_public_key + "data": data }); Ok(warp::reply::with_status(warp::reply::json(&response), StatusCode::OK)) } diff --git a/shinkai-bin/shinkai-node/src/network/v2_api/api_v2_commands.rs b/shinkai-bin/shinkai-node/src/network/v2_api/api_v2_commands.rs index d988107af..e5f452fa7 100644 --- a/shinkai-bin/shinkai-node/src/network/v2_api/api_v2_commands.rs +++ b/shinkai-bin/shinkai-node/src/network/v2_api/api_v2_commands.rs @@ -49,21 +49,39 @@ use super::api_v2_handlers_general::InitialRegistrationRequest; impl Node { pub async fn validate_bearer_token( bearer: &str, - _db: Arc, + db: Arc, res: &Sender>, ) -> Result<(), ()> { // Compare bearer token to the environment variable API_V2_KEY - match env::var("API_V2_KEY") { - Ok(api_key) if api_key == bearer => Ok(()), - _ => { - let api_error = APIError { - code: StatusCode::UNAUTHORIZED.as_u16(), - error: "Unauthorized".to_string(), - message: "Invalid bearer token".to_string(), - }; - let _ = res.send(Err(api_error)).await; - Err(()) + let api_key = match env::var("API_V2_KEY") { + Ok(api_key) => api_key, + Err(_) => { + // If the environment variable is not set, read from the database + match db.read_api_v2_key() { + Ok(Some(api_key)) => api_key, + Ok(None) | Err(_) => { + let api_error = APIError { + code: StatusCode::UNAUTHORIZED.as_u16(), + error: "Unauthorized".to_string(), + message: "Invalid bearer token".to_string(), + }; + let _ = res.send(Err(api_error)).await; + return Err(()); + } + } } + }; + + if api_key == bearer { + Ok(()) + } else { + let api_error = APIError { + code: StatusCode::UNAUTHORIZED.as_u16(), + error: "Unauthorized".to_string(), + message: "Invalid bearer token".to_string(), + }; + let _ = res.send(Err(api_error)).await; + Err(()) } } diff --git a/shinkai-bin/shinkai-node/src/runner.rs b/shinkai-bin/shinkai-node/src/runner.rs index b827e4dab..2af370709 100644 --- a/shinkai-bin/shinkai-node/src/runner.rs +++ b/shinkai-bin/shinkai-node/src/runner.rs @@ -185,6 +185,7 @@ pub async fn initialize_node() -> Result< node_env.ws_address, node_env.default_embedding_model.clone(), node_env.supported_embedding_models.clone(), + node_env.api_v2_key.clone(), ) .await; diff --git a/shinkai-bin/shinkai-node/src/utils/environment.rs b/shinkai-bin/shinkai-node/src/utils/environment.rs index 868d1c55f..cedb01152 100644 --- a/shinkai-bin/shinkai-node/src/utils/environment.rs +++ b/shinkai-bin/shinkai-node/src/utils/environment.rs @@ -28,6 +28,7 @@ pub struct NodeEnvironment { pub proxy_identity: Option, pub default_embedding_model: EmbeddingModelType, pub supported_embedding_models: Vec, + pub api_v2_key: Option, } #[derive(Debug, Clone)] @@ -193,7 +194,10 @@ pub fn fetch_node_environment() -> NodeEnvironment { OllamaTextEmbeddingsInference::SnowflakeArcticEmbed_M, )] }); - + + // Fetch the API_V2_KEY environment variable + let api_v2_key: Option = env::var("API_V2_KEY").ok(); + NodeEnvironment { global_identity_name, listen_address, @@ -213,6 +217,7 @@ pub fn fetch_node_environment() -> NodeEnvironment { proxy_identity, default_embedding_model, supported_embedding_models, + api_v2_key } } diff --git a/shinkai-bin/shinkai-node/tests/it/a0_subscription_manager_tests.rs b/shinkai-bin/shinkai-node/tests/it/a0_subscription_manager_tests.rs index 87e42c267..f7ad0a1d0 100644 --- a/shinkai-bin/shinkai-node/tests/it/a0_subscription_manager_tests.rs +++ b/shinkai-bin/shinkai-node/tests/it/a0_subscription_manager_tests.rs @@ -105,6 +105,7 @@ fn subscription_manager_test() { None, default_embedding_model(), supported_embedding_models(), + None, ) .await; @@ -127,6 +128,7 @@ fn subscription_manager_test() { None, default_embedding_model(), supported_embedding_models(), + None, ) .await; diff --git a/shinkai-bin/shinkai-node/tests/it/a1_http_subscription_tests.rs b/shinkai-bin/shinkai-node/tests/it/a1_http_subscription_tests.rs index 92cb6abb2..c2883c969 100644 --- a/shinkai-bin/shinkai-node/tests/it/a1_http_subscription_tests.rs +++ b/shinkai-bin/shinkai-node/tests/it/a1_http_subscription_tests.rs @@ -116,6 +116,7 @@ fn http_subscription_manager_test() { None, default_embedding_model(), supported_embedding_models(), + None, ) .await; @@ -138,6 +139,7 @@ fn http_subscription_manager_test() { None, default_embedding_model(), supported_embedding_models(), + None, ) .await; diff --git a/shinkai-bin/shinkai-node/tests/it/change_nodes_name_tests.rs b/shinkai-bin/shinkai-node/tests/it/change_nodes_name_tests.rs index ff396f039..780ca99f6 100644 --- a/shinkai-bin/shinkai-node/tests/it/change_nodes_name_tests.rs +++ b/shinkai-bin/shinkai-node/tests/it/change_nodes_name_tests.rs @@ -69,6 +69,7 @@ fn change_nodes_name_test() { None, default_embedding_model(), supported_embedding_models(), + None, ); let node1_handler = tokio::spawn(async move { @@ -163,6 +164,7 @@ fn change_nodes_name_test() { None, default_embedding_model(), supported_embedding_models(), + None, ); let node1_handler = tokio::spawn(async move { diff --git a/shinkai-bin/shinkai-node/tests/it/db_restore_tests.rs b/shinkai-bin/shinkai-node/tests/it/db_restore_tests.rs index a7cd74d13..fd4ea926e 100644 --- a/shinkai-bin/shinkai-node/tests/it/db_restore_tests.rs +++ b/shinkai-bin/shinkai-node/tests/it/db_restore_tests.rs @@ -61,6 +61,7 @@ mod tests { None, default_embedding_model(), supported_embedding_models(), + None, ); let node1_handler = tokio::spawn(async move { diff --git a/shinkai-bin/shinkai-node/tests/it/llm_provider_integration_tests.rs b/shinkai-bin/shinkai-node/tests/it/llm_provider_integration_tests.rs index fd8890ff4..1fac66c8e 100644 --- a/shinkai-bin/shinkai-node/tests/it/llm_provider_integration_tests.rs +++ b/shinkai-bin/shinkai-node/tests/it/llm_provider_integration_tests.rs @@ -145,6 +145,7 @@ fn node_llm_provider_registration() { None, default_embedding_model(), supported_embedding_models(), + None, ); let node1_handler = tokio::spawn(async move { diff --git a/shinkai-bin/shinkai-node/tests/it/node_integration_tests.rs b/shinkai-bin/shinkai-node/tests/it/node_integration_tests.rs index c22b06467..dcaff9bfb 100644 --- a/shinkai-bin/shinkai-node/tests/it/node_integration_tests.rs +++ b/shinkai-bin/shinkai-node/tests/it/node_integration_tests.rs @@ -106,6 +106,7 @@ fn subidentity_registration() { None, default_embedding_model(), supported_embedding_models(), + None, ) .await; @@ -128,6 +129,7 @@ fn subidentity_registration() { None, default_embedding_model(), supported_embedding_models(), + None, ) .await; diff --git a/shinkai-bin/shinkai-node/tests/it/node_retrying_tests.rs b/shinkai-bin/shinkai-node/tests/it/node_retrying_tests.rs index 242474000..e75fb6e5b 100644 --- a/shinkai-bin/shinkai-node/tests/it/node_retrying_tests.rs +++ b/shinkai-bin/shinkai-node/tests/it/node_retrying_tests.rs @@ -103,6 +103,7 @@ fn node_retrying_test() { None, default_embedding_model(), supported_embedding_models(), + None, ) .await; @@ -125,6 +126,7 @@ fn node_retrying_test() { None, default_embedding_model(), supported_embedding_models(), + None, ) .await; diff --git a/shinkai-bin/shinkai-node/tests/it/tcp_proxy_tests.rs b/shinkai-bin/shinkai-node/tests/it/tcp_proxy_tests.rs index 8299be6fb..e6fcd0a47 100644 --- a/shinkai-bin/shinkai-node/tests/it/tcp_proxy_tests.rs +++ b/shinkai-bin/shinkai-node/tests/it/tcp_proxy_tests.rs @@ -118,6 +118,7 @@ fn tcp_proxy_test_identity() { None, default_embedding_model(), supported_embedding_models(), + None, ) .await; @@ -140,6 +141,7 @@ fn tcp_proxy_test_identity() { None, default_embedding_model(), supported_embedding_models(), + None, ) .await; @@ -524,6 +526,7 @@ fn tcp_proxy_test_localhost() { None, default_embedding_model(), supported_embedding_models(), + None, ) .await; @@ -546,6 +549,7 @@ fn tcp_proxy_test_localhost() { None, default_embedding_model(), supported_embedding_models(), + None, ) .await; diff --git a/shinkai-bin/shinkai-node/tests/it/utils/test_boilerplate.rs b/shinkai-bin/shinkai-node/tests/it/utils/test_boilerplate.rs index 1a53b9a6a..d1bf23582 100644 --- a/shinkai-bin/shinkai-node/tests/it/utils/test_boilerplate.rs +++ b/shinkai-bin/shinkai-node/tests/it/utils/test_boilerplate.rs @@ -112,6 +112,8 @@ where // Fetch the PROXY_ADDRESS environment variable let proxy_identity: Option = env::var("PROXY_IDENTITY").ok().and_then(|addr| addr.parse().ok()); + let api_v2_key = env::var("API_V2_KEY").unwrap_or_else(|_| "SUPER_SECRET".to_string()); + // Create node1 and node2 let addr1 = SocketAddr::new(IpAddr::V4(Ipv4Addr::new(127, 0, 0, 1)), 8080); let node1 = Node::new( @@ -132,6 +134,7 @@ where None, default_embedding_model(), supported_embedding_models(), + Some(api_v2_key), ) .await; diff --git a/shinkai-bin/shinkai-node/tests/it/workflow_integration_tests.rs b/shinkai-bin/shinkai-node/tests/it/workflow_integration_tests.rs index 4b6a711f7..ae84558c6 100644 --- a/shinkai-bin/shinkai-node/tests/it/workflow_integration_tests.rs +++ b/shinkai-bin/shinkai-node/tests/it/workflow_integration_tests.rs @@ -137,6 +137,7 @@ fn workflow_integration_test() { None, default_embedding_model(), supported_embedding_models(), + None, ); let node1_handler = tokio::spawn(async move { @@ -398,6 +399,7 @@ fn workflow_complex_integration_test() { None, default_embedding_model(), supported_embedding_models(), + None, ); let node1_handler = tokio::spawn(async move { diff --git a/shinkai-bin/shinkai-node/tests/it/z_shinkai_mirror_tests.rs b/shinkai-bin/shinkai-node/tests/it/z_shinkai_mirror_tests.rs index 4d35d2ba9..527c8d1de 100644 --- a/shinkai-bin/shinkai-node/tests/it/z_shinkai_mirror_tests.rs +++ b/shinkai-bin/shinkai-node/tests/it/z_shinkai_mirror_tests.rs @@ -216,6 +216,7 @@ fn mirror_sync_tests() { None, default_embedding_model(), supported_embedding_models(), + None, ) .await;