diff --git a/cli/src/main.rs b/cli/src/main.rs index 9ffed2d..8ee8d94 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -44,7 +44,7 @@ async fn chat_loop(mut human_interface: impl HumanInterface, mut agent_chat: Age loop { let user_message = human_interface.expect_user_message().await.unwrap(); let immutable_interface = &human_interface; - let agent_message = agent_chat + let (agent_message, _) = agent_chat .message(user_message, async |mcp_server_name, tool_name| { immutable_interface .ask_for_permission(mcp_server_name, tool_name) diff --git a/own_mcp/src/mcp/chat.rs b/own_mcp/src/mcp/chat.rs index 826d485..3a2ca11 100644 --- a/own_mcp/src/mcp/chat.rs +++ b/own_mcp/src/mcp/chat.rs @@ -1,10 +1,10 @@ use crate::mcp::MCPClient; use ollama_rs::Ollama; -use ollama_rs::error::OllamaError; +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::ToolInfo; +use ollama_rs::generation::tools::{ToolCall, ToolInfo}; use rmcp::ServiceError; use rmcp::model::{CallToolRequestParams, CallToolResult}; use serde::Deserialize; @@ -99,6 +99,18 @@ pub enum ChatError { 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, @@ -205,7 +217,7 @@ impl AgentChat { &mut self, user_message: String, mut permission_request_callback: impl FnMut(String, String) -> C, - ) -> Result { + ) -> Result<(ChatMessage, Vec), ChatError> { let all_tools: Vec = self .get_all_tools() .filter_map(|restricted_tool| { @@ -229,6 +241,8 @@ impl AgentChat { ) .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); @@ -237,27 +251,28 @@ impl AgentChat { for tool_call in &response.message.tool_calls { log::debug!("calling tool {}", tool_call.function.name); - let result = self + 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 Some(structured_content) = - result.clone().and_then(|result| result.structured_content) + 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 result { - Some(result) => result + 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::>(), - None => vec![String::from("Tool Permission Denied by the user")], + RestrictedToolCallResult::Denied => vec![String::from("Tool Permission Denied by the user")], }; log::debug!("contents: {contents:#?}"); @@ -278,7 +293,7 @@ impl AgentChat { log::debug!("received message: {:#?}", response.message); } - Ok(response.message) + Ok((response.message, tool_usages)) } /// @@ -292,9 +307,9 @@ impl AgentChat { /// if permission was denied: Ok(None) async fn call_tool>( &self, - tool_call: &ollama_rs::generation::tools::ToolCall, + tool_call: &ToolCall, permission_request_callback: &mut impl FnMut(String, String) -> C, - ) -> Result, ChatError> { + ) -> Result { let arguments_json_object = tool_call .function .arguments @@ -303,7 +318,7 @@ impl AgentChat { .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 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())?; @@ -321,7 +336,7 @@ impl AgentChat { } ToolPermission::Denied => { log::info!("denied restricted tool: {:?}", restricted_tool); - return Ok(None); + return Ok(RestrictedToolCallResult::Denied); } ToolPermission::Ask => { log::trace!("ask restricted tool: {:?}", restricted_tool); @@ -333,7 +348,7 @@ impl AgentChat { } PermissionAnswer::Denied => { log::info!("denied restricted tool: {:?}", restricted_tool); - return Ok(None); + return Ok(RestrictedToolCallResult::Denied); } } } @@ -348,7 +363,7 @@ impl AgentChat { request_params ); - Ok(Some( + Ok(RestrictedToolCallResult::Granted( mcp_server_data.client.call_tool(request_params).await?, )) }