refactor: Agents that depend on global tools now have all binaries compiled and stored in the agent's bin directory so multiple agents can run at once

This commit is contained in:
2025-11-04 11:29:59 -07:00
parent 843abe0621
commit 2f3586cbbf
20 changed files with 342 additions and 48 deletions
+1 -2
View File
@@ -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
+1 -2
View File
@@ -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
+1
View File
@@ -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() {
+3 -3
View File
@@ -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()
main()
+3 -2
View File
@@ -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() {
+1 -2
View File
@@ -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
+1 -2
View File
@@ -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() {
+1 -2
View File
@@ -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() {
+1 -2
View File
@@ -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() {
+1 -2
View File
@@ -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() {
+98
View File
@@ -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 "<current command line>"`, 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<Alt-e>
# 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).
+111
View File
@@ -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 <SECRET_NAME> Add a secret to the Loki vault
--get-secret <SECRET_NAME> Decrypt a secret from the Loki vault and print the plaintext
--update-secret <SECRET_NAME> Update an existing secret in the Loki vault
--delete-secret <SECRET_NAME> 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}')
```
Binary file not shown.

After

Width:  |  Height:  |  Size: 30 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 78 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 490 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 5.8 KiB

Binary file not shown.

After

Width:  |  Height:  |  Size: 4.2 KiB

+1
View File
@@ -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),
+10
View File
@@ -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")),
+108 -29
View File
@@ -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> {
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> {
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<String>,
mut envs: HashMap<String, String>,
agent_name: Option<String>,
) -> Result<Option<String>> {
let mut bin_dirs: Vec<PathBuf> = 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()