agent chat now returns tool usage info

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<ChatMessage, ChatError> {
+    ) -> Result<(ChatMessage, Vec<ToolUsage>), ChatError> {
         let all_tools: Vec<ToolInfo> = 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::<Vec<_>>(),
-                        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<C: Future<Output = PermissionAnswer>>(
         &self,
-        tool_call: &ollama_rs::generation::tools::ToolCall,
+        tool_call: &ToolCall,
         permission_request_callback: &mut impl FnMut(String, String) -> C,
-    ) -> Result<Option<CallToolResult>, ChatError> {
+    ) -> Result<RestrictedToolCallResult, ChatError> {
         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?,
         ))
     }