Skip to content

Commit

Permalink
Refactor
Browse files Browse the repository at this point in the history
  • Loading branch information
@acedward
acedward committed Jan 14, 2025
1 parent 9ab2691 commit 0114317
Show file tree
Hide file tree
Showing 4 changed files with 208 additions and 185 deletions.
75 changes: 31 additions & 44 deletions ...n/shinkai-node/src/llm_provider/execution/chains/generic_chain/generic_inference_chain.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ use crate::network::agent_payments_manager::my_agent_offerings_manager::MyAgentO
use crate::utils::environment::{fetch_node_environment, NodeEnvironment};
use async_trait::async_trait;
use shinkai_embedding::embedding_generator::RemoteEmbeddingGenerator;
use shinkai_fs::shinkai_file_manager::ShinkaiFileManager;
use shinkai_message_primitives::schemas::inbox_name::InboxName;
use shinkai_message_primitives::schemas::job::{Job, JobLike};
use shinkai_message_primitives::schemas::llm_providers::common_agent_llm_provider::ProviderOrAgent;
Expand Down Expand Up @@ -344,29 +345,23 @@ impl GenericInferenceChain {
}
});

let mut additional_files = vec![];
additional_files.extend(
merged_fs_files_paths
.clone()
.iter()
.map(|path| format!("{:?}", path.path.clone())),
);
additional_files.extend(
merged_fs_folder_paths
.clone()
.iter()
.map(|path| format!("{:?}", path.path.clone())),
);
if !job_filenames.is_empty() {
let folder_path = db.get_job_folder_name(&full_job.job_id.clone());
if let Ok(folder_path) = folder_path {
additional_files.extend(job_filenames.iter().map(|path| {
let folder_path = fs::canonicalize(folder_path.path.clone()).unwrap_or_default();
format!("{}/{}", folder_path.to_string_lossy().replace("\"", ""), path.clone())
}));
}
let mut additional_files: Vec<String> = vec![];
let f = ShinkaiFileManager::get_absolute_path_for_additional_files(
merged_fs_files_paths.clone(),
merged_fs_folder_paths.clone(),
)?;
additional_files.extend(f);
let folder_path: Result<ShinkaiPath, shinkai_sqlite::errors::SqliteManagerError> =
db.get_job_folder_name(&full_job.job_id.clone());
if let Ok(folder_path) = folder_path {
additional_files.extend(ShinkaiFileManager::get_absolute_paths_with_folder(
job_filenames.clone(),
folder_path.path.clone(),
));
}

let mut filled_prompt = JobPromptGenerator::generic_inference_prompt(
db.clone(),
custom_system_prompt.clone(),
custom_prompt.clone(),
user_message.clone(),
Expand All @@ -376,7 +371,6 @@ impl GenericInferenceChain {
Some(full_job.step_history.clone()),
tools.clone(),
None,
full_job.scope().clone(),
full_job.job_id.clone(),
additional_files,
node_env.clone(),
Expand Down Expand Up @@ -499,30 +493,24 @@ impl GenericInferenceChain {
last_function_response = Some(function_response);
}

let mut additional_files = vec![];
additional_files.extend(
merged_fs_files_paths
.clone()
.iter()
.map(|path| format!("{:?}", path.path.clone())),
);
additional_files.extend(
merged_fs_folder_paths
.clone()
.iter()
.map(|path| format!("{:?}", path.path.clone())),
);
if !job_filenames.is_empty() {
let folder_path = db.get_job_folder_name(&full_job.job_id.clone());
if let Ok(folder_path) = folder_path {
additional_files.extend(job_filenames.iter().map(|path| {
let folder_path = fs::canonicalize(folder_path.path.clone()).unwrap_or_default();
format!("{}/{}", folder_path.to_string_lossy().replace("\"", ""), path.clone())
}));
}
let mut additional_files: Vec<String> = vec![];
let f = ShinkaiFileManager::get_absolute_path_for_additional_files(
merged_fs_files_paths.clone(),
merged_fs_folder_paths.clone(),
)?;
additional_files.extend(f);
let folder_path: Result<ShinkaiPath, shinkai_sqlite::errors::SqliteManagerError> =
db.get_job_folder_name(&full_job.job_id.clone());
if let Ok(folder_path) = folder_path {
additional_files.extend(ShinkaiFileManager::get_absolute_paths_with_folder(
job_filenames.clone(),
folder_path.path.clone(),
));
}

// 7) Call LLM again with the response (for formatting)
filled_prompt = JobPromptGenerator::generic_inference_prompt(
db.clone(),
custom_system_prompt.clone(),
custom_prompt.clone(),
user_message.clone(),
Expand All @@ -532,7 +520,6 @@ impl GenericInferenceChain {
Some(full_job.step_history.clone()),
tools.clone(),
last_function_response,
full_job.scope().clone(),
full_job.job_id.clone(),
additional_files,
node_env.clone(),
Expand Down
122 changes: 9 additions & 113 deletions shinkai-bin/shinkai-node/src/llm_provider/execution/chains/generic_chain/generic_prompts.rs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
use serde_json::json;
use shinkai_fs::shinkai_file_manager::ShinkaiFileManager;
use shinkai_sqlite::SqliteManager;
use std::{collections::HashMap, fs};

Expand All @@ -20,6 +21,7 @@ impl JobPromptGenerator {
/// summary_text is the content generated by an LLM on parsing (if exist)
#[allow(clippy::too_many_arguments)]
pub fn generic_inference_prompt(
db: Arc<SqliteManager>,
custom_system_prompt: Option<String>,
custom_user_prompt: Option<String>,
user_message: String,
Expand All @@ -29,7 +31,6 @@ impl JobPromptGenerator {
job_step_history: Option<Vec<ShinkaiMessage>>,
tools: Vec<ShinkaiTool>,
function_call: Option<ToolCallFunctionResponse>,
job_scope: MinimalJobScope,
job_id: String,
additional_files: Vec<String>,
_node_env: NodeEnvironment,
Expand Down Expand Up @@ -94,124 +95,19 @@ impl JobPromptGenerator {
}
}

// job_scope.
// {
// // Process fs_files_paths
// for path in &fs_files_paths {
// if let Some(parsed_file) = sqlite_manager.get_parsed_file_by_shinkai_path(path).unwrap() {
// let file_id = parsed_file.id.unwrap();
// parsed_file_ids.push(file_id);
// paths_map.insert(file_id, path.clone());

// if let Some(file_tokens) = parsed_file.total_tokens {
// total_tokens += file_tokens;
// } else {
// all_files_have_token_count = false;
// }
// }
// }

// // Process job_filenames
// for filename in &job_filenames {
// let file_path =
// match ShinkaiFileManager::construct_job_file_path(&job_id, filename, &sqlite_manager) {
// Ok(path) => path,
// Err(_) => continue,
// };

// if let Some(parsed_file) = sqlite_manager.get_parsed_file_by_shinkai_path(&file_path).unwrap() {
// let file_id = parsed_file.id.unwrap();
// parsed_file_ids.push(file_id);
// paths_map.insert(file_id, file_path);

// if let Some(file_tokens) = parsed_file.total_tokens {
// total_tokens += file_tokens;
// } else {
// all_files_have_token_count = false;
// }
// }
// }

// Retrieve each file in the job scope
// for path in &job_scope.vector_fs_items {
// if let Some(parsed_file) = db.get_parsed_file_by_shinkai_path(path).unwrap() {
// let file_id = parsed_file.id.unwrap();
// parsed_file_ids.push(file_id);
// paths_map.insert(file_id, path.clone());

// // Count tokens while we're here
// if let Some(file_tokens) = parsed_file.total_tokens {
// total_tokens += file_tokens;
// } else {
// all_files_have_token_count = false;
// }
// }
// }

// // Process fs_folder_paths
// for folder in &fs_folder_paths {
// Self::process_folder_contents(
// folder,
// &sqlite_manager,
// &mut parsed_file_ids,
// &mut paths_map,
// &mut total_tokens,
// &mut all_files_have_token_count,
// )
// .await?;
// }

// // Retrieve files inside vector_fs_folders
// for folder in &scope.vector_fs_folders {
// Self::process_folder_contents(
// folder,
// &sqlite_manager,
// &mut parsed_file_ids,
// &mut paths_map,
// &mut total_tokens,
// &mut all_files_have_token_count,
// )
// .await?;
// }
let file_in_scope = (ShinkaiFileManager::get_absolute_path_for_job_scope(&db, &job_id));
let mut all_files = additional_files.clone();
if let Ok(file_in_scope) = file_in_scope {
all_files.extend(file_in_scope);
}

// }
// job_id -> full-job -> job.scope -> shinkai-pathws -> file-info
let mut file_in_scope = additional_files.clone();
// let folder_path = db.get_job_folder_name(&job_id);
// if let Ok(folder_path) = folder_path {
job_scope.vector_fs_items.iter().for_each(|path| {
let file_path = format!("{:?}", fs::canonicalize(path.path.clone()).unwrap_or_default());
file_in_scope.push(file_path.replace("\"", ""));
});
// }
if !file_in_scope.is_empty() {
if !additional_files.is_empty() {
prompt.add_content(
format!("<cursrent_files>\n{}\n</current_files>\n", file_in_scope.join("\n")),
format!("<current_files>\n{}\n</current_files>\n", additional_files.join("\n")),
SubPromptType::ExtraContext,
97,
);
}

// let files_result = ShinkaiFileManager::get_all_files_and_folders_for_job(&job_id, &db);
// if let Ok(folder_path) = folder_path {
// if let Ok(files_result) = files_result {
// let files_result_str = files_result
// .iter()
// .map(|file| {
// format!(
// "{:?}/{}",
// fs::canonicalize(folder_path.path.clone()).unwrap_or_default(),
// file.name
// )
// })
// .collect::<Vec<String>>();
// prompt.add_content(
// format!("<current_files>\n{}\n</current_files>\n", files_result_str.join("\n")),
// SubPromptType::ExtraContext,
// 97,
// );
// }
// }
}

// Parses the retrieved nodes as individual sub-prompts, to support priority pruning
Expand Down
4 changes: 2 additions & 2 deletions ...shinkai-node/src/llm_provider/execution/chains/sheet_ui_chain/sheet_ui_inference_chain.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -281,6 +281,7 @@ impl SheetUIInferenceChain {
eprintln!("Extended user message: {}", extended_user_message);

let mut filled_prompt = JobPromptGenerator::generic_inference_prompt(
db.clone(),
None, // No custom prompt
None, // TODO: connect later on
extended_user_message.clone(),
Expand All @@ -290,7 +291,6 @@ impl SheetUIInferenceChain {
Some(full_job.step_history.clone()),
tools.clone(),
None,
full_job.scope().clone(),
full_job.job_id.clone(),
vec![],
node_env.clone(),
Expand Down Expand Up @@ -431,6 +431,7 @@ impl SheetUIInferenceChain {

// 7) Call LLM again with the response (for formatting)
filled_prompt = JobPromptGenerator::generic_inference_prompt(
db.clone(),
None, // TODO: connect later on
None, // TODO: connect later on
user_message.clone(),
Expand All @@ -440,7 +441,6 @@ impl SheetUIInferenceChain {
Some(full_job.step_history.clone()),
tools.clone(),
last_function_response,
full_job.scope().clone(),
full_job.job_id.clone(),
vec![],
node_env.clone(),
Expand Down
Loading

0 comments on commit 0114317

Please sign in to comment.