own_assist/own_mcp/src/mcp/chat.rs

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::completion::request::GenerationRequest;
use ollama_rs::generation::tools::{ToolCall, ToolInfo};
use rmcp::ServiceError;
use rmcp::model::{CallToolRequestParams, CallToolResult};
use serde::Deserialize;
use std::collections::HashMap;
use thiserror::Error;

#[derive(Debug, Copy, Clone, Default, Deserialize, PartialEq)]
pub enum ToolPermission {
    Allowed,
    #[default]
    Ask,
    Denied,
}

#[derive(Debug, Copy, Clone)]
pub enum PermissionAnswer {
    Granted,
    Denied,
}

#[derive(Debug, Clone)]
pub struct RestrictedTool {
    permission: ToolPermission,
    tool_info: ToolInfo,
}

#[derive(Debug)]
pub struct MCPServerData {
    name: String,
    client: MCPClient,
    /// tools in the ollama format. mcp resources are also translated into tools (not yet)
    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,
            })
        }

        // 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 {
                    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),
    #[error("tool could not be found")]
    ToolNotFoundError(String),
}

#[derive(Debug, Clone)]
pub enum RestrictedToolCallResult {
    Granted(CallToolResult),
    Denied,
}

#[derive(Debug, Clone)]
pub struct ToolUsage {
    pub call: ToolCall,
    pub result: RestrictedToolCallResult,
}

impl AgentChat {
    pub async fn new(
        ollama_client: Ollama,
        model: String,
        mcp_clients: HashMap<String, MCPClient>,
    ) -> Result<Self, ServiceError> {
        let history = vec![];

        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);
        }

        // lets the ollama server load the model
        let ollama_client_clone = ollama_client.clone();
        let model_clone = model.clone();
        tokio::spawn(async move {
            let _ = ollama_client_clone
                .generate(GenerationRequest::new(model_clone, ""))
                .await;
        });

        Ok(Self {
            ollama_client,
            model,
            mcp_servers: servers,
            message_history: history,
        })
    }

    pub fn get_all_tools(&self) -> impl Iterator<Item = &RestrictedTool> {
        self.mcp_servers
            .values()
            .flat_map(|server| server.translated_tools.iter())
    }

    ///
    ///
    /// # 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)
    ///
    /// # Examples
    ///
    /// ```ignore
    /// let result = AgentChat::parse_tool_name("example:get_foo".to_string());
    /// assert_eq!(result.unwrap(), ("example".to_string(), "get_foo".to_string()))
    /// ```
    pub 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()))?;

        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(mpc_server_data: &MCPServerData, name: String) -> Option<&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(())
    }

    // TODO: ChatMessage only returns last tool usage, which is always empty. instead the tool usages should be collected and returned in total
    pub async fn message<C: Future<Output = PermissionAnswer>>(
        &mut self,
        user_message: String,
        mut permission_request_callback: impl FnMut(String, String) -> C,
    ) -> Result<(ChatMessage, Vec<ToolUsage>), ChatError> {
        let all_tools: Vec<ToolInfo> = self
            .get_all_tools()
            .filter_map(|restricted_tool| {
                if restricted_tool.permission == ToolPermission::Denied {
                    // filters denied tools to save tokens
                    None
                } else {
                    Some(restricted_tool.tool_info.clone())
                }
            })
            .collect();

        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?;

        let mut tool_usages = Vec::new();

        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 restricted_result = self
                    .call_tool(tool_call, &mut permission_request_callback)
                    .await?;

                tool_usages.push(ToolUsage {
                    call: tool_call.clone(),
                    result: restricted_result.clone(),
                });

                // serialize structured content if it exists
                if let RestrictedToolCallResult::Granted(result) = restricted_result.clone()
                    && let Some(structured_content) = result.structured_content
                {
                    log::debug!("structured content: {structured_content:#?}");

                    self.message_history
                        .push(ChatMessage::tool(structured_content.to_string()));
                } else {
                    let contents = match restricted_result {
                        RestrictedToolCallResult::Granted(result) => result
                            .content
                            .iter()
                            .filter_map(|content| content.as_text())
                            .map(|text_content| text_content.text.clone())
                            .collect::<Vec<_>>(),
                        RestrictedToolCallResult::Denied => {
                            vec![String::from("Tool Permission Denied by the user")]
                        }
                    };

                    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, tool_usages))
    }

    ///
    ///
    /// # Arguments
    ///
    /// * `tool_call`:
    /// * `permission_request_callback`:
    ///
    /// returns: Result<Option<CallToolResult>, ChatError>
    /// if permission was denied: Ok(None)
    async fn call_tool<C: Future<Output = PermissionAnswer>>(
        &self,
        tool_call: &ToolCall,
        permission_request_callback: &mut impl FnMut(String, String) -> C,
    ) -> Result<RestrictedToolCallResult, 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 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 = AgentChat::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(RestrictedToolCallResult::Denied);
            }
            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(RestrictedToolCallResult::Denied);
                    }
                }
            }
        }

        let request_params =
            CallToolRequestParams::new(tool_name).with_arguments(arguments_json_object.clone());

        log::debug!(
            "calling tool {} with request_params: {:#?}",
            tool_call.function.name,
            request_params
        );

        Ok(RestrictedToolCallResult::Granted(
            mcp_server_data.client.call_tool(request_params).await?,
        ))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::mcp::chat::RestrictedTool;
    use crate::mcp::get_client;
    use rmcp::model::Implementation;
    use std::collections::HashMap;

    const MCP_TEST_SERVER_URL: &str = "https://mcpplaygroundonline.com/mcp-echo-server";

    #[tokio::test]
    async fn test_permission_setting() {
        let ollama = Ollama::default();

        println!(
            "Test requires ollama running on {} and connection to {MCP_TEST_SERVER_URL}",
            ollama.url()
        );

        let model_name = "lfm2.5-thinking:1.2b".into();

        ollama.pull_model(model_name, false).await.unwrap();

        let mcp_clients = HashMap::from([(
            "test".to_string(),
            get_client(
                MCP_TEST_SERVER_URL,
                None::<String>,
                Implementation::new("test", env!("CARGO_PKG_VERSION")),
            )
            .await
            .unwrap(),
        )]);

        let mut chat = AgentChat::new(ollama, "lfm2.5-thinking:1.2b".to_string(), mcp_clients)
            .await
            .unwrap();

        let tools: Vec<RestrictedTool> = chat.get_all_tools().cloned().collect();
        assert_eq!(tools.len(), 1);

        assert_eq!(tools[0].permission, ToolPermission::Ask);
        chat.set_permission("test:echo".to_string(), ToolPermission::Allowed)
            .unwrap();

        let tools: Vec<RestrictedTool> = chat.get_all_tools().cloned().collect();
        assert_eq!(tools[0].permission, ToolPermission::Allowed);
    }
}