Skip to content

Commit

Permalink
Implemented proper local/db storage
Browse files Browse the repository at this point in the history
  • Loading branch information
@robkorn
robkorn committed Oct 12, 2023
1 parent 4868bae commit 49ce36c
Show file tree
Hide file tree
Showing 2 changed files with 55 additions and 21 deletions.
56 changes: 38 additions & 18 deletions src/agent/execution/job_execution_core.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ use crate::db::ShinkaiDB;
use crate::resources::bert_cpp::BertCPPProcess;
use crate::resources::file_parsing::FileParser;
use serde_json::Value as JsonValue;
use shinkai_message_primitives::shinkai_utils::job_scope::{LocalScopeEntry, ScopeEntry};
use shinkai_message_primitives::shinkai_utils::job_scope::{DBScopeEntry, LocalScopeEntry, ScopeEntry};
use shinkai_message_primitives::{
schemas::shinkai_name::ShinkaiName,
shinkai_message::{shinkai_message::ShinkaiMessage, shinkai_message_schemas::JobMessage},
Expand All @@ -16,6 +16,7 @@ use shinkai_message_primitives::{
use shinkai_vector_resources::base_vector_resources::BaseVectorResource;
use shinkai_vector_resources::embedding_generator::RemoteEmbeddingGenerator;
use shinkai_vector_resources::source::{SourceDocumentType, SourceFile, SourceFileType};
use shinkai_vector_resources::vector_resource::VectorResource;
use std::result::Result::Ok;
use std::time::Instant;
use std::{collections::HashMap, sync::Arc};
Expand Down Expand Up @@ -71,7 +72,7 @@ impl AgentManager {
self.fetch_relevant_job_data(job.job_id()).await?;

// Processes any files which were sent with the job message
self.process_job_message_files(&job_message, agent_found.clone(), &mut full_job, profile)
self.process_job_message_files(&job_message, agent_found.clone(), &mut full_job, profile, false)
.await?;

// TODO(Nico): move this to a parallel thread that runs in the background
Expand Down Expand Up @@ -148,6 +149,7 @@ impl AgentManager {
agent_found: Option<Arc<Mutex<Agent>>>,
full_job: &mut Job,
profile: ShinkaiName,
save_to_db_directly: bool,
) -> Result<(), AgentError> {
if !job_message.files_inbox.is_empty() {
println!(
Expand All @@ -156,7 +158,13 @@ impl AgentManager {
);
// TODO: later we should able to grab errors and return them to the user
let new_scope_entries = self
.process_files_inbox(self.db.clone(), agent_found, job_message.files_inbox.clone(), profile)
.process_files_inbox(
self.db.clone(),
agent_found,
job_message.files_inbox.clone(),
profile,
save_to_db_directly,
)
.await?;
eprintln!(">>> new_scope_entries: {:?}", new_scope_entries.keys());

Expand Down Expand Up @@ -203,6 +211,7 @@ impl AgentManager {
agent: Option<Arc<Mutex<Agent>>>,
files_inbox: String,
profile: ShinkaiName,
save_to_db_directly: bool,
) -> Result<HashMap<String, ScopeEntry>, AgentError> {
// Handle the None case if the agent is not found
let agent = match agent {
Expand Down Expand Up @@ -231,9 +240,11 @@ impl AgentManager {
let pdf_overview = FileParser::parse_pdf_for_keywords_and_description(&content, 3, 200)?;
let grouped_text_list_clone = pdf_overview.grouped_text_list.clone();
let prompt = JobPromptGenerator::simple_doc_description(grouped_text_list_clone);
let description_response = self
.inference_agent_and_extract(agent.clone(), prompt, "answer")
.await?;
let description_response = Self::ending_stripper(
&self
.inference_agent_and_extract(agent.clone(), prompt, "answer")
.await?,
);

let vrsource = Self::create_vrsource(
&filename,
Expand All @@ -246,26 +257,35 @@ impl AgentManager {
&*generator,
&filename,
Some(&description_response),
vrsource,
vrsource.clone(),
&vec![],
)?;

// TODO: Maybe add: "\nKeywords: keywords_generated_by_RAKE"?
eprintln!("description_response: {:?}", description_response);

// Now create Local/DBScopeEntry
let resource = BaseVectorResource::from(doc.clone());
shinkai_db.init_profile_resource_router(&profile)?;
shinkai_db.save_resource(resource, &profile).unwrap();
if save_to_db_directly {
shinkai_db.init_profile_resource_router(&profile)?;
shinkai_db.save_resource(resource, &profile).unwrap();

let local_scope_entry = LocalScopeEntry {
resource: BaseVectorResource::from(doc.clone()),
source: SourceFile::new(
filename.clone(),
SourceFileType::Document(SourceDocumentType::Pdf),
content,
),
};
files_map.insert(filename, ScopeEntry::Local(local_scope_entry));
let db_scope_entry = DBScopeEntry {
resource_pointer: doc.get_resource_pointer(),
source: vrsource,
};
files_map.insert(filename, ScopeEntry::Database(db_scope_entry));
} else {
let local_scope_entry = LocalScopeEntry {
resource: BaseVectorResource::from(doc.clone()),
source: SourceFile::new(
filename.clone(),
SourceFileType::Document(SourceDocumentType::Pdf),
content,
),
};
files_map.insert(filename, ScopeEntry::Local(local_scope_entry));
}
}
}

Expand Down
20 changes: 17 additions & 3 deletions src/agent/execution/job_prompts.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -76,7 +76,7 @@ impl JobPromptGenerator {
) -> Prompt {
let mut prompt = Prompt::new();
prompt.add_content(
"You are an advanced assistant who only has access to the provided content and your own knowledge to answer any question the user provides.".to_string(),
"You are an advanced assistant who only has access to the provided content and your own knowledge to answer any question the user provides. Do not ask for further context or information in your answer to the user, but simply tell the user as much information as possible.".to_string(),
SubPromptType::System,
);

Expand Down Expand Up @@ -132,6 +132,11 @@ impl JobPromptGenerator {
SubPromptType::System,
);

prompt.add_content(
format!("Do not mention needing further context, or information, or ask for more research, just directly provide as much information as you know:"),
SubPromptType::System,
);

prompt
}

Expand All @@ -144,7 +149,7 @@ impl JobPromptGenerator {
) -> Prompt {
let mut prompt = Prompt::new();
prompt.add_content(
"You are an advanced assistant who only has access to the provided content and your own knowledge to answer any question the user provides.".to_string(),
"You are an advanced assistant who only has access to the provided content and your own knowledge to answer any question the user provides. Do not ask for further context or information in your answer to the user, but simply tell the user as much information as possible.".to_string(),
SubPromptType::System,
);

Expand Down Expand Up @@ -177,6 +182,10 @@ impl JobPromptGenerator {

prompt.add_ebnf(String::from(r#""{" "answer" ":" string "}""#), SubPromptType::System);

prompt.add_content(
format!("Do not mention needing further context, or information, or ask for more research, just directly provide as much information as you know:"),
SubPromptType::System,
);
prompt
}

Expand Down Expand Up @@ -209,7 +218,7 @@ impl JobPromptGenerator {
pub fn simple_doc_description(chunks: Vec<String>) -> Prompt {
let mut prompt = Prompt::new();
prompt.add_content(
r#"You are an advanced assistant who is specialized in summarizing information."#.to_string(),
"You are an advanced assistant who who is specialized in summarizing information. Do not ask for further context or information in your answer, simply summarize as much information as possible.".to_string(),
SubPromptType::System,
);

Expand All @@ -223,6 +232,11 @@ impl JobPromptGenerator {
);
prompt.add_ebnf(String::from(r#""{" "answer" ":" string "}""#), SubPromptType::System);

prompt.add_content(
format!("Do not mention needing further context, or information, or ask for more research, just directly provide as much information as you know:"),
SubPromptType::System,
);

prompt
}

Expand Down

0 comments on commit 49ce36c

Please sign in to comment.