Move whole rmcp crate to offcial rust sdk (#44)

* get rid of async-trait

* move rmcp to official-rust-sdk

* replace master with main

* update readme, change license

* fix typo

* Update README.md

Co-authored-by: Jefry Dewangga <jefrydco@gmail.com>

* fix typos

---------

Co-authored-by: = <=>
Co-authored-by: Jefry Dewangga <jefrydco@gmail.com>

examples/rig-integration/Cargo.toml

@@ -0,0 +1,26 @@
+[package]
+name = "rig-integration"
+edition.workspace = true
+version.workspace = true
+authors.workspace = true
+license.workspace = true
+repository.workspace = true
+description.workspace = true
+keywords.workspace = true
+homepage.workspace = true
+categories.workspace = true
+readme.workspace = true
+
+[dependencies]
+rig-core = "0.10.0"
+tokio = { version = "1", features = ["full"] }
+rmcp = { path = "../../crates/rmcp", features = [
+    "client",
+    "transport-child-process",
+    "transport-sse",
+], no-default-features = true }
+anyhow = "1.0"
+serde_json = "1"
+serde = { version = "1", features = ["derive"] }
+toml = "0.8"
+futures = "0.3"

examples/rig-integration/config.toml

@@ -0,0 +1,10 @@
+deepseek_key = ""
+cohere_key = ""
+
+[mcp]
+
+[[mcp.server]]
+name = "git"
+protocol = "stdio"
+command = "uvx"
+args = ["mcp-server-git"]

examples/rig-integration/src/chat.rs

@@ -0,0 +1,114 @@
+use futures::StreamExt;
+use rig::{
+    agent::Agent,
+    message::Message,
+    providers::deepseek::DeepSeekCompletionModel,
+    streaming::{StreamingChat, StreamingChoice},
+};
+use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter};
+
+pub async fn cli_chatbot(chatbot: Agent<DeepSeekCompletionModel>) -> anyhow::Result<()> {
+    let mut chat_log = vec![];
+
+    let mut output = BufWriter::new(tokio::io::stdout());
+    let mut input = BufReader::new(tokio::io::stdin());
+    output.write_all(b"Enter :q to quit\n").await?;
+    loop {
+        output.write_all(b"\x1b[32muser>\x1b[0m ").await?;
+        // Flush stdout to ensure the prompt appears before input
+        output.flush().await?;
+        let mut input_buf = String::new();
+        input.read_line(&mut input_buf).await?;
+        // Remove the newline character from the input
+        let input = input_buf.trim();
+        // Check for a command to exit
+        if input == ":q" {
+            break;
+        }
+        match chatbot.stream_chat(input, chat_log.clone()).await {
+            Ok(mut response) => {
+                chat_log.push(Message::user(input));
+                stream_output_agent_start(&mut output).await?;
+                let mut message_buf = String::new();
+                while let Some(message) = response.next().await {
+                    match message {
+                        Ok(StreamingChoice::Message(text)) => {
+                            message_buf.push_str(&text);
+                            output_agent(text, &mut output).await?;
+                        }
+                        Ok(StreamingChoice::ToolCall(name, _, param)) => {
+                            chat_log.push(Message::assistant(format!(
+                                "Calling tool: {name} with args: {param}"
+                            )));
+                            let result = chatbot.tools.call(&name, param.to_string()).await;
+                            match result {
+                                Ok(tool_call_result) => {
+                                    stream_output_agent_finished(&mut output).await?;
+                                    stream_output_toolcall(&tool_call_result, &mut output).await?;
+                                    stream_output_agent_start(&mut output).await?;
+                                    chat_log.push(Message::user(tool_call_result));
+                                }
+                                Err(e) => {
+                                    output_error(e, &mut output).await?;
+                                }
+                            }
+                        }
+                        Err(error) => {
+                            output_error(error, &mut output).await?;
+                        }
+                    }
+                }
+                chat_log.push(Message::assistant(message_buf));
+                stream_output_agent_finished(&mut output).await?;
+            }
+            Err(error) => {
+                output_error(error, &mut output).await?;
+            }
+        }
+    }
+
+    Ok(())
+}
+
+pub async fn output_error(
+    e: impl std::fmt::Display,
+    output: &mut BufWriter<tokio::io::Stdout>,
+) -> std::io::Result<()> {
+    output.write_all(b"\x1b[31mERROR: \x1b[0m").await?;
+    output.write_all(e.to_string().as_bytes()).await?;
+    output.write_all(b"\n").await?;
+    Ok(())
+}
+
+pub async fn output_agent(
+    content: impl std::fmt::Display,
+    output: &mut BufWriter<tokio::io::Stdout>,
+) -> std::io::Result<()> {
+    output.write_all(content.to_string().as_bytes()).await?;
+    output.flush().await?;
+    Ok(())
+}
+
+pub async fn stream_output_toolcall(
+    content: impl std::fmt::Display,
+    output: &mut BufWriter<tokio::io::Stdout>,
+) -> std::io::Result<()> {
+    output.write_all(b"\x1b[33mtool>\x1b[0m ").await?;
+    output.write_all(content.to_string().as_bytes()).await?;
+    output.write_all(b"\n").await?;
+    Ok(())
+}
+
+pub async fn stream_output_agent_start(
+    output: &mut BufWriter<tokio::io::Stdout>,
+) -> std::io::Result<()> {
+    output.write_all(b"\x1b[34magent>\x1b[0m ").await?;
+    Ok(())
+}
+
+pub async fn stream_output_agent_finished(
+    output: &mut BufWriter<tokio::io::Stdout>,
+) -> std::io::Result<()> {
+    output.write_all(b"\n").await?;
+    Ok(())
+}

