milan/own_assist
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add permission support

Browse Source 
refactoring
This commit is contained in:
milan
2026-04-07 17:53:32 +02:00
parent ee17e74753
commit 94811c5a89
8 changed files with 566 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
own_mcp/src/mcp/chat.rs
+147 -44
View File
@@ -5,9 +5,9 @@ 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 rmcp::model::{CallToolRequestParams, CallToolResult};
use std::collections::HashMap;
use std::ops::Index;
use thiserror::Error;
#[derive(Debug, Copy, Clone, Default)]
@@ -18,9 +18,15 @@ pub enum ToolPermission {
Denied,
}
#[derive(Debug, Copy, Clone)]
pub enum PermissionAnswer {
Granted,
Denied,
}
#[derive(Debug, Clone)]
pub struct RestrictedTool {
permission: ToolPermission, // TODO: USE PERMISSION
permission: ToolPermission,
tool_info: ToolInfo,
}
@@ -28,7 +34,7 @@ pub struct RestrictedTool {
pub struct MCPServerData {
name: String,
client: MCPClient,
/// tools in the ollama format. mcp resources are also translated into tools
/// tools in the ollama format. mcp resources are also translated into tools (not yet)
translated_tools: Vec<RestrictedTool>,
}
@@ -43,6 +49,8 @@ impl MCPServerData {
})
}
// Resource into tool translation commented out, since resource reading is not yet solved. `AgentChat` would instead try to call tool with the resource name.
/*
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 {
@@ -56,6 +64,7 @@ impl MCPServerData {
log::warn!("Server `{name}` does not support resources");
}
}
*/
Ok(MCPServerData {
name,
@@ -85,6 +94,8 @@ pub enum ChatError {
ArgumentParsingError,
#[error("service could not be found")]
ServiceNotFoundError(String),
#[error("tool could not be found")]
ToolNotFoundError(String),
}
impl AgentChat {
@@ -92,7 +103,7 @@ impl AgentChat {
ollama_client: Ollama,
model: String,
mcp_clients: HashMap<String, MCPClient>,
system_prompt: String
system_prompt: String,
) -> Result<Self, ServiceError> {
let history = vec![ChatMessage::system(system_prompt)];
@@ -111,22 +122,74 @@ impl AgentChat {
})
}
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 fn get_all_tools(&self) -> impl Iterator<Item = &RestrictedTool> {
self.mcp_servers
.values()
.map(|server| server.translated_tools.iter())
.flatten()
}
pub async fn message(&mut self, user_message: String) -> Result<ChatMessage, ChatError> {
let all_tools = self.get_all_tools();
///
///
/// # Arguments
///
/// * `name`: must be in the format "mcp_server_name::tool_or_resource_name".
///
/// returns:
/// Ok(tuple) => tuple of the server name and the tool or resource name
/// Err(error) => ChatError::FunctionParseError(name)
fn parse_tool_name(name: String) -> Result<(String, String), ChatError> {
let (mcp_server_name, tool_name) = name
.split_once("::")
.ok_or(ChatError::FunctionParseError(name.clone()))?;
log::debug!("all tools: {:#?}", all_tools);
Ok((mcp_server_name.to_string(), tool_name.to_string()))
}
pub fn get_mcp_server_by_name(&self, name: String) -> Option<&MCPServerData> {
self.mcp_servers.get(&name)
}
///
///
/// # Arguments
///
/// * `mpc_server_data`:
/// * `name`: must be in the format "mcp_server_name::tool_or_resource_name".
///
/// returns: Option<&RestrictedTool>
pub fn get_tool<'a>(
&self,
mpc_server_data: &'a MCPServerData,
name: String,
) -> Option<&'a RestrictedTool> {
mpc_server_data
.translated_tools
.iter()
.find(|tool| tool.tool_info.function.name == name)
}
pub fn set_permission(&mut self, full_tool_name: String, permission: ToolPermission) -> Result<(), ChatError> {
let (mcp_server_name, tool_name) = Self::parse_tool_name(full_tool_name.clone())?;
let mcp_server_data: &mut MCPServerData = self.mcp_servers.get_mut(&mcp_server_name.clone()).ok_or(ChatError::ServiceNotFoundError(mcp_server_name.clone()))?;
let tool_index = mcp_server_data.translated_tools.iter().position(|tool| tool.tool_info.function.name == full_tool_name).ok_or(ChatError::ToolNotFoundError(tool_name))?;
mcp_server_data.translated_tools[tool_index].permission = permission;
Ok(())
}
pub async fn message<C: Future<Output = PermissionAnswer>>(
&mut self,
user_message: String,
permission_request_callback: fn(mcp_server_name: String, tool_name: String) -> C,
) -> Result<ChatMessage, ChatError>
{
let all_tools: Vec<ToolInfo> = self
.get_all_tools()
.map(|restricted_tool| restricted_tool.tool_info.clone())
.collect();
log::debug!("all tools: {all_tools:#?}");
let mut response = self
.ollama_client
@@ -145,13 +208,16 @@ impl AgentChat {
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<_>>();
let result = self.call_tool(tool_call, permission_request_callback).await?;
let contents = match result {
Some(result) => result.content
.iter()
.filter_map(|content| content.as_text())
.map(|text_content| text_content.text.clone())
.collect::<Vec<_>>(),
None => vec![String::from("Tool Permission Denied by the user")],
};
log::debug!("contents: {:#?}", contents);
@@ -163,7 +229,8 @@ impl AgentChat {
.ollama_client
.send_chat_messages_with_history(
&mut self.message_history,
ChatMessageRequest::new(self.model.clone(), Vec::new()).tools(all_tools.clone()),
ChatMessageRequest::new(self.model.clone(), Vec::new())
.tools(all_tools.clone()),
)
.await?;
log::debug!("received message: {:#?}", response.message);
@@ -172,10 +239,25 @@ impl AgentChat {
Ok(response.message)
}
async fn call_tool(
///
///
/// # Arguments
///
/// * `tool_call`:
/// * `permission_request_callback`:
///
/// returns: Result<Option<CallToolResult>, ChatError>
/// if permission was denied: Ok(None)
/// # Examples
///
/// ```
///
/// ```
async fn call_tool<C: Future<Output = PermissionAnswer>>(
&self,
tool_call: &ollama_rs::generation::tools::ToolCall,
) -> Result<CallToolResult, ChatError> {
permission_request_callback: fn(mcp_server_name: String, tool_name: String) -> C,
) -> Result<Option<CallToolResult>, ChatError> {
let arguments_json_object = tool_call
.function
.arguments
@@ -184,18 +266,44 @@ impl AgentChat {
.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 full_unparsed_tool_name = &tool_call.function.name; // "mcp_server_name::tool_name"
let (mcp_server_name, tool_name) = Self::parse_tool_name(full_unparsed_tool_name.clone())?;
let mcp_server_data = self
.get_mcp_server_by_name(mcp_server_name.clone())
.ok_or(ChatError::ServiceNotFoundError(mcp_server_name.clone()))?;
let restricted_tool = self
.get_tool(&mcp_server_data, full_unparsed_tool_name.clone())
.ok_or(ChatError::ToolNotFoundError(
full_unparsed_tool_name.clone(),
))?;
match restricted_tool.permission {
ToolPermission::Allowed => {
log::trace!("allowed restricted tool: {:?}", restricted_tool.tool_info);
}
ToolPermission::Denied => {
log::info!("denied restricted tool: {:?}", restricted_tool);
return Ok(None);
}
ToolPermission::Ask => {
log::trace!("ask restricted tool: {:?}", restricted_tool);
let response = permission_request_callback(mcp_server_name, tool_name.clone()).await;
match response {
PermissionAnswer::Granted => {
log::trace!("granted restricted tool: {:?}", restricted_tool);
}
PermissionAnswer::Denied => {
log::info!("denied restricted tool: {:?}", restricted_tool);
return Ok(None);
}
}
}
}
let request_params =
CallToolRequestParams::new(function_name).with_arguments(arguments_json_object.clone());
CallToolRequestParams::new(tool_name).with_arguments(arguments_json_object.clone());
log::debug!(
"calling tool {} with request_params: {:#?}",
@@ -203,11 +311,6 @@ impl AgentChat {
request_params
);
Ok(self
.mcp_servers
.get(&mcp_name)
.ok_or(ChatError::ServiceNotFoundError(mcp_name))?
.client.call_tool(request_params)
.await?)
Ok(Some(mcp_server_data.client.call_tool(request_params).await?))
}
}