f63718d202* chore: add #[non_exhaustive] to reduce backwards-incompatible changes going forward * fix: remove ProtocolVersion import * fix: add a few more with_ mutator methods --------- Co-authored-by: Alex Hancock <alexhancock@block.xyz>
119 lines
3.5 KiB
Rust
119 lines
3.5 KiB
Rust
use std::collections::HashMap;
|
|
|
|
use rig::tool::{ToolDyn as RigTool, ToolEmbeddingDyn, ToolSet};
|
|
use rmcp::{
|
|
RoleClient,
|
|
model::{CallToolRequestParams, CallToolResult, Tool as McpTool},
|
|
service::{RunningService, ServerSink},
|
|
};
|
|
|
|
pub struct McpToolAdaptor {
|
|
tool: McpTool,
|
|
server: ServerSink,
|
|
}
|
|
|
|
impl RigTool for McpToolAdaptor {
|
|
fn name(&self) -> String {
|
|
self.tool.name.to_string()
|
|
}
|
|
|
|
fn definition(
|
|
&self,
|
|
_prompt: String,
|
|
) -> std::pin::Pin<Box<dyn Future<Output = rig::completion::ToolDefinition> + Send + '_>> {
|
|
Box::pin(std::future::ready(rig::completion::ToolDefinition {
|
|
name: self.name(),
|
|
description: self
|
|
.tool
|
|
.description
|
|
.as_deref()
|
|
.unwrap_or_default()
|
|
.to_string(),
|
|
parameters: self.tool.schema_as_json_value(),
|
|
}))
|
|
}
|
|
|
|
fn call(
|
|
&self,
|
|
args: String,
|
|
) -> std::pin::Pin<Box<dyn Future<Output = Result<String, rig::tool::ToolError>> + Send + '_>>
|
|
{
|
|
let server = self.server.clone();
|
|
Box::pin(async move {
|
|
let call_mcp_tool_result = server
|
|
.call_tool(
|
|
CallToolRequestParams::new(self.tool.name.clone()).with_arguments(
|
|
serde_json::from_str(&args).map_err(rig::tool::ToolError::JsonError)?,
|
|
),
|
|
)
|
|
.await
|
|
.inspect(|result| tracing::info!(?result))
|
|
.inspect_err(|error| tracing::error!(%error))
|
|
.map_err(|e| rig::tool::ToolError::ToolCallError(Box::new(e)))?;
|
|
|
|
Ok(convert_mcp_call_tool_result_to_string(call_mcp_tool_result))
|
|
})
|
|
}
|
|
}
|
|
|
|
impl ToolEmbeddingDyn for McpToolAdaptor {
|
|
fn context(&self) -> serde_json::Result<serde_json::Value> {
|
|
serde_json::to_value(self.tool.clone())
|
|
}
|
|
|
|
fn embedding_docs(&self) -> Vec<String> {
|
|
vec![
|
|
self.tool
|
|
.description
|
|
.as_deref()
|
|
.unwrap_or_default()
|
|
.to_string(),
|
|
]
|
|
}
|
|
}
|
|
|
|
pub struct McpManager {
|
|
pub clients: HashMap<String, RunningService<RoleClient, ()>>,
|
|
}
|
|
|
|
impl McpManager {
|
|
pub async fn get_tool_set(&self) -> anyhow::Result<ToolSet> {
|
|
let mut tool_set = ToolSet::default();
|
|
let mut task = tokio::task::JoinSet::<anyhow::Result<_>>::new();
|
|
for client in self.clients.values() {
|
|
let server = client.peer().clone();
|
|
task.spawn(get_tool_set(server));
|
|
}
|
|
let results = task.join_all().await;
|
|
for result in results {
|
|
match result {
|
|
Err(e) => {
|
|
tracing::error!(error = %e, "Failed to get tool set");
|
|
}
|
|
Ok(tools) => {
|
|
tool_set.add_tools(tools);
|
|
}
|
|
}
|
|
}
|
|
Ok(tool_set)
|
|
}
|
|
}
|
|
|
|
pub fn convert_mcp_call_tool_result_to_string(result: CallToolResult) -> String {
|
|
serde_json::to_string(&result).unwrap()
|
|
}
|
|
|
|
pub async fn get_tool_set(server: ServerSink) -> anyhow::Result<ToolSet> {
|
|
let tools = server.list_all_tools().await?;
|
|
let mut tool_builder = ToolSet::builder();
|
|
for tool in tools {
|
|
tracing::info!("get tool: {}", tool.name);
|
|
let adaptor = McpToolAdaptor {
|
|
tool: tool.clone(),
|
|
server: server.clone(),
|
|
};
|
|
tool_builder = tool_builder.dynamic_tool(adaptor);
|
|
}
|
|
let tool_set = tool_builder.build();
|
|
Ok(tool_set)
|
|
}
|