examples/rig-integration/src/config.rs

@@ -0,0 +1,20 @@
+use std::path::Path;
+
+use serde::{Deserialize, Serialize};
+
+pub mod mcp;
+
+#[derive(Debug, Deserialize, Serialize)]
+pub struct Config {
+    pub mcp: mcp::McpConfig,
+    pub deepseek_key: Option<String>,
+    pub cohere_key: Option<String>,
+}
+
+impl Config {
+    pub async fn retrive(path: impl AsRef<Path>) -> anyhow::Result<Self> {
+        let content = tokio::fs::read_to_string(path).await?;
+        let config: Self = toml::from_str(&content)?;
+        Ok(config)
+    }
+}

examples/rig-integration/src/config/mcp.rs

@@ -0,0 +1,83 @@
+use std::{collections::HashMap, process::Stdio};
+
+use rmcp::{RoleClient, ServiceExt, service::RunningService};
+use serde::{Deserialize, Serialize};
+
+use crate::mcp_adaptor::McpManager;
+#[derive(Debug, Serialize, Deserialize, Clone)]
+pub struct McpServerConfig {
+    name: String,
+    #[serde(flatten)]
+    transport: McpServerTransportConfig,
+}
+
+#[derive(Debug, Serialize, Deserialize, Clone)]
+#[serde(tag = "protocol", rename_all = "lowercase")]
+pub enum McpServerTransportConfig {
+    Sse {
+        url: String,
+    },
+    Stdio {
+        command: String,
+        #[serde(default)]
+        args: Vec<String>,
+        #[serde(default)]
+        envs: HashMap<String, String>,
+    },
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct McpConfig {
+    server: Vec<McpServerConfig>,
+}
+
+impl McpConfig {
+    pub async fn create_manager(&self) -> anyhow::Result<McpManager> {
+        let mut clients = HashMap::new();
+        let mut task_set = tokio::task::JoinSet::<anyhow::Result<_>>::new();
+        for server in &self.server {
+            let server = server.clone();
+            task_set.spawn(async move {
+                let client = server.transport.start().await?;
+                anyhow::Result::Ok((server.name.clone(), client))
+            });
+        }
+        let start_up_result = task_set.join_all().await;
+        for result in start_up_result {
+            match result {
+                Ok((name, client)) => {
+                    clients.insert(name, client);
+                }
+                Err(e) => {
+                    eprintln!("Failed to start server: {:?}", e);
+                }
+            }
+        }
+        Ok(McpManager { clients })
+    }
+}
+
+impl McpServerTransportConfig {
+    pub async fn start(&self) -> anyhow::Result<RunningService<RoleClient, ()>> {
+        let client = match self {
+            McpServerTransportConfig::Sse { url } => {
+                let transport = rmcp::transport::SseTransport::start(url).await?;
+                ().serve(transport).await?
+            }
+            McpServerTransportConfig::Stdio {
+                command,
+                args,
+                envs,
+            } => {
+                let transport = rmcp::transport::TokioChildProcess::new(
+                    tokio::process::Command::new(command)
+                        .args(args)
+                        .envs(envs)
+                        .stderr(Stdio::null()),
+                )?;
+                ().serve(transport).await?
+            }
+        };
+        Ok(client)
+    }
+}

examples/rig-integration/src/main.rs

@@ -0,0 +1,64 @@
+use rig::{
+    embeddings::EmbeddingsBuilder,
+    providers::{
+        cohere,
+        deepseek::{self, DEEPSEEK_CHAT},
+    },
+    vector_store::in_memory_store::InMemoryVectorStore,
+};
+pub mod chat;
+pub mod config;
+pub mod mcp_adaptor;
+#[tokio::main]
+async fn main() -> anyhow::Result<()> {
+    let config = config::Config::retrive("config.toml").await?;
+    let openai_client = {
+        if let Some(key) = config.deepseek_key {
+            deepseek::Client::new(&key)
+        } else {
+            deepseek::Client::from_env()
+        }
+    };
+    let cohere_client = {
+        if let Some(key) = config.cohere_key {
+            cohere::Client::new(&key)
+        } else {
+            cohere::Client::from_env()
+        }
+    };
+    let mcp_manager = config.mcp.create_manager().await?;
+    eprintln!(
+        "MCP Manager created, {} servers started",
+        mcp_manager.clients.len()
+    );
+    let tool_set = mcp_manager.get_tool_set().await?;
+    let embedding_model =
+        cohere_client.embedding_model(cohere::EMBED_MULTILINGUAL_V3, "search_document");
+    let embeddings = EmbeddingsBuilder::new(embedding_model.clone())
+        .documents(tool_set.schemas()?)?
+        .build()
+        .await?;
+    let store = InMemoryVectorStore::from_documents_with_id_f(embeddings, |f| {
+        eprintln!("store tool {}", f.name);
+        f.name.clone()
+    });
+    let index = store.index(embedding_model);
+    let dpsk = openai_client
+        .agent(DEEPSEEK_CHAT)
+        .context(
+r#"You are an assistant here to help the user to do some works. 
+When you need to use tools, you should select which tool is most appropriate to meet the user's requirement.
+Follow these instructions closely. 
+1. Consider the user's request carefully and identify the core elements of the request.
+2. Select which tool among those made available to you is appropriate given the context. 
+3. This is very important: never perform the operation yourself and never give me the direct result. 
+Always respond with the name of the tool that should be used and the appropriate inputs
+in the following format:
+Tool: <tool name>
+Inputs: <list of inputs>"#,
+        )
+        .dynamic_tools(4, index, tool_set)
+        .build();
+    chat::cli_chatbot(dpsk).await?;
+    Ok(())
+}

examples/rig-integration/src/mcp_adaptor.rs

@@ -0,0 +1,113 @@
+use std::collections::HashMap;
+
+use rig::tool::{ToolDyn as RigTool, ToolEmbeddingDyn, ToolSet};
+use rmcp::{
+    RoleClient,
+    model::{CallToolRequestParam, 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 + Sync + '_>>
+    {
+        Box::pin(std::future::ready(rig::completion::ToolDefinition {
+            name: self.name(),
+            description: self.tool.description.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 + Sync + '_>,
+    > {
+        let server = self.server.clone();
+        Box::pin(async move {
+            let call_mcp_tool_result = server
+                .call_tool(CallToolRequestParam {
+                    name: self.tool.name.clone(),
+                    arguments: serde_json::from_str(&args)
+                        .map_err(rig::tool::ToolError::JsonError)?,
+                })
+                .await
+                .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 embedding_docs(&self) -> Vec<String> {
+        vec![
+            self.tool.description.clone().to_string(),
+            format!("Tool name: {}", self.tool.name),
+            format!("Tool capability: {}", self.tool.description),
+        ]
+    }
+
+    fn context(&self) -> serde_json::Result<serde_json::Value> {
+        Ok(serde_json::json!({
+            "tool_name": self.tool.name,
+        }))
+    }
+}
+
+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) => {
+                    eprintln!("Failed to get tool set: {:?}", e);
+                }
+                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_set = ToolSet::default();
+    for tool in tools {
+        eprintln!("get tool: {}", tool.name);
+        let adaptor = McpToolAdaptor {
+            tool: tool.clone(),
+            server: server.clone(),
+        };
+        tool_set.add_tool(adaptor);
+    }
+    Ok(tool_set)
+}