Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat: add cli option --agent-variable <name> <value> #1027

Merged
merged 1 commit into from
Dec 1, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion scripts/completions/aichat.bash
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ _aichat() {

case "${cmd}" in
aichat)
opts="-m -r -s -a -R -e -c -f -S -h -V --model --prompt --role --session --empty-session --save-session --agent --rag --serve --execute --code --file --no-stream --dry-run --info --list-models --list-roles --list-sessions --list-agents --list-rags --help --version"
opts="-m -r -s -a -e -c -f -S -h -V --model --prompt --role --session --empty-session --save-session --agent --agent-variable --rag --serve --execute --code --file --no-stream --dry-run --info --list-models --list-roles --list-sessions --list-agents --list-rags --help --version"
if [[ ${cur} == -* || ${cword} -eq 1 ]] ; then
COMPREPLY=( $(compgen -W "${opts}" -- "${cur}") )
return 0
Expand Down
7 changes: 4 additions & 3 deletions scripts/completions/aichat.fish
Original file line number Diff line number Diff line change
@@ -1,11 +1,12 @@
complete -c aichat -s m -l model -x -a "(aichat --list-models)" -d 'Select a LLM model' -r
complete -c aichat -l prompt -d 'Use the system prompt'
complete -c aichat -s r -l role -x -a "(aichat --list-roles)" -d 'Select a role' -r
complete -c aichat -s s -l session -x -a"(aichat --list-sessions)" -d 'Start or join a session' -r
complete -c aichat -s s -l session -x -a "(aichat --list-sessions)" -d 'Start or join a session' -r
complete -c aichat -l empty-session -d 'Ensure the session is empty'
complete -c aichat -l save-session -d 'Ensure the new conversation is saved to the session'
complete -c aichat -s a -l agent -x -a"(aichat --list-agents)" -d 'Start a agent' -r
complete -c aichat -s R -l rag -x -a"(aichat --list-rags)" -d 'Start a RAG' -r
complete -c aichat -s a -l agent -x -a "(aichat --list-agents)" -d 'Start a agent' -r
complete -c aichat -l agent-variable -d 'Set agent variables'
complete -c aichat -l rag -x -a"(aichat --list-rags)" -d 'Start a RAG' -r
complete -c aichat -l serve -d 'Serve the LLM API and WebAPP'
complete -c aichat -s e -l execute -d 'Execute commands in natural language'
complete -c aichat -s c -l code -d 'Output code only'
Expand Down
3 changes: 2 additions & 1 deletion scripts/completions/aichat.nu
Original file line number Diff line number Diff line change
Expand Up @@ -43,7 +43,8 @@ module completions {
--empty-session # Ensure the session is empty
--save-session # Ensure the new conversation is saved to the session
--agent(-a): string@"nu-complete aichat agent" # Start a agent
--rag(-R): string@"nu-complete aichat rag" # Start a RAG
--agent-variable # Set agent variables
--rag: string@"nu-complete aichat rag" # Start a RAG
--serve # Serve the LLM API and WebAPP
--execute(-e) # Execute commands in natural language
--code(-c) # Output code only
Expand Down
2 changes: 1 addition & 1 deletion scripts/completions/aichat.ps1
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ Register-ArgumentCompleter -Native -CommandName 'aichat' -ScriptBlock {
[CompletionResult]::new('--save-session', '--save-session', [CompletionResultType]::ParameterName, 'Ensure the new conversation is saved to the session')
[CompletionResult]::new('-a', '-a', [CompletionResultType]::ParameterName, 'Start a agent')
[CompletionResult]::new('--agent', '--agent', [CompletionResultType]::ParameterName, 'Start a agent')
[CompletionResult]::new('-R', '-R', [CompletionResultType]::ParameterName, 'Start a RAG')
[CompletionResult]::new('--agent-variable', '--agent-variable', [CompletionResultType]::ParameterName, 'Set agent variables')
[CompletionResult]::new('--rag', '--rag', [CompletionResultType]::ParameterName, 'Start a RAG')
[CompletionResult]::new('--serve', '--serve', [CompletionResultType]::ParameterName, 'Serve the LLM API and WebAPP')
[CompletionResult]::new('-e', '-e', [CompletionResultType]::ParameterName, 'Execute commands in natural language')
Expand Down
2 changes: 1 addition & 1 deletion scripts/completions/aichat.zsh
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@ _aichat() {
'--save-session[Ensure the new conversation is saved to the session]' \
'-a[Start a agent]:AGENT:->agents' \
'--agent[Start a agent]:AGENT:->agents' \
'-R[Start a RAG]:RAG:->rags' \
'--agent-variable[Set agent variables]' \
'--rag[Start a RAG]:RAG:->rags' \
'--serve[Serve the LLM API and WebAPP]' \
'-e[Execute commands in natural language]' \
Expand Down
5 changes: 4 additions & 1 deletion src/cli.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,11 @@ pub struct Cli {
/// Start a agent
#[clap(short = 'a', long)]
pub agent: Option<String>,
/// Set agent variables
#[clap(long, value_names = ["NAME", "VALUE"], num_args = 2)]
pub agent_variable: Vec<String>,
/// Start a RAG
#[clap(short = 'R', long)]
#[clap(long)]
pub rag: Option<String>,
/// Serve the LLM API and WebAPP
#[clap(long, value_name = "ADDRESS")]
Expand Down
16 changes: 6 additions & 10 deletions src/config/agent.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,24 +15,17 @@ const DEFAULT_AGENT_NAME: &str = "rag";

pub type AgentVariables = IndexMap<String, String>;

#[derive(Debug, Clone, Serialize)]
#[derive(Debug, Clone)]
pub struct Agent {
name: String,
config: AgentConfig,
definition: AgentDefinition,
#[serde(skip)]
shared_variables: AgentVariables,
#[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>>,
#[serde(skip)]
model: Model,
}

Expand Down Expand Up @@ -75,7 +68,7 @@ impl Agent {

let rag = if rag_path.exists() {
Some(Arc::new(Rag::load(config, DEFAULT_AGENT_NAME, &rag_path)?))
} else if !definition.documents.is_empty() && !config.read().print_info_only {
} else if !definition.documents.is_empty() && !config.read().cli_info_flag {
let mut ans = false;
if *IS_STDOUT_TERMINAL {
ans = Confirm::new("The agent has the documents, init RAG?")
Expand Down Expand Up @@ -182,11 +175,14 @@ impl Agent {
pub fn export(&self) -> Result<String> {
let mut agent = self.clone();
agent.definition.instructions = self.interpolated_instructions();
let mut value = serde_json::json!(agent);
let mut value = json!({});
value["name"] = json!(self.name());
let variables = self.variables();
if !variables.is_empty() {
value["variables"] = serde_json::to_value(variables)?;
}
value["config"] = json!(self.config);
value["definition"] = json!(self.definition);
value["functions_dir"] = Config::agent_functions_dir(&self.name)
.display()
.to_string()
Expand Down
36 changes: 27 additions & 9 deletions src/config/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -156,9 +156,12 @@ pub struct Config {
#[serde(skip)]
pub working_mode: WorkingMode,
#[serde(skip)]
pub print_info_only: bool,
#[serde(skip)]
pub last_message: Option<(Input, String)>,

#[serde(skip)]
pub cli_info_flag: bool,
#[serde(skip)]
pub cli_agent_variables: Option<AgentVariables>,
}

impl Default for Config {
Expand Down Expand Up @@ -218,8 +221,10 @@ impl Default for Config {
model: Default::default(),
functions: Default::default(),
working_mode: WorkingMode::Cmd,
print_info_only: false,
last_message: None,

cli_info_flag: false,
cli_agent_variables: None,
}
}
}
Expand Down Expand Up @@ -1508,6 +1513,7 @@ impl Config {
if self.agent.take().is_some() {
self.rag.take();
self.last_message = None;
self.cli_agent_variables = None;
}
Ok(())
}
Expand Down Expand Up @@ -1997,14 +2003,20 @@ impl Config {
None => return Ok(()),
};
if !agent.defined_variables().is_empty() && agent.shared_variables().is_empty() {
let mut config_variables = agent.config_variables().clone();
if let Some(v) = &self.cli_agent_variables {
config_variables.extend(v.clone());
}
let new_variables = Agent::init_agent_variables(
agent.defined_variables(),
agent.config_variables(),
self.print_info_only,
&config_variables,
self.cli_info_flag,
)?;
agent.set_shared_variables(new_variables);
}
agent.update_shared_dynamic_instructions(false)?;
if !self.cli_info_flag {
agent.update_shared_dynamic_instructions(false)?;
}
Ok(())
}

Expand All @@ -2017,18 +2029,24 @@ impl Config {
let shared_variables = agent.shared_variables().clone();
let session_variables =
if !agent.defined_variables().is_empty() && shared_variables.is_empty() {
let mut config_variables = agent.config_variables().clone();
if let Some(v) = &self.cli_agent_variables {
config_variables.extend(v.clone());
}
let new_variables = Agent::init_agent_variables(
agent.defined_variables(),
agent.config_variables(),
self.print_info_only,
&config_variables,
self.cli_info_flag,
)?;
agent.set_shared_variables(new_variables.clone());
new_variables
} else {
shared_variables
};
agent.set_session_variables(session_variables);
agent.update_session_dynamic_instructions(None)?;
if !self.cli_info_flag {
agent.update_session_dynamic_instructions(None)?;
}
session.sync_agent(agent);
} else {
let variables = session.agent_variables();
Expand Down
11 changes: 10 additions & 1 deletion src/main.rs
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ async fn run(config: GlobalConfig, cli: Cli, text: Option<String>) -> Result<()>
return serve::run(config, addr).await;
}
if cli.info {
config.write().print_info_only = true;
config.write().cli_info_flag = true;
}

if cli.list_models {
Expand Down Expand Up @@ -98,6 +98,15 @@ async fn run(config: GlobalConfig, cli: Cli, text: Option<String>) -> Result<()>
Some(v) => v.as_str(),
None => TEMP_SESSION_NAME,
});
if !cli.agent_variable.is_empty() {
config.write().cli_agent_variables = Some(
cli.agent_variable
.chunks(2)
.map(|v| (v[0].to_string(), v[1].to_string()))
.collect(),
);
}

Config::use_agent(&config, agent, session, abort_signal.clone()).await?
} else {
if let Some(prompt) = &cli.prompt {
Expand Down