feat: agent supports dynamic instructions (#1023)

* feat: agent supports dynamic instructions

* change tool calls' null output to 'TODO'

* REPL don't print banner if use agent/rag

3 files changed, 148 insertions, 46 deletions

diff --git a/src/config/agent.rs b/src/config/agent.rs
index 29121ee..87ac719 100644
--- a/src/config/agent.rs
+++ b/src/config/agent.rs
@@ -1,6 +1,9 @@
 use super::*;
 
-use crate::{client::Model, function::Functions};
+use crate::{
+    client::Model,
+    function::{run_llm_function, Functions},
+};
 
 use anyhow::{Context, Result};
 use inquire::{validator::Validation, Text};
@@ -22,6 +25,10 @@ pub struct Agent {
     #[serde(skip)]
     session_variables: Option<AgentVariables>,
     #[serde(skip)]
+    shared_dynamic_instructions: Option<String>,
+    #[serde(skip)]
+    session_dynamic_instructions: Option<String>,
+    #[serde(skip)]
     functions: Functions,
     #[serde(skip)]
     rag: Option<Arc<Rag>>,
@@ -102,6 +109,8 @@ impl Agent {
             definition,
             shared_variables: Default::default(),
             session_variables: None,
+            shared_dynamic_instructions: None,
+            session_dynamic_instructions: None,
             functions,
             rag,
             model,
@@ -215,7 +224,16 @@ impl Agent {
     }
 
     pub fn interpolated_instructions(&self) -> String {
-        self.definition.interpolated_instructions(self.variables())
+        let mut output = self
+            .session_dynamic_instructions
+            .clone()
+            .or_else(|| self.shared_dynamic_instructions.clone())
+            .unwrap_or_else(|| self.definition.instructions.clone());
+        for (k, v) in self.variables() {
+            output = output.replace(&format!("{{{{{k}}}}}"), v)
+        }
+        interpolate_variables(&mut output);
+        output
     }
 
     pub fn agent_prelude(&self) -> Option<&str> {
@@ -257,11 +275,8 @@ impl Agent {
         self.shared_variables = shared_variables;
     }
 
-    pub fn set_session_variables(&mut self, session_variables: Option<AgentVariables>) {
-        if self.shared_variables.is_empty() {
-            self.shared_variables = session_variables.clone().unwrap_or_default();
-        }
-        self.session_variables = session_variables;
+    pub fn set_session_variables(&mut self, session_variables: AgentVariables) {
+        self.session_variables = Some(session_variables);
     }
 
     pub fn set_variable(&mut self, key: &str, value: &str) -> Result<()> {
@@ -269,16 +284,67 @@ impl Agent {
             Some(v) => v,
             None => &mut self.shared_variables,
         };
-        if !variables.contains_key(key) {
+        let Some(old_value) = variables.get(key) else {
             bail!("Unknown variable: '{key}'")
+        };
+        if old_value == value {
+            return Ok(());
         }
         variables.insert(key.to_string(), value.to_string());
+        if self.session_variables.is_some() {
+            self.update_session_dynamic_instructions(None)?;
+        } else {
+            self.update_shared_dynamic_instructions(true)?;
+        }
+
         Ok(())
     }
 
     pub fn defined_variables(&self) -> &[AgentVariable] {
         &self.definition.variables
     }
+
+    pub fn exit_session(&mut self) {
+        self.session_variables = None;
+        self.session_dynamic_instructions = None;
+    }
+
+    pub fn is_dynamic_instructions(&self) -> bool {
+        self.definition.dynamic_instructions
+    }
+
+    pub fn update_shared_dynamic_instructions(&mut self, force: bool) -> Result<()> {
+        if self.is_dynamic_instructions() && (force || self.shared_dynamic_instructions.is_none()) {
+            self.shared_dynamic_instructions = Some(self.run_instructions_fn()?);
+        }
+        Ok(())
+    }
+
+    pub fn update_session_dynamic_instructions(&mut self, value: Option<String>) -> Result<()> {
+        if self.is_dynamic_instructions() {
+            let value = match value {
+                Some(v) => v,
+                None => self.run_instructions_fn()?,
+            };
+            self.session_dynamic_instructions = Some(value);
+        }
+        Ok(())
+    }
+
+    fn run_instructions_fn(&self) -> Result<String> {
+        let value = run_llm_function(
+            self.name().to_string(),
+            vec!["_instructions".into(), "{}".into()],
+            self.variable_envs(),
+        )?;
+        match value {
+            Some(v) => {
+                println!();
+                Ok(v)
+            }
+            _ => bail!("No return value from '_instructions' function"),
+        }
+    }
 }
 
 impl RoleLike for Agent {
@@ -331,11 +397,15 @@ impl RoleLike for Agent {
 pub struct AgentConfig {
     #[serde(rename(serialize = "model", deserialize = "model"))]
     pub model_id: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub temperature: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub top_p: Option<f64>,
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub use_tools: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
     pub agent_prelude: Option<String>,
-    #[serde(default)]
+    #[serde(default, skip_serializing_if = "IndexMap::is_empty")]
     pub variables: AgentVariables,
 }
 
@@ -389,8 +459,11 @@ pub struct AgentDefinition {
     pub description: String,
     #[serde(default)]
     pub version: String,
+    #[serde(default)]
     pub instructions: String,
     #[serde(default)]
+    pub dynamic_instructions: bool,
+    #[serde(default)]
     pub variables: Vec<AgentVariable>,
     #[serde(default)]
     pub conversation_starters: Vec<String>,
@@ -436,15 +509,6 @@ impl AgentDefinition {
         )
     }
 
-    fn interpolated_instructions(&self, variables: &AgentVariables) -> String {
-        let mut output = self.instructions.clone();
-        for (k, v) in variables {
-            output = output.replace(&format!("{{{{{k}}}}}"), v)
-        }
-        interpolate_variables(&mut output);
-        output
-    }
-
     fn replace_tools_placeholder(&mut self, functions: &Functions) {
         let tools_placeholder: &str = "{{__tools__}}";
         if self.instructions.contains(tools_placeholder) {
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 31337cd..39c63a3 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -1076,9 +1076,6 @@ impl Config {
             session.exit(&sessions_dir, self.working_mode.is_repl())?;
             self.last_message = None;
         }
-        if let Some(agent) = self.agent.as_mut() {
-            agent.set_session_variables(None);
-        }
         Ok(())
     }
 
@@ -1122,7 +1119,7 @@ impl Config {
     pub fn empty_session(&mut self) -> Result<()> {
         if let Some(session) = self.session.as_mut() {
             if let Some(agent) = self.agent.as_ref() {
-                session.sync_agent(agent, false);
+                session.sync_agent(agent);
             }
             session.clear_messages();
         } else {
@@ -1496,7 +1493,7 @@ impl Config {
                 let value = parts[1];
                 agent.set_variable(key, value)?;
                 if let Some(session) = self.session.as_mut() {
-                    session.sync_agent(agent, true);
+                    session.sync_agent(agent);
                 }
             }
             None => bail!("No agent"),
@@ -1513,6 +1510,17 @@ impl Config {
         Ok(())
     }
 
+    pub fn exit_agent_session(&mut self) -> Result<()> {
+        self.exit_session()?;
+        if let Some(agent) = self.agent.as_mut() {
+            agent.exit_session();
+            if self.working_mode.is_repl() {
+                self.init_agent_shared_variables()?;
+            }
+        }
+        Ok(())
+    }
+
     pub fn apply_prelude(&mut self) -> Result<()> {
         if !self.state().is_empty() {
             return Ok(());
@@ -1986,12 +1994,15 @@ impl Config {
             Some(v) => v,
             None => return Ok(()),
         };
-        let new_variables = Agent::init_agent_variables(
-            agent.defined_variables(),
-            agent.config_variables(),
-            self.print_info_only,
-        )?;
-        agent.set_shared_variables(new_variables);
+        if !agent.defined_variables().is_empty() && agent.shared_variables().is_empty() {
+            let new_variables = Agent::init_agent_variables(
+                agent.defined_variables(),
+                agent.config_variables(),
+                self.print_info_only,
+            )?;
+            agent.set_shared_variables(new_variables);
+        }
+        agent.update_shared_dynamic_instructions(false)?;
         Ok(())
     }
 
@@ -2000,20 +2011,30 @@ impl Config {
             (Some(agent), Some(session)) => (agent, session),
             _ => return Ok(()),
         };
-        let shared_variables = agent.shared_variables();
-        let mut all_variables = if shared_variables.is_empty() {
-            agent.config_variables().clone()
+        if session.is_empty() {
+            let shared_variables = agent.shared_variables().clone();
+            let session_variables =
+                if !agent.defined_variables().is_empty() && shared_variables.is_empty() {
+                    let new_variables = Agent::init_agent_variables(
+                        agent.defined_variables(),
+                        agent.config_variables(),
+                        self.print_info_only,
+                    )?;
+                    agent.set_shared_variables(new_variables.clone());
+                    new_variables
+                } else {
+                    shared_variables
+                };
+            agent.set_session_variables(session_variables);
+            agent.update_session_dynamic_instructions(None)?;
+            session.sync_agent(agent);
         } else {
-            shared_variables.clone()
-        };
-        all_variables.extend(session.agent_variables().clone());
-        let new_variables = Agent::init_agent_variables(
-            agent.defined_variables(),
-            &all_variables,
-            self.print_info_only,
-        )?;
-        agent.set_session_variables(Some(new_variables));
-        session.sync_agent(agent, false);
+            let variables = session.agent_variables();
+            agent.set_session_variables(variables.clone());
+            agent.update_session_dynamic_instructions(Some(
+                session.agent_instructions().to_string(),
+            ))?;
+        }
         Ok(())
     }
 
@@ -2306,9 +2327,22 @@ impl AssertState {
     pub fn pass() -> Self {
         AssertState::False(StateFlags::empty())
     }
+
     pub fn bare() -> Self {
         AssertState::Equal(StateFlags::empty())
     }
+
+    pub fn assert(self, flags: StateFlags) -> bool {
+        match self {
+            AssertState::True(true_flags) => true_flags & flags != StateFlags::empty(),
+            AssertState::False(false_flags) => false_flags & flags == StateFlags::empty(),
+            AssertState::TrueFalse(true_flags, false_flags) => {
+                (true_flags & flags != StateFlags::empty())
+                    && (false_flags & flags == StateFlags::empty())
+            }
+            AssertState::Equal(check_flags) => check_flags == flags,
+        }
+    }
 }
 
 fn create_config_file(config_path: &Path) -> Result<()> {
diff --git a/src/config/session.rs b/src/config/session.rs
index 4b3da76..820fa12 100644
--- a/src/config/session.rs
+++ b/src/config/session.rs
@@ -36,6 +36,8 @@ pub struct Session {
     role_name: Option<String>,
     #[serde(default, skip_serializing_if = "IndexMap::is_empty")]
     agent_variables: AgentVariables,
+    #[serde(default, skip_serializing_if = "String::is_empty")]
+    agent_instructions: String,
 
     #[serde(default, skip_serializing_if = "Vec::is_empty")]
     compressed_messages: Vec<Message>,
@@ -275,19 +277,21 @@ impl Session {
         self.role_prompt.clear();
     }
 
-    pub fn sync_agent(&mut self, agent: &Agent, set_dirty: bool) {
+    pub fn sync_agent(&mut self, agent: &Agent) {
         self.role_name = None;
         self.role_prompt = agent.interpolated_instructions();
         self.agent_variables = agent.variables().clone();
-        if set_dirty {
-            self.dirty = true;
-        }
+        self.agent_instructions = self.role_prompt.clone();
     }
 
     pub fn agent_variables(&self) -> &AgentVariables {
         &self.agent_variables
     }
 
+    pub fn agent_instructions(&self) -> &str {
+        &self.agent_instructions
+    }
+
     pub fn set_save_session(&mut self, value: Option<bool>) {
         if self.save_session != value {
             self.save_session = value;