Compare commits
10 Commits
v0.1.1
..
9eeeb11871
| Author | SHA1 | Date | |
|---|---|---|---|
| 9eeeb11871 | |||
| b8db3f689d | |||
| 3b21ce2aa5 | |||
|
|
9bf4fcd943 | ||
|
|
c1f5cfbbda | ||
| 46517a4e15 | |||
|
|
efbe76e1fc | ||
| 245c567d30 | |||
|
|
cbb3d2c34a | ||
| bddec85fa5 |
+2
-1
@@ -2,5 +2,6 @@
|
|||||||
/tmp
|
/tmp
|
||||||
/.env
|
/.env
|
||||||
!cli/**
|
!cli/**
|
||||||
/.idea/
|
.idea/
|
||||||
/loki.iml
|
/loki.iml
|
||||||
|
/.idea/
|
||||||
|
|||||||
Generated
-10
@@ -1,10 +0,0 @@
|
|||||||
# Default ignored files
|
|
||||||
/shelf/
|
|
||||||
/workspace.xml
|
|
||||||
# Editor-based HTTP Client requests
|
|
||||||
/httpRequests/
|
|
||||||
# Datasource local storage ignored files
|
|
||||||
/dataSources/
|
|
||||||
/dataSources.local.xml
|
|
||||||
# Zeppelin ignored files
|
|
||||||
/ZeppelinRemoteNotebooks/
|
|
||||||
@@ -1,3 +1,9 @@
|
|||||||
|
## v0.1.2 (2025-11-08)
|
||||||
|
|
||||||
|
### Refactor
|
||||||
|
|
||||||
|
- Gave the GitHub MCP server a default placeholder value that doesn't require the vault
|
||||||
|
|
||||||
## v0.1.1 (2025-11-08)
|
## v0.1.1 (2025-11-08)
|
||||||
|
|
||||||
## v0.1.0 (2025-11-07)
|
## v0.1.0 (2025-11-07)
|
||||||
|
|||||||
Generated
+1
-1
@@ -3089,7 +3089,7 @@ dependencies = [
|
|||||||
|
|
||||||
[[package]]
|
[[package]]
|
||||||
name = "loki-ai"
|
name = "loki-ai"
|
||||||
version = "0.1.1"
|
version = "0.1.2"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"ansi_colours",
|
"ansi_colours",
|
||||||
"anyhow",
|
"anyhow",
|
||||||
|
|||||||
+1
-1
@@ -1,6 +1,6 @@
|
|||||||
[package]
|
[package]
|
||||||
name = "loki-ai"
|
name = "loki-ai"
|
||||||
version = "0.1.1"
|
version = "0.1.2"
|
||||||
edition = "2024"
|
edition = "2024"
|
||||||
authors = ["Alex Clarke <alex.j.tusa@gmail.com>"]
|
authors = ["Alex Clarke <alex.j.tusa@gmail.com>"]
|
||||||
description = "An all-in-one, batteries included LLM CLI Tool"
|
description = "An all-in-one, batteries included LLM CLI Tool"
|
||||||
|
|||||||
@@ -19,7 +19,6 @@ Coming from [AIChat](https://github.com/sigoden/aichat)? Follow the [migration g
|
|||||||
|
|
||||||
## Quick Links
|
## Quick Links
|
||||||
* [AIChat Migration Guide](./docs/AICHAT-MIGRATION.md): Coming from AIChat? Follow the migration guide to get started.
|
* [AIChat Migration Guide](./docs/AICHAT-MIGRATION.md): Coming from AIChat? Follow the migration guide to get started.
|
||||||
* [History](#history): A history of how Loki came to be.
|
|
||||||
* [Installation](#install): Install Loki
|
* [Installation](#install): Install Loki
|
||||||
* [Getting Started](#getting-started): Get started with Loki by doing first-run setup steps.
|
* [Getting Started](#getting-started): Get started with Loki by doing first-run setup steps.
|
||||||
* [REPL](./docs/REPL.md): Interactive Read-Eval-Print Loop for conversational interactions with LLMs and Loki.
|
* [REPL](./docs/REPL.md): Interactive Read-Eval-Print Loop for conversational interactions with LLMs and Loki.
|
||||||
@@ -41,21 +40,7 @@ Coming from [AIChat](https://github.com/sigoden/aichat)? Follow the [migration g
|
|||||||
* [Client Configurations](./docs/clients/CLIENTS.md): Configuration instructions for various LLM providers.
|
* [Client Configurations](./docs/clients/CLIENTS.md): Configuration instructions for various LLM providers.
|
||||||
* [Patching API Requests](./docs/clients/PATCHES.md): Learn how to patch API requests for advanced customization.
|
* [Patching API Requests](./docs/clients/PATCHES.md): Learn how to patch API requests for advanced customization.
|
||||||
* [Custom Themes](./docs/THEMES.md): Change the look and feel of Loki to your preferences with custom themes.
|
* [Custom Themes](./docs/THEMES.md): Change the look and feel of Loki to your preferences with custom themes.
|
||||||
|
* [History](#history): A history of how Loki came to be.
|
||||||
---
|
|
||||||
|
|
||||||
## History
|
|
||||||
Loki originally started as a fork of the fantastic [AIChat CLI](https://github.com/sigoden/aichat). The purpose was to
|
|
||||||
simply fix a bug in how MCP servers worked with AIChat so that I could specify different ones for agents. However, it
|
|
||||||
has since evolved far beyond that and become a passion project with a life of its own!
|
|
||||||
|
|
||||||
Loki now has first class MCP server support (with support for local and remote servers alike), a built-in vault for
|
|
||||||
interpolating secrets in configuration files, built-in agents, built-in macros, dynamic tab completions, integrated
|
|
||||||
custom functions (no `argc` dependency), improved documentation, and much more with many more plans for the future!
|
|
||||||
|
|
||||||
The original kudos goes out to all the developers of the wonderful AIChat project!
|
|
||||||
|
|
||||||
---
|
|
||||||
|
|
||||||
## Prerequisites
|
## Prerequisites
|
||||||
Loki requires the following tools to be installed on your system:
|
Loki requires the following tools to be installed on your system:
|
||||||
@@ -164,21 +149,6 @@ guide you through the process when you first attempt to access the vault. So, to
|
|||||||
loki --list-secrets
|
loki --list-secrets
|
||||||
```
|
```
|
||||||
|
|
||||||
### First Time Setup
|
|
||||||
In order for Loki to function correctly, you'll need to add a few secrets to the Loki vault so the MCP servers can
|
|
||||||
function.
|
|
||||||
|
|
||||||
**GitHub MCP Server:**
|
|
||||||
* `GITHUB_PERSONAL_ACCESS_TOKEN` - A GitHub Personal Access Token with `repo` and `workflow` scopes.
|
|
||||||
See [Creating a GitHub Personal Access Token](https://docs.github.com/en/authentication/keeping-your-account-and-data-secure/managing-your-personal-access-tokens)
|
|
||||||
|
|
||||||
#### Add the secrets to the Loki vault
|
|
||||||
You can add the secrets to the Loki vault using the following commands (First time use will prompt you to create a vault
|
|
||||||
password file):
|
|
||||||
```sh
|
|
||||||
loki --add-secret GITHUB_PERSONAL_ACCESS_TOKEN
|
|
||||||
```
|
|
||||||
|
|
||||||
### Tab-Completions
|
### Tab-Completions
|
||||||
You can also enable tab completions to make using Loki easier. To do so, add the following to your shell profile:
|
You can also enable tab completions to make using Loki easier. To do so, add the following to your shell profile:
|
||||||
```shell
|
```shell
|
||||||
@@ -272,5 +242,20 @@ The appearance of Loki can be modified using the following settings:
|
|||||||
| `user_agent` | `null` | The name of the `User-Agent` that should be passed in the `User-Agent` header on all requests to model providers |
|
| `user_agent` | `null` | The name of the `User-Agent` that should be passed in the `User-Agent` header on all requests to model providers |
|
||||||
| `save_shell_history` | `true` | Enables or disables REPL command history |
|
| `save_shell_history` | `true` | Enables or disables REPL command history |
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
|
## History
|
||||||
|
Loki originally started as a fork of the fantastic [AIChat CLI](https://github.com/sigoden/aichat). The purpose was to
|
||||||
|
simply fix a bug in how MCP servers worked with AIChat so that I could specify different ones for agents. However, it
|
||||||
|
has since evolved far beyond that and become a passion project with a life of its own!
|
||||||
|
|
||||||
|
Loki now has first class MCP server support (with support for local and remote servers alike), a built-in vault for
|
||||||
|
interpolating secrets in configuration files, built-in agents, built-in macros, dynamic tab completions, integrated
|
||||||
|
custom functions (no `argc` dependency), improved documentation, and much more with many more plans for the future!
|
||||||
|
|
||||||
|
The original kudos goes out to all the developers of the wonderful AIChat project!
|
||||||
|
|
||||||
|
---
|
||||||
|
|
||||||
## Creator
|
## Creator
|
||||||
* [Alex Clarke](https://github.com/Dark-Alex-17)
|
* [Alex Clarke](https://github.com/Dark-Alex-17)
|
||||||
|
|||||||
@@ -11,21 +11,12 @@
|
|||||||
"ghcr.io/github/github-mcp-server"
|
"ghcr.io/github/github-mcp-server"
|
||||||
],
|
],
|
||||||
"env": {
|
"env": {
|
||||||
"GITHUB_PERSONAL_ACCESS_TOKEN": "{{GITHUB_PERSONAL_ACCESS_TOKEN}}"
|
"GITHUB_PERSONAL_ACCESS_TOKEN": "YOUR_GITHUB_TOKEN"
|
||||||
}
|
}
|
||||||
},
|
},
|
||||||
"docker": {
|
"docker": {
|
||||||
"command": "uvx",
|
"command": "uvx",
|
||||||
"args": ["mcp-server-docker"]
|
"args": ["mcp-server-docker"]
|
||||||
},
|
|
||||||
"slack": {
|
|
||||||
"command": "npx",
|
|
||||||
"args": ["-y", "slack-mcp-server@latest", "--transport", "stdio"],
|
|
||||||
"env": {
|
|
||||||
"SLACK_MCP_XOXC_TOKEN": "{{SLACK_MCP_XOXC_TOKEN}}",
|
|
||||||
"SLACK_MCP_XOXD_TOKEN": "{{SLACK_MCP_XOXD_TOKEN}}",
|
|
||||||
"SLACK_MCP_ADD_MESSAGE_TOOL": true
|
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -3,8 +3,8 @@ Loki originally started as a fork of AIChat but has since evolved into its own s
|
|||||||
|
|
||||||
As a result, there's some changes you'll need to make to your AIChat configuration to be able to use Loki.
|
As a result, there's some changes you'll need to make to your AIChat configuration to be able to use Loki.
|
||||||
|
|
||||||
Be sure you've followed the [first-time setup steps](../README.md#first-time-setup) so that the Loki configuration
|
Be sure you've run `loki` at least once so that the Loki configuration directory and subdirectories exist and is
|
||||||
directory and subdirectories exist and is populated with the built-in defaults.
|
populated with the built-in defaults.
|
||||||
|
|
||||||
## Global Configuration File
|
## Global Configuration File
|
||||||
You should be able to copy/paste your AIChat configuration file into your Loki configuration directory. Since the
|
You should be able to copy/paste your AIChat configuration file into your Loki configuration directory. Since the
|
||||||
@@ -208,4 +208,4 @@ Here's how to migrate your functions over to Loki from the `llm-functions` repos
|
|||||||
```
|
```
|
||||||
|
|
||||||
Refer to the [custom bash tools docs](./function-calling/CUSTOM-BASH-TOOLS.md) to learn how to compile and test bash
|
Refer to the [custom bash tools docs](./function-calling/CUSTOM-BASH-TOOLS.md) to learn how to compile and test bash
|
||||||
tools in Loki without needing to use `argc`.
|
tools in Loki without needing to use `argc`.
|
||||||
|
|||||||
+22
-9
@@ -24,7 +24,8 @@ use crate::utils::*;
|
|||||||
|
|
||||||
use crate::config::macros::Macro;
|
use crate::config::macros::Macro;
|
||||||
use crate::mcp::{
|
use crate::mcp::{
|
||||||
MCP_INVOKE_META_FUNCTION_NAME_PREFIX, MCP_LIST_META_FUNCTION_NAME_PREFIX, McpRegistry,
|
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX,
|
||||||
|
MCP_SEARCH_META_FUNCTION_NAME_PREFIX, McpRegistry,
|
||||||
};
|
};
|
||||||
use crate::vault::{GlobalVault, Vault, create_vault_password_file, interpolate_secrets};
|
use crate::vault::{GlobalVault, Vault, create_vault_password_file, interpolate_secrets};
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
use anyhow::{Context, Result, anyhow, bail};
|
||||||
@@ -1972,7 +1973,8 @@ impl Config {
|
|||||||
.iter()
|
.iter()
|
||||||
.filter(|v| {
|
.filter(|v| {
|
||||||
!v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
!v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
||||||
&& !v.name.starts_with(MCP_LIST_META_FUNCTION_NAME_PREFIX)
|
&& !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
||||||
|
&& !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
||||||
})
|
})
|
||||||
.map(|v| v.name.to_string())
|
.map(|v| v.name.to_string())
|
||||||
.collect();
|
.collect();
|
||||||
@@ -2015,7 +2017,8 @@ impl Config {
|
|||||||
.into_iter()
|
.into_iter()
|
||||||
.filter(|v| {
|
.filter(|v| {
|
||||||
!v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
!v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
||||||
&& !v.name.starts_with(MCP_LIST_META_FUNCTION_NAME_PREFIX)
|
&& !v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
||||||
|
&& !v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
let tool_names: HashSet<String> = agent_functions
|
let tool_names: HashSet<String> = agent_functions
|
||||||
@@ -2051,7 +2054,8 @@ impl Config {
|
|||||||
.iter()
|
.iter()
|
||||||
.filter(|v| {
|
.filter(|v| {
|
||||||
v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
||||||
|| v.name.starts_with(MCP_LIST_META_FUNCTION_NAME_PREFIX)
|
|| v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
||||||
|
|| v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
||||||
})
|
})
|
||||||
.map(|v| v.name.to_string())
|
.map(|v| v.name.to_string())
|
||||||
.collect();
|
.collect();
|
||||||
@@ -2062,8 +2066,10 @@ impl Config {
|
|||||||
let item = item.trim();
|
let item = item.trim();
|
||||||
let item_invoke_name =
|
let item_invoke_name =
|
||||||
format!("{}_{item}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX);
|
format!("{}_{item}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX);
|
||||||
let item_list_name =
|
let item_search_name =
|
||||||
format!("{}_{item}", MCP_LIST_META_FUNCTION_NAME_PREFIX);
|
format!("{}_{item}", MCP_SEARCH_META_FUNCTION_NAME_PREFIX);
|
||||||
|
let item_describe_name =
|
||||||
|
format!("{}_{item}", MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX);
|
||||||
if let Some(values) = self.mapping_mcp_servers.get(item) {
|
if let Some(values) = self.mapping_mcp_servers.get(item) {
|
||||||
server_names.extend(
|
server_names.extend(
|
||||||
values
|
values
|
||||||
@@ -2077,7 +2083,12 @@ impl Config {
|
|||||||
),
|
),
|
||||||
format!(
|
format!(
|
||||||
"{}_{}",
|
"{}_{}",
|
||||||
MCP_LIST_META_FUNCTION_NAME_PREFIX,
|
MCP_SEARCH_META_FUNCTION_NAME_PREFIX,
|
||||||
|
v.to_string()
|
||||||
|
),
|
||||||
|
format!(
|
||||||
|
"{}_{}",
|
||||||
|
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX,
|
||||||
v.to_string()
|
v.to_string()
|
||||||
),
|
),
|
||||||
]
|
]
|
||||||
@@ -2086,7 +2097,8 @@ impl Config {
|
|||||||
)
|
)
|
||||||
} else if mcp_declaration_names.contains(&item_invoke_name) {
|
} else if mcp_declaration_names.contains(&item_invoke_name) {
|
||||||
server_names.insert(item_invoke_name);
|
server_names.insert(item_invoke_name);
|
||||||
server_names.insert(item_list_name);
|
server_names.insert(item_search_name);
|
||||||
|
server_names.insert(item_describe_name);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
@@ -2112,7 +2124,8 @@ impl Config {
|
|||||||
.into_iter()
|
.into_iter()
|
||||||
.filter(|v| {
|
.filter(|v| {
|
||||||
v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
v.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
||||||
|| v.name.starts_with(MCP_LIST_META_FUNCTION_NAME_PREFIX)
|
|| v.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
||||||
|
|| v.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
||||||
})
|
})
|
||||||
.collect();
|
.collect();
|
||||||
let tool_names: HashSet<String> = agent_functions
|
let tool_names: HashSet<String> = agent_functions
|
||||||
|
|||||||
+150
-48
@@ -4,7 +4,10 @@ use crate::{
|
|||||||
};
|
};
|
||||||
|
|
||||||
use crate::config::ensure_parent_exists;
|
use crate::config::ensure_parent_exists;
|
||||||
use crate::mcp::{MCP_INVOKE_META_FUNCTION_NAME_PREFIX, MCP_LIST_META_FUNCTION_NAME_PREFIX};
|
use crate::mcp::{
|
||||||
|
MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX, MCP_INVOKE_META_FUNCTION_NAME_PREFIX,
|
||||||
|
MCP_SEARCH_META_FUNCTION_NAME_PREFIX,
|
||||||
|
};
|
||||||
use crate::parsers::{bash, python};
|
use crate::parsers::{bash, python};
|
||||||
use anyhow::{Context, Result, anyhow, bail};
|
use anyhow::{Context, Result, anyhow, bail};
|
||||||
use indexmap::IndexMap;
|
use indexmap::IndexMap;
|
||||||
@@ -247,19 +250,13 @@ impl Functions {
|
|||||||
pub fn clear_mcp_meta_functions(&mut self) {
|
pub fn clear_mcp_meta_functions(&mut self) {
|
||||||
self.declarations.retain(|d| {
|
self.declarations.retain(|d| {
|
||||||
!d.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
!d.name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX)
|
||||||
&& !d.name.starts_with(MCP_LIST_META_FUNCTION_NAME_PREFIX)
|
&& !d.name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX)
|
||||||
|
&& !d.name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX)
|
||||||
});
|
});
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec<String>) {
|
pub fn append_mcp_meta_functions(&mut self, mcp_servers: Vec<String>) {
|
||||||
let mut invoke_function_properties = IndexMap::new();
|
let mut invoke_function_properties = IndexMap::new();
|
||||||
invoke_function_properties.insert(
|
|
||||||
"server".to_string(),
|
|
||||||
JsonSchema {
|
|
||||||
type_value: Some("string".to_string()),
|
|
||||||
..Default::default()
|
|
||||||
},
|
|
||||||
);
|
|
||||||
invoke_function_properties.insert(
|
invoke_function_properties.insert(
|
||||||
"tool".to_string(),
|
"tool".to_string(),
|
||||||
JsonSchema {
|
JsonSchema {
|
||||||
@@ -275,32 +272,86 @@ impl Functions {
|
|||||||
},
|
},
|
||||||
);
|
);
|
||||||
|
|
||||||
|
let mut search_function_properties = IndexMap::new();
|
||||||
|
search_function_properties.insert(
|
||||||
|
"query".to_string(),
|
||||||
|
JsonSchema {
|
||||||
|
type_value: Some("string".to_string()),
|
||||||
|
description: Some("Generalized explanation of what you want to do".into()),
|
||||||
|
..Default::default()
|
||||||
|
},
|
||||||
|
);
|
||||||
|
search_function_properties.insert(
|
||||||
|
"top_k".to_string(),
|
||||||
|
JsonSchema {
|
||||||
|
type_value: Some("integer".to_string()),
|
||||||
|
description: Some("How many results to return, between 1 and 20".into()),
|
||||||
|
default: Some(Value::from(8usize)),
|
||||||
|
..Default::default()
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
|
let mut describe_function_properties = IndexMap::new();
|
||||||
|
describe_function_properties.insert(
|
||||||
|
"tool".to_string(),
|
||||||
|
JsonSchema {
|
||||||
|
type_value: Some("string".to_string()),
|
||||||
|
description: Some("The name of the tool; e.g., search_issues".into()),
|
||||||
|
..Default::default()
|
||||||
|
},
|
||||||
|
);
|
||||||
|
|
||||||
for server in mcp_servers {
|
for server in mcp_servers {
|
||||||
|
let search_function_name = format!("{}_{server}", MCP_SEARCH_META_FUNCTION_NAME_PREFIX);
|
||||||
|
let describe_function_name =
|
||||||
|
format!("{}_{server}", MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX);
|
||||||
let invoke_function_name = format!("{}_{server}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX);
|
let invoke_function_name = format!("{}_{server}", MCP_INVOKE_META_FUNCTION_NAME_PREFIX);
|
||||||
let invoke_function_declaration = FunctionDeclaration {
|
let invoke_function_declaration = FunctionDeclaration {
|
||||||
name: invoke_function_name.clone(),
|
name: invoke_function_name.clone(),
|
||||||
description: formatdoc!(
|
description: formatdoc!(
|
||||||
r#"
|
r#"
|
||||||
Invoke the specified tool on the {server} MCP server. Always call {invoke_function_name} first to find the
|
Invoke the specified tool on the {server} MCP server. Always call {describe_function_name} first to
|
||||||
correct names of tools before calling '{invoke_function_name}'.
|
find the correct invocation schema for the given tool.
|
||||||
"#
|
"#
|
||||||
),
|
),
|
||||||
parameters: JsonSchema {
|
parameters: JsonSchema {
|
||||||
type_value: Some("object".to_string()),
|
type_value: Some("object".to_string()),
|
||||||
properties: Some(invoke_function_properties.clone()),
|
properties: Some(invoke_function_properties.clone()),
|
||||||
required: Some(vec!["server".to_string(), "tool".to_string()]),
|
required: Some(vec!["tool".to_string()]),
|
||||||
..Default::default()
|
..Default::default()
|
||||||
},
|
},
|
||||||
agent: false,
|
agent: false,
|
||||||
};
|
};
|
||||||
let list_functions_declaration = FunctionDeclaration {
|
let search_functions_declaration = FunctionDeclaration {
|
||||||
name: format!("{}_{}", MCP_LIST_META_FUNCTION_NAME_PREFIX, server),
|
name: search_function_name.clone(),
|
||||||
description: format!("List all the available tools for the {server} MCP server"),
|
description: formatdoc!(
|
||||||
parameters: JsonSchema::default(),
|
r#"
|
||||||
|
Find candidate tools by keywords for the {server} MCP server. Returns small suggestions; fetch
|
||||||
|
schemas with {describe_function_name}.
|
||||||
|
"#
|
||||||
|
),
|
||||||
|
parameters: JsonSchema {
|
||||||
|
type_value: Some("object".to_string()),
|
||||||
|
properties: Some(search_function_properties.clone()),
|
||||||
|
required: Some(vec!["query".to_string()]),
|
||||||
|
..Default::default()
|
||||||
|
},
|
||||||
|
agent: false,
|
||||||
|
};
|
||||||
|
let describe_functions_declaration = FunctionDeclaration {
|
||||||
|
name: describe_function_name.clone(),
|
||||||
|
description: "Get the full JSON schema for exactly one MCP tool.".to_string(),
|
||||||
|
parameters: JsonSchema {
|
||||||
|
type_value: Some("object".to_string()),
|
||||||
|
properties: Some(describe_function_properties.clone()),
|
||||||
|
required: Some(vec!["tool".to_string()]),
|
||||||
|
..Default::default()
|
||||||
|
},
|
||||||
agent: false,
|
agent: false,
|
||||||
};
|
};
|
||||||
self.declarations.push(invoke_function_declaration);
|
self.declarations.push(invoke_function_declaration);
|
||||||
self.declarations.push(list_functions_declaration);
|
self.declarations.push(search_functions_declaration);
|
||||||
|
self.declarations.push(describe_functions_declaration);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -771,39 +822,14 @@ impl ToolCall {
|
|||||||
}
|
}
|
||||||
|
|
||||||
let output = match cmd_name.as_str() {
|
let output = match cmd_name.as_str() {
|
||||||
_ if cmd_name.starts_with(MCP_LIST_META_FUNCTION_NAME_PREFIX) => {
|
_ if cmd_name.starts_with(MCP_SEARCH_META_FUNCTION_NAME_PREFIX) => {
|
||||||
let registry_arc = {
|
Self::search_mcp_tools(config, &cmd_name, &json_data)?
|
||||||
let cfg = config.read();
|
}
|
||||||
cfg.mcp_registry
|
_ if cmd_name.starts_with(MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX) => {
|
||||||
.clone()
|
Self::describe_mcp_tool(config, &cmd_name, json_data).await?
|
||||||
.with_context(|| "MCP is not configured")?
|
|
||||||
};
|
|
||||||
|
|
||||||
registry_arc.catalog().await?
|
|
||||||
}
|
}
|
||||||
_ if cmd_name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) => {
|
_ if cmd_name.starts_with(MCP_INVOKE_META_FUNCTION_NAME_PREFIX) => {
|
||||||
let server = json_data
|
Self::invoke_mcp_tool(config, &cmd_name, &json_data).await?
|
||||||
.get("server")
|
|
||||||
.ok_or_else(|| anyhow!("Missing 'server' in arguments"))?
|
|
||||||
.as_str()
|
|
||||||
.ok_or_else(|| anyhow!("Invalid 'server' in arguments"))?;
|
|
||||||
let tool = json_data
|
|
||||||
.get("tool")
|
|
||||||
.ok_or_else(|| anyhow!("Missing 'tool' in arguments"))?
|
|
||||||
.as_str()
|
|
||||||
.ok_or_else(|| anyhow!("Invalid 'tool' in arguments"))?;
|
|
||||||
let arguments = json_data
|
|
||||||
.get("arguments")
|
|
||||||
.cloned()
|
|
||||||
.unwrap_or_else(|| json!({}));
|
|
||||||
let registry_arc = {
|
|
||||||
let cfg = config.read();
|
|
||||||
cfg.mcp_registry
|
|
||||||
.clone()
|
|
||||||
.with_context(|| "MCP is not configured")?
|
|
||||||
};
|
|
||||||
let result = registry_arc.invoke(server, tool, arguments).await?;
|
|
||||||
serde_json::to_value(result)?
|
|
||||||
}
|
}
|
||||||
_ => match run_llm_function(cmd_name, cmd_args, envs, agent_name)? {
|
_ => match run_llm_function(cmd_name, cmd_args, envs, agent_name)? {
|
||||||
Some(contents) => serde_json::from_str(&contents)
|
Some(contents) => serde_json::from_str(&contents)
|
||||||
@@ -816,6 +842,82 @@ impl ToolCall {
|
|||||||
Ok(output)
|
Ok(output)
|
||||||
}
|
}
|
||||||
|
|
||||||
|
async fn describe_mcp_tool(
|
||||||
|
config: &GlobalConfig,
|
||||||
|
cmd_name: &str,
|
||||||
|
json_data: Value,
|
||||||
|
) -> Result<Value> {
|
||||||
|
let server_id = cmd_name.replace(&format!("{MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX}_"), "");
|
||||||
|
let tool = json_data
|
||||||
|
.get("tool")
|
||||||
|
.ok_or_else(|| anyhow!("Missing 'tool' in arguments"))?
|
||||||
|
.as_str()
|
||||||
|
.ok_or_else(|| anyhow!("Invalid 'tool' in arguments"))?;
|
||||||
|
let registry_arc = {
|
||||||
|
let cfg = config.read();
|
||||||
|
cfg.mcp_registry
|
||||||
|
.clone()
|
||||||
|
.with_context(|| "MCP is not configured")?
|
||||||
|
};
|
||||||
|
|
||||||
|
let result = registry_arc.describe(&server_id, tool).await?;
|
||||||
|
Ok(serde_json::to_value(result)?)
|
||||||
|
}
|
||||||
|
|
||||||
|
fn search_mcp_tools(config: &GlobalConfig, cmd_name: &str, json_data: &Value) -> Result<Value> {
|
||||||
|
let server = cmd_name.replace(&format!("{MCP_SEARCH_META_FUNCTION_NAME_PREFIX}_"), "");
|
||||||
|
let query = json_data
|
||||||
|
.get("query")
|
||||||
|
.ok_or_else(|| anyhow!("Missing 'query' in arguments"))?
|
||||||
|
.as_str()
|
||||||
|
.ok_or_else(|| anyhow!("Invalid 'query' in arguments"))?;
|
||||||
|
let top_k = json_data
|
||||||
|
.get("top_k")
|
||||||
|
.cloned()
|
||||||
|
.unwrap_or_else(|| Value::from(8u64))
|
||||||
|
.as_u64()
|
||||||
|
.ok_or_else(|| anyhow!("Invalid 'top_k' in arguments"))? as usize;
|
||||||
|
let registry_arc = {
|
||||||
|
let cfg = config.read();
|
||||||
|
cfg.mcp_registry
|
||||||
|
.clone()
|
||||||
|
.with_context(|| "MCP is not configured")?
|
||||||
|
};
|
||||||
|
|
||||||
|
let catalog_items = registry_arc
|
||||||
|
.search_tools_server(&server, query, top_k)
|
||||||
|
.into_iter()
|
||||||
|
.map(|it| serde_json::to_value(&it).unwrap_or_default())
|
||||||
|
.collect();
|
||||||
|
Ok(Value::Array(catalog_items))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn invoke_mcp_tool(
|
||||||
|
config: &GlobalConfig,
|
||||||
|
cmd_name: &str,
|
||||||
|
json_data: &Value,
|
||||||
|
) -> Result<Value> {
|
||||||
|
let server = cmd_name.replace(&format!("{MCP_INVOKE_META_FUNCTION_NAME_PREFIX}_"), "");
|
||||||
|
let tool = json_data
|
||||||
|
.get("tool")
|
||||||
|
.ok_or_else(|| anyhow!("Missing 'tool' in arguments"))?
|
||||||
|
.as_str()
|
||||||
|
.ok_or_else(|| anyhow!("Invalid 'tool' in arguments"))?;
|
||||||
|
let arguments = json_data
|
||||||
|
.get("arguments")
|
||||||
|
.cloned()
|
||||||
|
.unwrap_or_else(|| json!({}));
|
||||||
|
let registry_arc = {
|
||||||
|
let cfg = config.read();
|
||||||
|
cfg.mcp_registry
|
||||||
|
.clone()
|
||||||
|
.with_context(|| "MCP is not configured")?
|
||||||
|
};
|
||||||
|
|
||||||
|
let result = registry_arc.invoke(&server, tool, arguments).await?;
|
||||||
|
Ok(serde_json::to_value(result)?)
|
||||||
|
}
|
||||||
|
|
||||||
fn extract_call_config_from_agent(
|
fn extract_call_config_from_agent(
|
||||||
&self,
|
&self,
|
||||||
config: &GlobalConfig,
|
config: &GlobalConfig,
|
||||||
|
|||||||
+117
-27
@@ -2,6 +2,7 @@ use crate::config::Config;
|
|||||||
use crate::utils::{AbortSignal, abortable_run_with_spinner};
|
use crate::utils::{AbortSignal, abortable_run_with_spinner};
|
||||||
use crate::vault::interpolate_secrets;
|
use crate::vault::interpolate_secrets;
|
||||||
use anyhow::{Context, Result, anyhow};
|
use anyhow::{Context, Result, anyhow};
|
||||||
|
use bm25::{Document, Language, SearchEngine, SearchEngineBuilder};
|
||||||
use futures_util::future::BoxFuture;
|
use futures_util::future::BoxFuture;
|
||||||
use futures_util::{StreamExt, TryStreamExt, stream};
|
use futures_util::{StreamExt, TryStreamExt, stream};
|
||||||
use indoc::formatdoc;
|
use indoc::formatdoc;
|
||||||
@@ -9,7 +10,7 @@ use rmcp::model::{CallToolRequestParam, CallToolResult};
|
|||||||
use rmcp::service::RunningService;
|
use rmcp::service::RunningService;
|
||||||
use rmcp::transport::TokioChildProcess;
|
use rmcp::transport::TokioChildProcess;
|
||||||
use rmcp::{RoleClient, ServiceExt};
|
use rmcp::{RoleClient, ServiceExt};
|
||||||
use serde::Deserialize;
|
use serde::{Deserialize, Serialize};
|
||||||
use serde_json::{Value, json};
|
use serde_json::{Value, json};
|
||||||
use std::borrow::Cow;
|
use std::borrow::Cow;
|
||||||
use std::collections::{HashMap, HashSet};
|
use std::collections::{HashMap, HashSet};
|
||||||
@@ -20,10 +21,46 @@ use std::sync::Arc;
|
|||||||
use tokio::process::Command;
|
use tokio::process::Command;
|
||||||
|
|
||||||
pub const MCP_INVOKE_META_FUNCTION_NAME_PREFIX: &str = "mcp_invoke";
|
pub const MCP_INVOKE_META_FUNCTION_NAME_PREFIX: &str = "mcp_invoke";
|
||||||
pub const MCP_LIST_META_FUNCTION_NAME_PREFIX: &str = "mcp_list";
|
pub const MCP_SEARCH_META_FUNCTION_NAME_PREFIX: &str = "mcp_search";
|
||||||
|
pub const MCP_DESCRIBE_META_FUNCTION_NAME_PREFIX: &str = "mcp_describe";
|
||||||
|
|
||||||
type ConnectedServer = RunningService<RoleClient, ()>;
|
type ConnectedServer = RunningService<RoleClient, ()>;
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, Default, Serialize)]
|
||||||
|
pub struct CatalogItem {
|
||||||
|
pub name: String,
|
||||||
|
pub server: String,
|
||||||
|
pub description: String,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug)]
|
||||||
|
struct ServerCatalog {
|
||||||
|
engine: SearchEngine<String>,
|
||||||
|
items: HashMap<String, CatalogItem>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl ServerCatalog {
|
||||||
|
pub fn build_bm25(items: &HashMap<String, CatalogItem>) -> SearchEngine<String> {
|
||||||
|
let docs = items.values().map(|it| {
|
||||||
|
let contents = format!("{}\n{}\nserver:{}", it.name, it.description, it.server);
|
||||||
|
Document {
|
||||||
|
id: it.name.clone(),
|
||||||
|
contents,
|
||||||
|
}
|
||||||
|
});
|
||||||
|
SearchEngineBuilder::<String>::with_documents(Language::English, docs).build()
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Clone for ServerCatalog {
|
||||||
|
fn clone(&self) -> Self {
|
||||||
|
Self {
|
||||||
|
engine: Self::build_bm25(&self.items),
|
||||||
|
items: self.items.clone(),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[derive(Debug, Clone, Deserialize)]
|
#[derive(Debug, Clone, Deserialize)]
|
||||||
struct McpServersConfig {
|
struct McpServersConfig {
|
||||||
#[serde(rename = "mcpServers")]
|
#[serde(rename = "mcpServers")]
|
||||||
@@ -50,7 +87,8 @@ enum JsonField {
|
|||||||
pub struct McpRegistry {
|
pub struct McpRegistry {
|
||||||
log_path: Option<PathBuf>,
|
log_path: Option<PathBuf>,
|
||||||
config: Option<McpServersConfig>,
|
config: Option<McpServersConfig>,
|
||||||
servers: HashMap<String, Arc<RunningService<RoleClient, ()>>>,
|
servers: HashMap<String, Arc<ConnectedServer>>,
|
||||||
|
catalogs: HashMap<String, ServerCatalog>,
|
||||||
}
|
}
|
||||||
|
|
||||||
impl McpRegistry {
|
impl McpRegistry {
|
||||||
@@ -173,7 +211,7 @@ impl McpRegistry {
|
|||||||
.collect()
|
.collect()
|
||||||
};
|
};
|
||||||
|
|
||||||
let results: Vec<(String, Arc<_>)> = stream::iter(
|
let results: Vec<(String, Arc<_>, ServerCatalog)> = stream::iter(
|
||||||
server_ids
|
server_ids
|
||||||
.into_iter()
|
.into_iter()
|
||||||
.map(|id| async { self.start_server(id).await }),
|
.map(|id| async { self.start_server(id).await }),
|
||||||
@@ -182,13 +220,24 @@ impl McpRegistry {
|
|||||||
.try_collect()
|
.try_collect()
|
||||||
.await?;
|
.await?;
|
||||||
|
|
||||||
self.servers = results.into_iter().collect();
|
self.servers = results
|
||||||
|
.clone()
|
||||||
|
.into_iter()
|
||||||
|
.map(|(id, server, _)| (id, server))
|
||||||
|
.collect();
|
||||||
|
self.catalogs = results
|
||||||
|
.into_iter()
|
||||||
|
.map(|(id, _, catalog)| (id, catalog))
|
||||||
|
.collect();
|
||||||
}
|
}
|
||||||
|
|
||||||
Ok(())
|
Ok(())
|
||||||
}
|
}
|
||||||
|
|
||||||
async fn start_server(&self, id: String) -> Result<(String, Arc<ConnectedServer>)> {
|
async fn start_server(
|
||||||
|
&self,
|
||||||
|
id: String,
|
||||||
|
) -> Result<(String, Arc<ConnectedServer>, ServerCatalog)> {
|
||||||
let server = self
|
let server = self
|
||||||
.config
|
.config
|
||||||
.as_ref()
|
.as_ref()
|
||||||
@@ -231,14 +280,33 @@ impl McpRegistry {
|
|||||||
.await
|
.await
|
||||||
.with_context(|| format!("Failed to start MCP server: {}", &server.command))?,
|
.with_context(|| format!("Failed to start MCP server: {}", &server.command))?,
|
||||||
);
|
);
|
||||||
debug!(
|
let tools = service.list_tools(None).await?;
|
||||||
"Available tools for MCP server {id}: {:?}",
|
debug!("Available tools for MCP server {id}: {tools:?}");
|
||||||
service.list_tools(None).await?
|
|
||||||
);
|
let mut items_vec = Vec::new();
|
||||||
|
for t in tools.tools {
|
||||||
|
let name = t.name.to_string();
|
||||||
|
let description = t.description.unwrap_or_default().to_string();
|
||||||
|
items_vec.push(CatalogItem {
|
||||||
|
name,
|
||||||
|
server: id.clone(),
|
||||||
|
description,
|
||||||
|
});
|
||||||
|
}
|
||||||
|
|
||||||
|
let mut items_map = HashMap::new();
|
||||||
|
items_vec.into_iter().for_each(|it| {
|
||||||
|
items_map.insert(it.name.clone(), it);
|
||||||
|
});
|
||||||
|
|
||||||
|
let catalog = ServerCatalog {
|
||||||
|
engine: ServerCatalog::build_bm25(&items_map),
|
||||||
|
items: items_map,
|
||||||
|
};
|
||||||
|
|
||||||
info!("Started MCP server: {id}");
|
info!("Started MCP server: {id}");
|
||||||
|
|
||||||
Ok((id.to_string(), service))
|
Ok((id.to_string(), service, catalog))
|
||||||
}
|
}
|
||||||
|
|
||||||
pub async fn stop_all_servers(mut self) -> Result<Self> {
|
pub async fn stop_all_servers(mut self) -> Result<Self> {
|
||||||
@@ -268,26 +336,48 @@ impl McpRegistry {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn catalog(&self) -> BoxFuture<'static, Result<Value>> {
|
pub fn search_tools_server(&self, server: &str, query: &str, top_k: usize) -> Vec<CatalogItem> {
|
||||||
let servers: Vec<(String, Arc<ConnectedServer>)> = self
|
let Some(catalog) = self.catalogs.get(server) else {
|
||||||
|
return vec![];
|
||||||
|
};
|
||||||
|
let engine = &catalog.engine;
|
||||||
|
let raw = engine.search(query, top_k.min(20));
|
||||||
|
|
||||||
|
raw.into_iter()
|
||||||
|
.filter_map(|r| catalog.items.get(&r.document.id))
|
||||||
|
.take(top_k)
|
||||||
|
.cloned()
|
||||||
|
.collect()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn describe(&self, server_id: &str, tool: &str) -> Result<Value> {
|
||||||
|
let server = self
|
||||||
.servers
|
.servers
|
||||||
.iter()
|
.iter()
|
||||||
.map(|(id, s)| (id.clone(), s.clone()))
|
.filter(|(id, _)| &server_id == id)
|
||||||
.collect();
|
.map(|(_, s)| s.clone())
|
||||||
|
.next()
|
||||||
|
.ok_or(anyhow!("{server_id} MCP server not found in config"))?;
|
||||||
|
|
||||||
Box::pin(async move {
|
let tool_schema = server
|
||||||
let mut out = Vec::with_capacity(servers.len());
|
.list_tools(None)
|
||||||
for (id, server) in servers {
|
.await?
|
||||||
let tools = server.list_tools(None).await?;
|
.tools
|
||||||
let resources = server.list_resources(None).await.unwrap_or_default();
|
.into_iter()
|
||||||
out.push(json!({
|
.find(|it| it.name == tool)
|
||||||
"server": id,
|
.ok_or(anyhow!(
|
||||||
"tools": tools,
|
"{tool} not found in {server_id} MCP server catalog"
|
||||||
"resources": resources,
|
))?
|
||||||
}));
|
.input_schema;
|
||||||
|
Ok(json!({
|
||||||
|
"type": "object",
|
||||||
|
"properties": {
|
||||||
|
"tool": {
|
||||||
|
"type": "string",
|
||||||
|
},
|
||||||
|
"arguments": tool_schema
|
||||||
}
|
}
|
||||||
Ok(Value::Array(out))
|
}))
|
||||||
})
|
|
||||||
}
|
}
|
||||||
|
|
||||||
pub fn invoke(
|
pub fn invoke(
|
||||||
|
|||||||
Reference in New Issue
Block a user