chore: remove the rig example (#763)
The rig official has the mcp example, we need not to keep it just give the link. Signed-off-by: jokemanfire <hu.dingyang@zte.com.cn>
This commit is contained in:
jokemanfire
GitHub
parent
3ea8c3c555
commit
30cdc38c9a
8 changed files with 1 additions and 470 deletions
|
|
@ -70,7 +70,7 @@ see [servers/README.md](servers/README.md)
|
||||||
|
|
||||||
# Integration
|
# Integration
|
||||||
|
|
||||||
- [Rig](rig-integration) A stream chatbot with rig
|
- [Rig](https://github.com/0xPlaygrounds/rig/blob/main/rig/rig-core/examples/rmcp.rs) A stream chatbot with rig
|
||||||
- [Simple Chat Client](simple-chat-client) A simple chat client implementation using the Model Context Protocol (MCP) SDK.
|
- [Simple Chat Client](simple-chat-client) A simple chat client implementation using the Model Context Protocol (MCP) SDK.
|
||||||
|
|
||||||
# WASI
|
# WASI
|
||||||
|
|
|
||||||
|
|
@ -1,34 +0,0 @@
|
||||||
[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 }
|
|
||||||
publish = false
|
|
||||||
|
|
||||||
[dependencies]
|
|
||||||
rig-core = "0.32.0"
|
|
||||||
tokio = { version = "1", features = ["full"] }
|
|
||||||
rmcp = { workspace = true, features = [
|
|
||||||
"client",
|
|
||||||
"transport-child-process",
|
|
||||||
"transport-streamable-http-client-reqwest"
|
|
||||||
] }
|
|
||||||
anyhow = "1.0"
|
|
||||||
serde_json = "1"
|
|
||||||
serde = { version = "1", features = ["derive"] }
|
|
||||||
toml = "1.0"
|
|
||||||
futures = "0.3"
|
|
||||||
tracing = "0.1"
|
|
||||||
tracing-subscriber = { version = "0.3", features = [
|
|
||||||
"env-filter",
|
|
||||||
"std",
|
|
||||||
"fmt",
|
|
||||||
] }
|
|
||||||
tracing-appender = "0.2"
|
|
||||||
|
|
@ -1,10 +0,0 @@
|
||||||
deepseek_key = ""
|
|
||||||
cohere_key = ""
|
|
||||||
|
|
||||||
[mcp]
|
|
||||||
|
|
||||||
[[mcp.server]]
|
|
||||||
name = "git"
|
|
||||||
protocol = "stdio"
|
|
||||||
command = "uvx"
|
|
||||||
args = ["mcp-server-git"]
|
|
||||||
|
|
@ -1,134 +0,0 @@
|
||||||
use futures::StreamExt;
|
|
||||||
use rig::{
|
|
||||||
agent::{Agent, MultiTurnStreamItem},
|
|
||||||
completion::CompletionModel,
|
|
||||||
message::{Message, Text},
|
|
||||||
streaming::{StreamedAssistantContent, StreamingChat},
|
|
||||||
};
|
|
||||||
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter};
|
|
||||||
|
|
||||||
pub async fn cli_chatbot<M>(chatbot: Agent<M>) -> anyhow::Result<()>
|
|
||||||
where
|
|
||||||
M: CompletionModel + 'static,
|
|
||||||
M::StreamingResponse: Send,
|
|
||||||
{
|
|
||||||
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;
|
|
||||||
}
|
|
||||||
|
|
||||||
tracing::info!(%input);
|
|
||||||
chat_log.push(Message::user(input));
|
|
||||||
|
|
||||||
let mut response = chatbot.stream_chat(input, chat_log.clone()).await;
|
|
||||||
stream_output_agent_start(&mut output).await?;
|
|
||||||
let mut message_buf = String::new();
|
|
||||||
|
|
||||||
while let Some(message) = response.next().await {
|
|
||||||
match message {
|
|
||||||
Ok(MultiTurnStreamItem::StreamAssistantItem(StreamedAssistantContent::Text(
|
|
||||||
Text { text },
|
|
||||||
))) => {
|
|
||||||
message_buf.push_str(&text);
|
|
||||||
output_agent(&text, &mut output).await?;
|
|
||||||
}
|
|
||||||
Ok(MultiTurnStreamItem::StreamAssistantItem(
|
|
||||||
StreamedAssistantContent::ToolCall { tool_call, .. },
|
|
||||||
)) => {
|
|
||||||
let name = &tool_call.function.name;
|
|
||||||
let arguments = &tool_call.function.arguments;
|
|
||||||
stream_output_toolcall(
|
|
||||||
format!("Calling tool: {name} with args: {arguments}"),
|
|
||||||
&mut output,
|
|
||||||
)
|
|
||||||
.await?;
|
|
||||||
}
|
|
||||||
Ok(MultiTurnStreamItem::StreamUserItem(user_content)) => {
|
|
||||||
// Tool results are streamed back as user items
|
|
||||||
stream_output_toolcall(format!("Tool result: {:?}", user_content), &mut output)
|
|
||||||
.await?;
|
|
||||||
}
|
|
||||||
Ok(MultiTurnStreamItem::FinalResponse(final_response)) => {
|
|
||||||
tracing::info!("Final response received: {:?}", final_response);
|
|
||||||
}
|
|
||||||
Ok(_) => {
|
|
||||||
// Handle other stream items (reasoning, deltas, etc.)
|
|
||||||
}
|
|
||||||
Err(error) => {
|
|
||||||
output_error(error, &mut output).await?;
|
|
||||||
}
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
||||||
chat_log.push(Message::assistant(message_buf));
|
|
||||||
stream_output_agent_finished(&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[1;31m\xE2\x9D\x8C ERROR: \x1b[0m")
|
|
||||||
.await?;
|
|
||||||
output.write_all(e.to_string().as_bytes()).await?;
|
|
||||||
output.write_all(b"\n").await?;
|
|
||||||
output.flush().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[1;33m\xF0\x9F\x9B\xA0 Tool Call: \x1b[0m")
|
|
||||||
.await?;
|
|
||||||
output.write_all(content.to_string().as_bytes()).await?;
|
|
||||||
output.write_all(b"\n").await?;
|
|
||||||
output.flush().await?;
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
pub async fn stream_output_agent_start(
|
|
||||||
output: &mut BufWriter<tokio::io::Stdout>,
|
|
||||||
) -> std::io::Result<()> {
|
|
||||||
output
|
|
||||||
.write_all(b"\x1b[1;34m\xF0\x9F\xA4\x96 Agent: \x1b[0m")
|
|
||||||
.await?;
|
|
||||||
output.flush().await?;
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
||||||
pub async fn stream_output_agent_finished(
|
|
||||||
output: &mut BufWriter<tokio::io::Stdout>,
|
|
||||||
) -> std::io::Result<()> {
|
|
||||||
output.write_all(b"\n").await?;
|
|
||||||
output.flush().await?;
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
@ -1,20 +0,0 @@
|
||||||
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 retrieve(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)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
@ -1,83 +0,0 @@
|
||||||
use std::{collections::HashMap, process::Stdio};
|
|
||||||
|
|
||||||
use rmcp::{RoleClient, ServiceExt, service::RunningService, transport::ConfigureCommandExt};
|
|
||||||
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 {
|
|
||||||
Streamable {
|
|
||||||
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::Streamable { url } => {
|
|
||||||
let transport =
|
|
||||||
rmcp::transport::StreamableHttpClientTransport::from_uri(url.to_string());
|
|
||||||
().serve(transport).await?
|
|
||||||
}
|
|
||||||
McpServerTransportConfig::Stdio {
|
|
||||||
command,
|
|
||||||
args,
|
|
||||||
envs,
|
|
||||||
} => {
|
|
||||||
let transport = rmcp::transport::TokioChildProcess::new(
|
|
||||||
tokio::process::Command::new(command).configure(|cmd| {
|
|
||||||
cmd.args(args).envs(envs).stderr(Stdio::null());
|
|
||||||
}),
|
|
||||||
)?;
|
|
||||||
().serve(transport).await?
|
|
||||||
}
|
|
||||||
};
|
|
||||||
Ok(client)
|
|
||||||
}
|
|
||||||
}
|
|
||||||
|
|
@ -1,69 +0,0 @@
|
||||||
use rig::{
|
|
||||||
client::{CompletionClient, ProviderClient},
|
|
||||||
embeddings::EmbeddingsBuilder,
|
|
||||||
providers::{cohere, deepseek},
|
|
||||||
vector_store::in_memory_store::InMemoryVectorStore,
|
|
||||||
};
|
|
||||||
use tracing_appender::rolling::{RollingFileAppender, Rotation};
|
|
||||||
pub mod chat;
|
|
||||||
pub mod config;
|
|
||||||
pub mod mcp_adaptor;
|
|
||||||
|
|
||||||
#[tokio::main]
|
|
||||||
async fn main() -> anyhow::Result<()> {
|
|
||||||
let file_appender = RollingFileAppender::new(
|
|
||||||
Rotation::DAILY,
|
|
||||||
"logs",
|
|
||||||
format!("{}.log", env!("CARGO_CRATE_NAME")),
|
|
||||||
);
|
|
||||||
tracing_subscriber::fmt()
|
|
||||||
.with_env_filter(
|
|
||||||
tracing_subscriber::EnvFilter::from_default_env()
|
|
||||||
.add_directive(tracing::Level::INFO.into()),
|
|
||||||
)
|
|
||||||
.with_writer(file_appender)
|
|
||||||
.with_file(false)
|
|
||||||
.with_ansi(false)
|
|
||||||
.init();
|
|
||||||
|
|
||||||
let config = config::Config::retrieve("config.toml").await?;
|
|
||||||
let deepseek_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?;
|
|
||||||
tracing::info!(
|
|
||||||
"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| {
|
|
||||||
tracing::info!("store tool {}", f.name);
|
|
||||||
f.name.clone()
|
|
||||||
});
|
|
||||||
let index = store.index(embedding_model);
|
|
||||||
let dpsk = deepseek_client
|
|
||||||
.agent(deepseek::DEEPSEEK_CHAT)
|
|
||||||
.dynamic_tools(4, index, tool_set)
|
|
||||||
.build();
|
|
||||||
|
|
||||||
chat::cli_chatbot(dpsk).await?;
|
|
||||||
|
|
||||||
Ok(())
|
|
||||||
}
|
|
||||||
|
|
@ -1,119 +0,0 @@
|
||||||
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)
|
|
||||||
}
|
|
||||||
Loading…
Reference in a new issue