diff --git a/assets/agents/coder/tools.sh b/assets/agents/coder/tools.sh index 407ff29..db61afc 100755 --- a/assets/agents/coder/tools.sh +++ b/assets/agents/coder/tools.sh @@ -3,9 +3,8 @@ set -e # @env LLM_OUTPUT=/dev/stdout The output path -PROMPT_UTILS="${LLM_ROOT_DIR:-$(dirname "${BASH_SOURCE[0]}")/..}/functions/utils/prompt-utils.sh" # shellcheck disable=SC1090 -source "$PROMPT_UTILS" +source "$LLM_PROMPT_UTILS_FILE" # @cmd Create a new file at the specified path with the given contents. # @option --path! The path where the file should be created diff --git a/assets/agents/sql/tools.sh b/assets/agents/sql/tools.sh index 20e5bf2..1cb0379 100755 --- a/assets/agents/sql/tools.sh +++ b/assets/agents/sql/tools.sh @@ -6,9 +6,8 @@ set -e # @env LLM_OUTPUT=/dev/stdout The output path # @env LLM_AGENT_VAR_DSN! The database connection url. e.g. pgsql://user:pass@host:port -PROMPT_UTILS="${LLM_ROOT_DIR:-$(dirname "${BASH_SOURCE[0]}")/..}/functions/utils/prompt-utils.sh" # shellcheck disable=SC1090 -source "$PROMPT_UTILS" +source "$LLM_PROMPT_UTILS_FILE" # @cmd Execute a SELECT query # @option --query! SELECT SQL query to execute diff --git a/assets/functions/scripts/run-agent.sh b/assets/functions/scripts/run-agent.sh index fac28af..5dd5179 100644 --- a/assets/functions/scripts/run-agent.sh +++ b/assets/functions/scripts/run-agent.sh @@ -27,6 +27,7 @@ setup_env() { export LLM_AGENT_FUNC="$agent_func" export LLM_AGENT_ROOT_DIR="$LLM_ROOT_DIR/agents/{agent_name}" export LLM_AGENT_CACHE_DIR="$LLM_ROOT_DIR/cache/{agent_name}" + export LLM_PROMPT_UTILS_FILE="{prompt_utils_file}" } load_env() { diff --git a/assets/functions/scripts/run-tool.py b/assets/functions/scripts/run-tool.py index b7098a3..d8c90b1 100644 --- a/assets/functions/scripts/run-tool.py +++ b/assets/functions/scripts/run-tool.py @@ -31,10 +31,10 @@ def main(): raw_data = parse_argv() tool_data = parse_raw_data(raw_data) - root_dir = "{config_dir}/functions" + root_dir = "{root_dir}" setup_env(root_dir) - tool_path = os.path.join(root_dir, "tools/{function_name}.py") + tool_path = "{tool_path}.py" run(tool_path, "run", tool_data) @@ -150,4 +150,4 @@ def dump_result(name): if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/assets/functions/scripts/run-tool.sh b/assets/functions/scripts/run-tool.sh index 02b9bb6..ead83f6 100644 --- a/assets/functions/scripts/run-tool.sh +++ b/assets/functions/scripts/run-tool.sh @@ -5,10 +5,10 @@ set -e main() { - root_dir="{config_dir}/functions" + root_dir="{root_dir}" parse_argv "$@" setup_env - tool_path="$root_dir/tools/{function_name}.sh" + tool_path="{tool_path}.sh" run } @@ -24,6 +24,7 @@ setup_env() { export LLM_ROOT_DIR="$root_dir" export LLM_TOOL_NAME="{function_name}" export LLM_TOOL_CACHE_DIR="$LLM_ROOT_DIR/cache/{function_name}" + export LLM_PROMPT_UTILS_FILE="{prompt_utils_file}" } load_env() { diff --git a/assets/functions/tools/execute_command.sh b/assets/functions/tools/execute_command.sh index 1850f1a..d96da5b 100755 --- a/assets/functions/tools/execute_command.sh +++ b/assets/functions/tools/execute_command.sh @@ -6,9 +6,8 @@ set -e # @env LLM_OUTPUT=/dev/stdout The output path -PROMPT_UTILS="${LLM_ROOT_DIR:-$(dirname "${BASH_SOURCE[0]}")/..}/functions/utils/prompt-utils.sh" # shellcheck disable=SC1090 -source "$PROMPT_UTILS" +source "$LLM_PROMPT_UTILS_FILE" main() { guard_operation diff --git a/assets/functions/tools/execute_sql_code.sh b/assets/functions/tools/execute_sql_code.sh index 02b687f..68753a1 100755 --- a/assets/functions/tools/execute_sql_code.sh +++ b/assets/functions/tools/execute_sql_code.sh @@ -9,9 +9,8 @@ set -e # @env USQL_DSN! The database connection url. e.g. pgsql://user:pass@host:port # @env LLM_OUTPUT=/dev/stdout The output path -PROMPT_UTILS="${LLM_ROOT_DIR:-$(dirname "${BASH_SOURCE[0]}")/..}/functions/utils/prompt-utils.sh" # shellcheck disable=SC1090 -source "$PROMPT_UTILS" +source "$LLM_PROMPT_UTILS_FILE" # shellcheck disable=SC2154 main() { diff --git a/assets/functions/tools/fs_patch.sh b/assets/functions/tools/fs_patch.sh index d973827..083119a 100755 --- a/assets/functions/tools/fs_patch.sh +++ b/assets/functions/tools/fs_patch.sh @@ -9,9 +9,8 @@ set -e # @env LLM_OUTPUT=/dev/stdout The output path -PROMPT_UTILS="${LLM_ROOT_DIR:-$(dirname "${BASH_SOURCE[0]}")/..}/functions/utils/prompt-utils.sh" # shellcheck disable=SC1090 -source "$PROMPT_UTILS" +source "$LLM_PROMPT_UTILS_FILE" # shellcheck disable=SC2154 main() { diff --git a/assets/functions/tools/fs_rm.sh b/assets/functions/tools/fs_rm.sh index d60149d..ac27c1a 100755 --- a/assets/functions/tools/fs_rm.sh +++ b/assets/functions/tools/fs_rm.sh @@ -7,9 +7,8 @@ set -e # @env LLM_OUTPUT=/dev/stdout The output path -PROMPT_UTILS="${LLM_ROOT_DIR:-$(dirname "${BASH_SOURCE[0]}")/..}/functions/utils/prompt-utils.sh" # shellcheck disable=SC1090 -source "$PROMPT_UTILS" +source "$LLM_PROMPT_UTILS_FILE" # shellcheck disable=SC2154 main() { diff --git a/assets/functions/tools/fs_write.sh b/assets/functions/tools/fs_write.sh index 8511702..235126f 100755 --- a/assets/functions/tools/fs_write.sh +++ b/assets/functions/tools/fs_write.sh @@ -8,9 +8,8 @@ set -e # @env LLM_OUTPUT=/dev/stdout The output path -PROMPT_UTILS="${LLM_ROOT_DIR:-$(dirname "${BASH_SOURCE[0]}")/..}/functions/utils/prompt-utils.sh" # shellcheck disable=SC1090 -source "$PROMPT_UTILS" +source "$LLM_PROMPT_UTILS_FILE" # shellcheck disable=SC2154 main() { diff --git a/docs/SHELL-INTEGRATIONS.md b/docs/SHELL-INTEGRATIONS.md new file mode 100644 index 0000000..59819ac --- /dev/null +++ b/docs/SHELL-INTEGRATIONS.md @@ -0,0 +1,98 @@ +# Loki Shell Integrations +Loki supports the following integrations with a handful of shell environments to enhance user experience and streamline workflows. + +## Tab Completions +### Dynamic +Dynamic tab completions are supported by Loki to assist users in quickly completing commands, options, and arguments. +You can enable it by using the corresponding command for your shell. To enable dynamic tab completions for every +shell session (i.e. persistently), add the corresponding command to your shell's configuration file as indicated: + +```shell +# Bash +# (add to: `~/.bashrc`) +source <(COMPLETE=bash loki) + +# Zsh +# (add to: `~/.zshrc`) +source <(COMPLETE=zsh loki) + +# Fish +# (add to: `~/.config/fish/config.fish`) +source <(COMPLETE=fish loki | psub) + +# Elvish +# (add to: `~/.elvish/rc.elv`) +eval (E:COMPLETE=elvish loki | slurp) + +# PowerShell +# (add to: `$PROFILE`) +$env:COMPLETE = "powershell" +loki | Out-String | Invoke-Expression +``` + +At the time of writing, `nushell` is not yet fully supported for dynamic tab completions due to limitations +in the [`clap`](https://crates.io/crates/clap) crate. However, `nushell` support is being actively developed, and will +be added in a future release. + +Progress on this feature can be tracked in the following issue: [Clap Issue #5840](https://github.com/clap-rs/clap/issues/5840). + +### Static +Static tab completions (i.e. pre-generated completion scripts that are not context aware) can also be generated using the +`--completions` flag. You can enable static tab completions by using the corresponding commands for your shell. These commands +will enable them for every shell session (i.e. persistently): + +```shell +# Bash +echo 'source <(loki --completions bash)' >> ~/.bashrc + +# Zsh +echo 'source <(loki --completions zsh)' >> ~/.zshrc + +# Fish +echo 'loki --completions fish | source' >> ~/.config/fish/config.fish + +# Elvish +echo 'eval (loki --completions elvish | slurp)' >> ~/.elvish/rc.elv + +# Nushell +[[ -d ~/.config/nushell/completions ]] || mkdir -p ~/.config/nushell/completions +loki --completions nushell | save -f ~/.config/nushell/completions/loki.nu +echo 'use ~/.config/nushell/completions/cli.nu *' >> ~/.config/nushell/config.nu + +# PowerShell +Add-content $PROFILE "loki --completions powershell | Out-String | Invoke-Expression" +``` + +## Shell Assistant +Loki has an `-e,--execute` flag that allows users to run natural language commands directly from the CLI. It accepts +natural language input and translates it into executable shell commands. + +![Shell Assistant Demo](./images/shell_integrations/assistant.gif) + +## Intelligent Command Completions +Loki also provides shell scripts that bind `Alt-e` to `loki -e ""`, allowing users to generate +commands from natural text directly without invoking the CLI. + +For example: + +```shell +$ find all typescript files with more than 100 lines +# Gets replaced with +$ find . -name '*.ts' -type f -exec awk 'NR>100{exit 1}' {} \; -print +``` + +To use the CLI helper, add the content of the appropriate integration script for your shell to your shell configuration file: +* [Bash Integration](../scripts/shell-integration/bash-integration.sh) (add to: `~/.bashrc`) +* [Zsh Integration](../scripts/shell-integration/zsh-integration.zsh) (add to: `~/.zshrc`) +* [Elvish Integration](../scripts/shell-integration/elvish-integration.elv) (add to: `~/.elvish/rc.elv`) +* [Fish Integration](../scripts/shell-integration/fish-integration.fish) (add to: `~/.config/fish/config.fish`) +* [Nushell Integration](../scripts/shell-integration/nushell-integration.nu) (add to: `~/.config/nushell/config.nu`) +* [PowerShell Integration](../scripts/shell-integration/powershell-integration.ps1) (add to: `$PROFILE`) + +## Code Generation +Users can also directly generate code snippets from natural language prompts using the `-c,--code` flag. + +![Code Generation Demo](./images/shell_integrations/code-generation.gif) + +**Pro Tip:** Pipe the output of the code generation directly into `tee` to ensure the generated code is properly extracted +from any generated Markdown (i.e. remove any triple backticks). diff --git a/docs/VAULT.md b/docs/VAULT.md new file mode 100644 index 0000000..84b9838 --- /dev/null +++ b/docs/VAULT.md @@ -0,0 +1,111 @@ +# The Loki Vault +The Loki vault lets users store sensitive secrets and credentials securely so that there's no plaintext secrets +anywhere in your configurations. + +It's based on the [G-Man library](https://github.com/Dark-Alex-17/gman) (which also comes in a binary format) which +functions as a universal secret management tool. + +![Vault Demo](./images/vault/vault-demo.gif) + +## Usage +The Loki vault can be used in one of two ways: via the CLI or via the REPL for interactive usage. + +### CLI Usage +The vault is utilized from the CLI with the following flags: + +```bash +--add-secret Add a secret to the Loki vault +--get-secret Decrypt a secret from the Loki vault and print the plaintext +--update-secret Update an existing secret in the Loki vault +--delete-secret Delete a secret from the Loki vault +--list-secrets List all secrets stored in the Loki vault +``` +(The above is also documented in `loki --help`) + +Loki will guide you through manipulating your secrets to make usage easier. + +### REPL Usage +The vault can be access from within the Loki REPL using the `.vault` commands: + +![Loki Vault REPL](./images/vault/vault-repl.png) +![Loki Vault REPL Commands](./images/vault/vault-repl-commands.png) + +The manipulation of your vault is guided in the same way as the CLI usage, ensuring ease of use. + +## Motivation +Loki is intended to be highly configurable and adaptable to many different use cases. This means that users of Loki +should be able to share configurations for agents, tools, roles, etc. with other users or even entire teams. + +My objective is to encourage this, and to make it so that users can easily version their configurations using version +control. Good VCS hygiene dictates that one *never* commits secrets or sensitive information to a repository. + +Since a number of files and configurations in Loki may contain sensitive information, the vault exists to solve this problem. + +Users can either share the vault password with a team, making it so a single configuration can be pulled from VCS and used +by said team. Alternatively, each user can maintain their own vault password and expect other users to replace secret values +with their user-specific secrets. + +## How it works +When you first start Loki, if you don't already have a vault password file, it will prompt you to create one. This file +houses the password that is used to encrypt and decrypt secrets within Loki. This file exists so that you are not prompted +for a password every time Loki attempts to decrypt a secret. + +When you encrypt a secret, it uses the local provider for `gman` to securely store those secrets in the Loki vault file. +This file is typically located at your Loki configuration directory under `vault.yml`. If you open this file, you'll see a +bunch of gibberish. This is because all secrets are encrypted using the password you provided, meaning only you can decrypt them. + +Secrets are specified in Loki configurations using the same variable templating as the [Jinja templating engine](https://jinja.palletsprojects.com/en/stable/): + +``` +{{some_variable}} +``` + +So whenever you want Loki to use a secret from the vault, you simply specify the secret name in this format in the applicable +file. + +**Example:** +Suppose my vault has a secret called `GITHUB_TOKEN` in it, and I want to use that in the MCP configuration. Then, I simply replace +the expected value in my `mcp.json` with the templated secret: + +```json +{ + "mcpServers": { + "atlassian": { + "command": "npx", + "args": ["-y", "mcp-remote", "https://mcp.atlassian.com/v1/sse"] + }, + "github": { + "command": "docker", + "args": [ + "run", + "-i", + "--rm", + "-e", + "GITHUB_PERSONAL_ACCESS_TOKEN", + "ghcr.io/github/github-mcp-server" + ], + "env": { + "GITHUB_PERSONAL_ACCESS_TOKEN": "{{GITHUB_TOKEN}}" + } + } + } +} +``` + +At runtime, Loki will detect the templated secret and replace it with the decrypted value from the vault before executing. + +## Supported Files +At the time of writing, the following files support Loki secret injection: + +| File Type | Description | Limitations | +|----------------------|-----------------------------------|----------------------------------------------------------------| +| `config.yaml` | The main Loki configuration file | Cannot use secret injection on the `vault_password_file` field | +| `functions/mcp.json` | The MCP server configuration file | | + + +Note that all paths are relative to the Loki configuration directory. The directory varies by system, so you can find yours by +running + +```shell +dirname $(loki --info | grep config_file | awk '{print $2}') +``` diff --git a/docs/images/shell_integrations/assistant.gif b/docs/images/shell_integrations/assistant.gif new file mode 100644 index 0000000..511730d Binary files /dev/null and b/docs/images/shell_integrations/assistant.gif differ diff --git a/docs/images/shell_integrations/code-generation.gif b/docs/images/shell_integrations/code-generation.gif new file mode 100644 index 0000000..a757bff Binary files /dev/null and b/docs/images/shell_integrations/code-generation.gif differ diff --git a/docs/images/vault/vault-demo.gif b/docs/images/vault/vault-demo.gif new file mode 100644 index 0000000..d3ce40a Binary files /dev/null and b/docs/images/vault/vault-demo.gif differ diff --git a/docs/images/vault/vault-repl-commands.png b/docs/images/vault/vault-repl-commands.png new file mode 100644 index 0000000..2fbed02 Binary files /dev/null and b/docs/images/vault/vault-repl-commands.png differ diff --git a/docs/images/vault/vault-repl.png b/docs/images/vault/vault-repl.png new file mode 100644 index 0000000..9c7d6aa Binary files /dev/null and b/docs/images/vault/vault-repl.png differ diff --git a/src/config/agent.rs b/src/config/agent.rs index 1c7f206..d7b4bb7 100644 --- a/src/config/agent.rs +++ b/src/config/agent.rs @@ -389,6 +389,7 @@ impl Agent { self.name().to_string(), vec!["_instructions".into(), "{}".into()], self.variable_envs(), + Some(self.name().to_string()), )?; match value { Some(v) => Ok(v), diff --git a/src/config/mod.rs b/src/config/mod.rs index 12d400f..4bb16bd 100644 --- a/src/config/mod.rs +++ b/src/config/mod.rs @@ -74,6 +74,8 @@ const FUNCTIONS_BIN_DIR_NAME: &str = "bin"; const AGENTS_DIR_NAME: &str = "agents"; const GLOBAL_TOOLS_DIR_NAME: &str = "tools"; const GLOBAL_TOOLS_FILE_NAME: &str = "tools.txt"; +const GLOBAL_TOOLS_UTILS_DIR_NAME: &str = "utils"; +const BASH_PROMPT_UTILS_FILE_NAME: &str = "prompt-utils.sh"; const MCP_FILE_NAME: &str = "mcp.json"; const CLIENTS_FIELD: &str = "clients"; @@ -492,6 +494,14 @@ impl Config { Self::functions_dir().join(GLOBAL_TOOLS_DIR_NAME) } + pub fn global_utils_dir() -> PathBuf { + Self::functions_dir().join(GLOBAL_TOOLS_UTILS_DIR_NAME) + } + + pub fn bash_prompt_utils_file() -> PathBuf { + Self::global_utils_dir().join(BASH_PROMPT_UTILS_FILE_NAME) + } + pub fn session_file(&self, name: &str) -> PathBuf { match name.split_once("/") { Some((dir, name)) => self.sessions_dir().join(dir).join(format!("{name}.yaml")), diff --git a/src/function.rs b/src/function/mod.rs similarity index 87% rename from src/function.rs rename to src/function/mod.rs index f3142ed..e9fd261 100644 --- a/src/function.rs +++ b/src/function/mod.rs @@ -32,8 +32,8 @@ const PATH_SEP: &str = ";"; const PATH_SEP: &str = ":"; #[derive(AsRefStr)] -enum BinaryType { - Tool, +enum BinaryType<'a> { + Tool(Option<&'a str>), Agent, } @@ -169,6 +169,7 @@ impl Functions { pub fn init() -> Result { Self::install_global_tools()?; + Self::clear_global_functions_bin_dir()?; info!( "Initializing global functions from {}", Config::global_tools_file().display() @@ -191,6 +192,8 @@ impl Functions { pub fn init_agent(name: &str, global_tools: &[String]) -> Result { Self::install_global_tools()?; + Self::clear_agent_bin_dir(name)?; + let global_tools_declarations = if !global_tools.is_empty() { let enabled_tools = global_tools.join("\n"); info!("Loading global tools for agent: {name}: {enabled_tools}"); @@ -200,7 +203,7 @@ impl Functions { "Building global function binaries required by agent: {name} in {}", Config::functions_bin_dir().display() ); - Self::build_global_function_binaries(&enabled_tools)?; + Self::build_global_function_binaries(&enabled_tools, Some(name))?; tools_declarations } else { debug!("No global tools found for agent: {}", name); @@ -381,17 +384,7 @@ impl Functions { } } - fn build_global_function_binaries(enabled_tools: &str) -> Result<()> { - let bin_dir = Config::functions_bin_dir(); - if !bin_dir.exists() { - fs::create_dir_all(&bin_dir)?; - } - info!( - "Clearing existing function binaries in {}", - bin_dir.display() - ); - clear_dir(&bin_dir)?; - + fn build_global_function_binaries(enabled_tools: &str, agent_name: Option<&str>) -> Result<()> { for line in enabled_tools.lines() { if line.starts_with('#') { continue; @@ -417,7 +410,7 @@ impl Functions { bail!("Unsupported tool file extension: {}", language.as_ref()); } - Self::build_binaries(binary_name, language, BinaryType::Tool)?; + Self::build_binaries(binary_name, language, BinaryType::Tool(agent_name))?; } Ok(()) @@ -427,10 +420,10 @@ impl Functions { let enabled_tools = fs::read_to_string(&tools_txt_path) .with_context(|| format!("failed to load functions at {}", tools_txt_path.display()))?; - Self::build_global_function_binaries(&enabled_tools) + Self::build_global_function_binaries(&enabled_tools, None) } - fn build_agent_tool_binaries(name: &str) -> Result<()> { + fn clear_agent_bin_dir(name: &str) -> Result<()> { let agent_bin_directory = Config::agent_bin_dir(name); if !agent_bin_directory.exists() { debug!( @@ -446,6 +439,25 @@ impl Functions { clear_dir(&agent_bin_directory)?; } + Ok(()) + } + + fn clear_global_functions_bin_dir() -> Result<()> { + let bin_dir = Config::functions_bin_dir(); + if !bin_dir.exists() { + fs::create_dir_all(&bin_dir)?; + } + + info!( + "Clearing existing function binaries in {}", + bin_dir.display() + ); + clear_dir(&bin_dir)?; + + Ok(()) + } + + fn build_agent_tool_binaries(name: &str) -> Result<()> { let language = Language::from( &Config::agent_functions_file(name)? .extension() @@ -471,11 +483,16 @@ impl Functions { ) -> Result<()> { use native::runtime; let (binary_file, binary_script_file) = match binary_type { - BinaryType::Tool => ( + BinaryType::Tool(None) => ( Config::functions_bin_dir().join(format!("{binary_name}.cmd")), Config::functions_bin_dir() .join(format!("run-{binary_name}.{}", language.to_extension())), ), + BinaryType::Tool(Some(agent_name)) => ( + Config::agent_bin_dir(agent_name).join(format!("{binary_name}.cmd")), + Config::agent_bin_dir(agent_name) + .join(format!("run-{binary_name}.{}", language.to_extension())), + ), BinaryType::Agent => ( Config::agent_bin_dir(binary_name).join(format!("{binary_name}.cmd")), Config::agent_bin_dir(binary_name) @@ -501,10 +518,36 @@ impl Functions { })?; let content_template = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) }; let content = match binary_type { - BinaryType::Tool => content_template.replace("{function_name}", binary_name), - BinaryType::Agent => content_template.replace("{agent_name}", binary_name), + BinaryType::Tool(None) => { + let root_dir = Config::functions_dir(); + let tool_path = format!( + "{}/{binary_name}", + &Config::global_tools_dir().to_string_lossy() + ); + content_template + .replace("{function_name}", binary_name) + .replace("{root_dir}", &root_dir.to_string_lossy()) + .replace("{tool_path}", &tool_path) + } + BinaryType::Tool(Some(agent_name)) => { + let root_dir = Config::agent_data_dir(agent_name); + let tool_path = format!( + "{}/{binary_name}", + &Config::global_tools_dir().to_string_lossy() + ); + content_template + .replace("{function_name}", binary_name) + .replace("{root_dir}", &root_dir.to_string_lossy()) + .replace("{tool_path}", &tool_path) + } + BinaryType::Agent => content_template + .replace("{agent_name}", binary_name) + .replace("{config_dir}", &Config::config_dir().to_string_lossy()), } - .replace("{config_dir}", &Config::config_dir().to_string_lossy()); + .replace( + "{prompt_utils_file}", + &Config::bash_prompt_utils_file().to_string_lossy(), + ); if binary_script_file.exists() { fs::remove_file(&binary_script_file)?; } @@ -578,7 +621,10 @@ impl Functions { use std::os::unix::prelude::PermissionsExt; let binary_file = match binary_type { - BinaryType::Tool => Config::functions_bin_dir().join(binary_name), + BinaryType::Tool(None) => Config::functions_bin_dir().join(binary_name), + BinaryType::Tool(Some(agent_name)) => { + Config::agent_bin_dir(agent_name).join(binary_name) + } BinaryType::Agent => Config::agent_bin_dir(binary_name).join(binary_name), }; info!( @@ -600,10 +646,36 @@ impl Functions { })?; let content_template = unsafe { std::str::from_utf8_unchecked(&embedded_file.data) }; let content = match binary_type { - BinaryType::Tool => content_template.replace("{function_name}", binary_name), - BinaryType::Agent => content_template.replace("{agent_name}", binary_name), + BinaryType::Tool(None) => { + let root_dir = Config::functions_dir(); + let tool_path = format!( + "{}/{binary_name}", + &Config::global_tools_dir().to_string_lossy() + ); + content_template + .replace("{function_name}", binary_name) + .replace("{root_dir}", &root_dir.to_string_lossy()) + .replace("{tool_path}", &tool_path) + } + BinaryType::Tool(Some(agent_name)) => { + let root_dir = Config::agent_data_dir(agent_name); + let tool_path = format!( + "{}/{binary_name}", + &Config::global_tools_dir().to_string_lossy() + ); + content_template + .replace("{function_name}", binary_name) + .replace("{root_dir}", &root_dir.to_string_lossy()) + .replace("{tool_path}", &tool_path) + } + BinaryType::Agent => content_template + .replace("{agent_name}", binary_name) + .replace("{config_dir}", &Config::config_dir().to_string_lossy()), } - .replace("{config_dir}", &Config::config_dir().to_string_lossy()); + .replace( + "{prompt_utils_file}", + &Config::bash_prompt_utils_file().to_string_lossy(), + ); if binary_file.exists() { fs::remove_file(&binary_file)?; } @@ -696,6 +768,11 @@ impl ToolCall { Some(agent) => self.extract_call_config_from_agent(config, agent)?, None => self.extract_call_config_from_config(config)?, }; + let agent_name = config + .read() + .agent + .as_ref() + .map(|agent| agent.name().to_owned()); let json_data = if self.arguments.is_object() { self.arguments.clone() @@ -754,7 +831,7 @@ impl ToolCall { let result = registry_arc.invoke(server, tool, arguments).await?; serde_json::to_value(result)? } - _ => match run_llm_function(cmd_name, cmd_args, envs)? { + _ => match run_llm_function(cmd_name, cmd_args, envs, agent_name)? { Some(contents) => serde_json::from_str(&contents) .ok() .unwrap_or_else(|| json!({"output": contents})), @@ -812,15 +889,17 @@ pub fn run_llm_function( cmd_name: String, cmd_args: Vec, mut envs: HashMap, + agent_name: Option, ) -> Result> { let mut bin_dirs: Vec = vec![]; - if cmd_args.len() > 1 { - let dir = Config::agent_bin_dir(&cmd_name); + if let Some(agent_name) = agent_name { + let dir = Config::agent_bin_dir(&agent_name); if dir.exists() { bin_dirs.push(dir); } + } else { + bin_dirs.push(Config::functions_bin_dir()); } - bin_dirs.push(Config::functions_bin_dir()); let current_path = env::var("PATH").context("No PATH environment variable")?; let prepend_path = bin_dirs .iter()