restructure into multiple crates
add mcp server collection
This commit is contained in:
milan
@@ -0,0 +1,3 @@
|
||||
pub mod mcp;
|
||||
|
||||
pub use mcp::chat::AgentChat;
|
||||
@@ -0,0 +1,213 @@
|
||||
use crate::mcp::MCPClient;
|
||||
use ollama_rs::Ollama;
|
||||
use ollama_rs::error::OllamaError;
|
||||
use ollama_rs::generation::chat::ChatMessage;
|
||||
use ollama_rs::generation::chat::request::ChatMessageRequest;
|
||||
use ollama_rs::generation::tools::ToolInfo;
|
||||
use rmcp::ServiceError;
|
||||
use rmcp::model::{CallToolRequestParams, CallToolResult, };
|
||||
use std::collections::HashMap;
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Debug, Copy, Clone, Default)]
|
||||
pub enum ToolPermission {
|
||||
Allowed,
|
||||
#[default]
|
||||
Ask,
|
||||
Denied,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct RestrictedTool {
|
||||
permission: ToolPermission, // TODO: USE PERMISSION
|
||||
tool_info: ToolInfo,
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct MCPServerData {
|
||||
name: String,
|
||||
client: MCPClient,
|
||||
/// tools in the ollama format. mcp resources are also translated into tools
|
||||
translated_tools: Vec<RestrictedTool>,
|
||||
}
|
||||
|
||||
impl MCPServerData {
|
||||
pub async fn new(client: MCPClient, name: String) -> Result<Self, ServiceError> {
|
||||
let mut tools: Vec<RestrictedTool> = Vec::new();
|
||||
|
||||
for tool_info in crate::mcp::translation::get_server_tool_info(&client, &name).await? {
|
||||
tools.push(RestrictedTool {
|
||||
permission: Default::default(),
|
||||
tool_info,
|
||||
})
|
||||
}
|
||||
|
||||
match crate::mcp::translation::get_server_resource_tool_info(&client, &name).await {
|
||||
Ok(resources_as_tool_infos) => {
|
||||
for tool_info in resources_as_tool_infos {
|
||||
tools.push(RestrictedTool {
|
||||
permission: Default::default(),
|
||||
tool_info,
|
||||
})
|
||||
}
|
||||
},
|
||||
Err(_) => {
|
||||
log::warn!("Server `{name}` does not support resources");
|
||||
}
|
||||
}
|
||||
|
||||
Ok(MCPServerData {
|
||||
name,
|
||||
client,
|
||||
translated_tools: tools,
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct AgentChat {
|
||||
ollama_client: Ollama,
|
||||
model: String,
|
||||
mcp_servers: HashMap<String, MCPServerData>,
|
||||
message_history: Vec<ChatMessage>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum ChatError {
|
||||
#[error(transparent)]
|
||||
OllamaError(#[from] OllamaError),
|
||||
#[error(transparent)]
|
||||
ServiceError(#[from] ServiceError),
|
||||
#[error("the function name could not be parsed")]
|
||||
FunctionParseError(String),
|
||||
#[error("error parsing provided arguments")]
|
||||
ArgumentParsingError,
|
||||
#[error("service could not be found")]
|
||||
ServiceNotFoundError(String),
|
||||
}
|
||||
|
||||
impl AgentChat {
|
||||
pub async fn new(
|
||||
ollama_client: Ollama,
|
||||
model: String,
|
||||
mcp_clients: HashMap<String, MCPClient>,
|
||||
system_prompt: String
|
||||
) -> Result<Self, ServiceError> {
|
||||
let history = vec![ChatMessage::system(system_prompt)];
|
||||
|
||||
let mut servers = HashMap::new();
|
||||
|
||||
for (name, client) in mcp_clients.into_iter() {
|
||||
let server = MCPServerData::new(client, name).await?;
|
||||
servers.insert(server.name.clone(), server);
|
||||
}
|
||||
|
||||
Ok(Self {
|
||||
ollama_client,
|
||||
model,
|
||||
mcp_servers: servers,
|
||||
message_history: history,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn get_all_tools(&self) -> Vec<ToolInfo> {
|
||||
let mut all_tools: Vec<ToolInfo> = Vec::new();
|
||||
|
||||
for server in self.mcp_servers.values() {
|
||||
for tool in server.translated_tools.iter() {
|
||||
all_tools.push(tool.tool_info.clone());
|
||||
}
|
||||
}
|
||||
|
||||
all_tools
|
||||
}
|
||||
|
||||
pub async fn message(&mut self, user_message: String) -> Result<ChatMessage, ChatError> {
|
||||
let all_tools = self.get_all_tools();
|
||||
|
||||
log::debug!("all tools: {:#?}", all_tools);
|
||||
|
||||
let mut response = self
|
||||
.ollama_client
|
||||
.send_chat_messages_with_history(
|
||||
&mut self.message_history,
|
||||
ChatMessageRequest::new(self.model.clone(), vec![ChatMessage::user(user_message)])
|
||||
.tools(all_tools.clone()),
|
||||
)
|
||||
.await?;
|
||||
|
||||
loop {
|
||||
if response.message.tool_calls.is_empty() {
|
||||
log::trace!("no tool was used for message: {:#?}", response.message);
|
||||
break;
|
||||
}
|
||||
|
||||
for tool_call in &response.message.tool_calls {
|
||||
log::debug!("calling tool {}", tool_call.function.name);
|
||||
let result = self.call_tool(tool_call).await?;
|
||||
let contents = result
|
||||
.content
|
||||
.iter()
|
||||
.filter_map(|content| content.as_text())
|
||||
.map(|text_content| text_content.text.clone())
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
log::debug!("contents: {:#?}", contents);
|
||||
|
||||
self.message_history
|
||||
.push(ChatMessage::tool(contents.join("\n")));
|
||||
}
|
||||
|
||||
response = self
|
||||
.ollama_client
|
||||
.send_chat_messages_with_history(
|
||||
&mut self.message_history,
|
||||
ChatMessageRequest::new(self.model.clone(), Vec::new()).tools(all_tools.clone()),
|
||||
)
|
||||
.await?;
|
||||
log::debug!("received message: {:#?}", response.message);
|
||||
}
|
||||
|
||||
Ok(response.message)
|
||||
}
|
||||
|
||||
async fn call_tool(
|
||||
&self,
|
||||
tool_call: &ollama_rs::generation::tools::ToolCall,
|
||||
) -> Result<CallToolResult, ChatError> {
|
||||
let arguments_json_object = tool_call
|
||||
.function
|
||||
.arguments
|
||||
.as_object()
|
||||
.ok_or(ChatError::ArgumentParsingError)
|
||||
.inspect_err(|_| log::error!("arguments are not of type object"))?;
|
||||
log::trace!("arguments_json_object: {:?}", arguments_json_object);
|
||||
|
||||
let (mcp_name, function_name) =
|
||||
tool_call
|
||||
.function
|
||||
.name
|
||||
.split_once("::")
|
||||
.ok_or(ChatError::FunctionParseError(
|
||||
tool_call.function.name.clone(),
|
||||
))?;
|
||||
let (mcp_name, function_name) = (mcp_name.to_string(), function_name.to_string());
|
||||
|
||||
let request_params =
|
||||
CallToolRequestParams::new(function_name).with_arguments(arguments_json_object.clone());
|
||||
|
||||
log::debug!(
|
||||
"calling tool {} with request_params: {:#?}",
|
||||
tool_call.function.name,
|
||||
request_params
|
||||
);
|
||||
|
||||
Ok(self
|
||||
.mcp_servers
|
||||
.get(&mcp_name)
|
||||
.ok_or(ChatError::ServiceNotFoundError(mcp_name))?
|
||||
.client.call_tool(request_params)
|
||||
.await?)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,57 @@
|
||||
pub mod chat;
|
||||
mod translation;
|
||||
|
||||
use rmcp::model::InitializeRequestParams;
|
||||
use rmcp::service::{ClientInitializeError, RunningService};
|
||||
use rmcp::transport::streamable_http_client::StreamableHttpClientTransportConfig;
|
||||
use rmcp::{
|
||||
RoleClient, ServiceExt,
|
||||
model::{ClientCapabilities, ClientInfo, Implementation},
|
||||
transport::StreamableHttpClientTransport,
|
||||
};
|
||||
use std::sync::Arc;
|
||||
|
||||
pub type MCPClient = RunningService<RoleClient, InitializeRequestParams>;
|
||||
|
||||
pub async fn get_client(
|
||||
uri: impl Into<Arc<str>>,
|
||||
authentication: Option<impl Into<String>>,
|
||||
implementation: Implementation,
|
||||
) -> Result<MCPClient, ClientInitializeError> {
|
||||
let mut config = StreamableHttpClientTransportConfig::with_uri(uri);
|
||||
|
||||
if let Some(authentication) = authentication {
|
||||
config = config.auth_header(authentication);
|
||||
}
|
||||
|
||||
let client_info = ClientInfo::new(ClientCapabilities::default(), implementation);
|
||||
|
||||
let transport = StreamableHttpClientTransport::from_config(config);
|
||||
|
||||
client_info.serve(transport).await.inspect_err(|e| {
|
||||
log::error!("client error: {:?}", e);
|
||||
})
|
||||
}
|
||||
|
||||
fn generate_random_mcp_server_name() -> String {
|
||||
use rand::distr::Alphanumeric;
|
||||
use rand::RngExt;
|
||||
|
||||
let mut rng = rand::rng();
|
||||
|
||||
(0..5)
|
||||
.map(|_| rng.sample(Alphanumeric) as char)
|
||||
.collect()
|
||||
}
|
||||
|
||||
pub fn guaranteed_mcp_server_name(user_specified_name: Option<String>, client: &MCPClient) -> String {
|
||||
if let Some(user_specified_name) = user_specified_name {
|
||||
user_specified_name
|
||||
} else if let Some(peer_info) = client.peer_info() {
|
||||
peer_info.server_info.name.clone()
|
||||
} else {
|
||||
let random_name = generate_random_mcp_server_name();
|
||||
log::warn!("no name was specified by the user or the client, so a random name had to be generated: `{}` for {:#?}", random_name, client);
|
||||
random_name
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
use crate::mcp::MCPClient;
|
||||
use ollama_rs::generation::tools::{ToolFunctionInfo, ToolInfo, ToolType};
|
||||
use ollama_rs::re_exports::schemars::_private::serde_json::{Map, Value};
|
||||
use ollama_rs::re_exports::schemars::Schema;
|
||||
use rmcp::ServiceError;
|
||||
use crate::mcp::chat::ChatError::ServiceError as ChatServiceError;
|
||||
|
||||
pub(in crate::mcp) async fn get_server_tool_info(
|
||||
client: &MCPClient,
|
||||
server_name: &String,
|
||||
) -> Result<Vec<ToolInfo>, ServiceError> {
|
||||
let tool_info = client
|
||||
.list_all_tools()
|
||||
.await?
|
||||
.iter()
|
||||
.map(|mcp_tool| tool_info_from_mcp_tool(&mcp_tool, server_name))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Ok(tool_info)
|
||||
}
|
||||
|
||||
fn tool_info_from_mcp_tool(mcp_tool: &rmcp::model::Tool, server_name: &String) -> ToolInfo {
|
||||
let mut value_map: Map<String, Value> = Map::new();
|
||||
|
||||
for key in mcp_tool.input_schema.keys() {
|
||||
let value = mcp_tool.input_schema.get(key).unwrap().clone();
|
||||
value_map.insert(key.clone(), value);
|
||||
}
|
||||
|
||||
let schema_value = Value::Object(value_map);
|
||||
let schema = Schema::try_from(schema_value).unwrap();
|
||||
|
||||
ToolInfo {
|
||||
tool_type: ToolType::Function,
|
||||
function: ToolFunctionInfo {
|
||||
name: format!("{}::{}", server_name, mcp_tool.name),
|
||||
description: mcp_tool
|
||||
.description
|
||||
.clone()
|
||||
.unwrap_or(std::borrow::Cow::from(""))
|
||||
.to_string(),
|
||||
parameters: schema,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
pub(in crate::mcp) async fn get_server_resource_tool_info(
|
||||
client: &MCPClient,
|
||||
server_name: &String,
|
||||
) -> Result<Vec<ToolInfo>, ServiceError> {
|
||||
let resource_tool_info = client
|
||||
.list_all_resources()
|
||||
.await?
|
||||
.iter()
|
||||
.map(|resource| tool_info_from_mcp_resource(resource, server_name))
|
||||
.collect::<Vec<_>>();
|
||||
|
||||
Ok(resource_tool_info)
|
||||
}
|
||||
|
||||
fn tool_info_from_mcp_resource(
|
||||
mcp_resource: &rmcp::model::Resource,
|
||||
server_name: &String,
|
||||
) -> ToolInfo {
|
||||
let schema = Schema::try_from(Value::Object(Map::new())).unwrap();
|
||||
|
||||
log::debug!("resource schema: {:#?}", mcp_resource);
|
||||
|
||||
ToolInfo {
|
||||
tool_type: ToolType::Function,
|
||||
function: ToolFunctionInfo {
|
||||
name: format!("{}::get_{}", server_name, mcp_resource.name),
|
||||
description: format!(
|
||||
"type: {} - {}",
|
||||
mcp_resource
|
||||
.mime_type.clone()
|
||||
.unwrap_or("No type provided".to_string()),
|
||||
mcp_resource.clone()
|
||||
.description.clone()
|
||||
.unwrap_or("No description".to_string())
|
||||
),
|
||||
parameters: schema,
|
||||
},
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user