From aed243c3aa0dd6d6c7dcba088304170f9d5cb696 Mon Sep 17 00:00:00 2001 From: sigoden Date: Sun, 10 Mar 2024 08:01:54 +0800 Subject: feat: allow use of temporary role in a session (#348) --- src/client/message.rs | 7 ++++++- src/config/input.rs | 46 ++++++++++++++++++++++++++++++++++++++++++++-- src/config/mod.rs | 18 +++++++++++------- src/main.rs | 10 +++++++--- src/repl/mod.rs | 41 +++++++++++++++-------------------------- 5 files changed, 83 insertions(+), 39 deletions(-) (limited to 'src') diff --git a/src/client/message.rs b/src/client/message.rs index 757b673..1a454ab 100644 --- a/src/client/message.rs +++ b/src/client/message.rs @@ -102,11 +102,16 @@ pub struct ImageUrl { #[cfg(test)] mod tests { use super::*; + use crate::config::InputContext; #[test] fn test_serde() { assert_eq!( - serde_json::to_string(&Message::new(&Input::from_str("Hello World"))).unwrap(), + serde_json::to_string(&Message::new(&Input::from_str( + "Hello World", + InputContext::default() + ))) + .unwrap(), "{\"role\":\"user\",\"content\":\"Hello World\"}" ); } diff --git a/src/config/input.rs b/src/config/input.rs index cb3cdf2..c5ba1a2 100644 --- a/src/config/input.rs +++ b/src/config/input.rs @@ -1,3 +1,6 @@ +use super::role::Role; +use super::session::Session; + use crate::client::{ImageUrl, MessageContent, MessageContentPart, ModelCapabilities}; use crate::utils::sha256sum; @@ -25,18 +28,20 @@ pub struct Input { text: String, medias: Vec, data_urls: HashMap, + context: InputContext, } impl Input { - pub fn from_str(text: &str) -> Self { + pub fn from_str(text: &str, context: InputContext) -> Self { Self { text: text.to_string(), medias: Default::default(), data_urls: Default::default(), + context, } } - pub fn new(text: &str, files: Vec) -> Result { + pub fn new(text: &str, files: Vec, context: InputContext) -> Result { let mut texts = vec![text.to_string()]; let mut medias = vec![]; let mut data_urls = HashMap::new(); @@ -72,13 +77,38 @@ impl Input { text: texts.join("\n"), medias, data_urls, + context, }) } + pub fn is_empty(&self) -> bool { + self.text.is_empty() && self.medias.is_empty() + } + pub fn data_urls(&self) -> HashMap { self.data_urls.clone() } + pub fn role(&self) -> Option<&Role> { + self.context.role.as_ref() + } + + pub fn session<'a>(&self, session: &'a Option) -> Option<&'a Session> { + if self.context.in_session { + session.as_ref() + } else { + None + } + } + + pub fn session_mut<'a>(&self, session: &'a mut Option) -> Option<&'a mut Session> { + if self.context.in_session { + session.as_mut() + } else { + None + } + } + pub fn summary(&self) -> String { let text: String = self .text @@ -154,6 +184,18 @@ impl Input { } } +#[derive(Debug, Clone, Default)] +pub struct InputContext { + role: Option, + in_session: bool, +} + +impl InputContext { + pub fn new(role: Option, in_session: bool) -> Self { + Self { role, in_session } + } +} + pub fn resolve_data_url(data_urls: &HashMap, data_url: String) -> String { if data_url.starts_with("data:") { let hash = sha256sum(&data_url); diff --git a/src/config/mod.rs b/src/config/mod.rs index 65b573a..4ee5319 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -2,7 +2,7 @@ mod input; mod role; mod session; -pub use self::input::Input; +pub use self::input::{Input, InputContext}; use self::role::Role; use self::session::{Session, TEMP_SESSION_NAME}; @@ -226,7 +226,7 @@ impl Config { return Ok(()); } - if let Some(session) = self.session.as_mut() { + if let Some(session) = input.session_mut(&mut self.session) { session.add_message(&input, output)?; return Ok(()); } @@ -241,7 +241,7 @@ impl Config { let timestamp = now(); let summary = input.summary(); let input_markdown = input.render(); - let output = match self.role.as_ref() { + let output = match input.role() { None => { format!("# CHAT: {summary} [{timestamp}]\n{input_markdown}\n--------\n{output}\n--------\n\n",) } @@ -369,9 +369,9 @@ impl Config { } pub fn echo_messages(&self, input: &Input) -> String { - if let Some(session) = self.session.as_ref() { + if let Some(session) = input.session(&self.session) { session.echo_messages(input) - } else if let Some(role) = self.role.as_ref() { + } else if let Some(role) = input.role() { role.echo_messages(input) } else { input.render() @@ -379,9 +379,9 @@ impl Config { } pub fn build_messages(&self, input: &Input) -> Result> { - let messages = if let Some(session) = self.session.as_ref() { + let messages = if let Some(session) = input.session(&self.session) { session.build_emssages(input) - } else if let Some(role) = self.role.as_ref() { + } else if let Some(role) = input.role() { role.build_messages(input) } else { let message = Message::new(input); @@ -762,6 +762,10 @@ impl Config { }) } + pub fn input_context(&self) -> InputContext { + InputContext::new(self.role.clone(), self.has_session()) + } + pub fn maybe_print_send_tokens(&self, input: &Input) { if self.dry_run { if let Ok(messages) = self.build_messages(input) { diff --git a/src/main.rs b/src/main.rs index c967611..589676e 100644 --- a/src/main.rs +++ b/src/main.rs @@ -114,7 +114,11 @@ fn start_directive( if let Some(session) = &config.read().session { session.guard_save()?; } - let input = Input::new(text, include.unwrap_or_default())?; + let input = Input::new( + text, + include.unwrap_or_default(), + config.read().input_context(), + )?; let mut client = init_client(config)?; ensure_model_capabilities(client.as_mut(), input.required_capabilities())?; config.read().maybe_print_send_tokens(&input); @@ -148,7 +152,7 @@ fn start_interactive(config: &GlobalConfig) -> Result<()> { } fn execute(config: &GlobalConfig, text: &str) -> Result<()> { - let input = Input::from_str(text); + let input = Input::from_str(text, config.read().input_context()); let client = init_client(config)?; config.read().maybe_print_send_tokens(&input); let mut eval_str = client.send_message(input.clone())?; @@ -192,7 +196,7 @@ fn execute(config: &GlobalConfig, text: &str) -> Result<()> { if !describe { config.write().set_describe_command_role()?; } - let input = Input::from_str(&eval_str); + let input = Input::from_str(&eval_str, config.read().input_context()); let abort = create_abort_signal(); render_stream(&input, client.as_ref(), config, abort)?; describe = true; diff --git a/src/repl/mod.rs b/src/repl/mod.rs index 51a88bb..b04907c 100644 --- a/src/repl/mod.rs +++ b/src/repl/mod.rs @@ -7,7 +7,7 @@ use self::highlighter::ReplHighlighter; use self::prompt::ReplPrompt; use crate::client::{ensure_model_capabilities, init_client}; -use crate::config::{GlobalConfig, Input, State}; +use crate::config::{GlobalConfig, Input, InputContext, State}; use crate::render::{render_error, render_stream}; use crate::utils::{create_abort_signal, set_text, AbortSignal}; @@ -174,21 +174,10 @@ impl Repl { ".role" => match args { Some(args) => match args.split_once(|c| c == '\n' || c == ' ') { Some((name, text)) => { - if self.config.read().has_session() { - bail!(r#"Cannot perform this action in a session"#); - } else { - let name = name.trim(); - let text = text.trim(); - let old_role = - self.config.read().role.as_ref().map(|v| v.name.to_string()); - self.config.write().set_role(name)?; - let ask_ret = self.ask(text, vec![]); - match old_role { - Some(old_role) => self.config.write().set_role(&old_role)?, - None => self.config.write().clear_role()?, - } - ask_ret?; - } + let role = self.config.read().retrieve_role(name.trim())?; + let input = + Input::from_str(text.trim(), InputContext::new(Some(role), false)); + self.ask(input)?; } None => { self.config.write().set_role(args)?; @@ -219,7 +208,8 @@ impl Repl { None => (args, ""), }; let files = shell_words::split(files).with_context(|| "Invalid args")?; - self.ask(text, files)?; + let input = Input::new(text, files, self.config.read().input_context())?; + self.ask(input)?; } None => println!("Usage: .file ...[ -- ...]"), }, @@ -250,7 +240,8 @@ impl Repl { _ => unknown_command()?, }, None => { - self.ask(line, vec![])?; + let input = Input::from_str(line, self.config.read().input_context()); + self.ask(input)?; } } @@ -259,18 +250,13 @@ impl Repl { Ok(false) } - fn ask(&self, text: &str, files: Vec) -> Result<()> { - if text.is_empty() && files.is_empty() { + fn ask(&self, input: Input) -> Result<()> { + if input.is_empty() { return Ok(()); } while self.config.read().is_compressing_session() { std::thread::sleep(std::time::Duration::from_millis(100)); } - let input = if files.is_empty() { - Input::from_str(text) - } else { - Input::new(text, files)? - }; self.config.read().maybe_print_send_tokens(&input); let mut client = init_client(&self.config)?; ensure_model_capabilities(client.as_mut(), input.required_capabilities())?; @@ -435,7 +421,10 @@ fn parse_command(line: &str) -> Option<(&str, Option<&str>)> { } fn compress_session(config: &GlobalConfig) -> Result<()> { - let input = Input::from_str(&config.read().summarize_prompt); + let input = Input::from_str( + &config.read().summarize_prompt, + config.read().input_context(), + ); let mut client = init_client(config)?; ensure_model_capabilities(client.as_mut(), input.required_capabilities())?; let summary = client.send_message(input)?; -- cgit v1.2.3