agent chat now returns tool usage info

This commit is contained in:
2026-05-25 15:17:36 +02:00
parent 9ba1c3bae3
commit 8532039199
2 changed files with 32 additions and 17 deletions
+1 -1
View File
@@ -44,7 +44,7 @@ async fn chat_loop(mut human_interface: impl HumanInterface, mut agent_chat: Age
loop { loop {
let user_message = human_interface.expect_user_message().await.unwrap(); let user_message = human_interface.expect_user_message().await.unwrap();
let immutable_interface = &human_interface; let immutable_interface = &human_interface;
let agent_message = agent_chat let (agent_message, _) = agent_chat
.message(user_message, async |mcp_server_name, tool_name| { .message(user_message, async |mcp_server_name, tool_name| {
immutable_interface immutable_interface
.ask_for_permission(mcp_server_name, tool_name) .ask_for_permission(mcp_server_name, tool_name)
+31 -16
View File
@@ -1,10 +1,10 @@
use crate::mcp::MCPClient; use crate::mcp::MCPClient;
use ollama_rs::Ollama; 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::ChatMessage;
use ollama_rs::generation::chat::request::ChatMessageRequest; use ollama_rs::generation::chat::request::ChatMessageRequest;
use ollama_rs::generation::completion::request::GenerationRequest; 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::ServiceError;
use rmcp::model::{CallToolRequestParams, CallToolResult}; use rmcp::model::{CallToolRequestParams, CallToolResult};
use serde::Deserialize; use serde::Deserialize;
@@ -99,6 +99,18 @@ pub enum ChatError {
ToolNotFoundError(String), 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 { impl AgentChat {
pub async fn new( pub async fn new(
ollama_client: Ollama, ollama_client: Ollama,
@@ -205,7 +217,7 @@ impl AgentChat {
&mut self, &mut self,
user_message: String, user_message: String,
mut permission_request_callback: impl FnMut(String, String) -> C, mut permission_request_callback: impl FnMut(String, String) -> C,
) -> Result<ChatMessage, ChatError> { ) -> Result<(ChatMessage, Vec<ToolUsage>), ChatError> {
let all_tools: Vec<ToolInfo> = self let all_tools: Vec<ToolInfo> = self
.get_all_tools() .get_all_tools()
.filter_map(|restricted_tool| { .filter_map(|restricted_tool| {
@@ -229,6 +241,8 @@ impl AgentChat {
) )
.await?; .await?;
let mut tool_usages = Vec::new();
loop { loop {
if response.message.tool_calls.is_empty() { if response.message.tool_calls.is_empty() {
log::trace!("no tool was used for message: {:#?}", response.message); log::trace!("no tool was used for message: {:#?}", response.message);
@@ -237,27 +251,28 @@ impl AgentChat {
for tool_call in &response.message.tool_calls { for tool_call in &response.message.tool_calls {
log::debug!("calling tool {}", tool_call.function.name); log::debug!("calling tool {}", tool_call.function.name);
let result = self let restricted_result = self
.call_tool(tool_call, &mut permission_request_callback) .call_tool(tool_call, &mut permission_request_callback)
.await?; .await?;
tool_usages.push(ToolUsage{call: tool_call.clone(), result: restricted_result.clone()});
// serialize structured content if it exists // serialize structured content if it exists
if let Some(structured_content) = if let RestrictedToolCallResult::Granted(result) = restricted_result.clone() && let Some(structured_content) = result.structured_content
result.clone().and_then(|result| result.structured_content)
{ {
log::debug!("structured content: {structured_content:#?}"); log::debug!("structured content: {structured_content:#?}");
self.message_history self.message_history
.push(ChatMessage::tool(structured_content.to_string())); .push(ChatMessage::tool(structured_content.to_string()));
} else { } else {
let contents = match result { let contents = match restricted_result {
Some(result) => result RestrictedToolCallResult::Granted(result) => result
.content .content
.iter() .iter()
.filter_map(|content| content.as_text()) .filter_map(|content| content.as_text())
.map(|text_content| text_content.text.clone()) .map(|text_content| text_content.text.clone())
.collect::<Vec<_>>(), .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:#?}"); log::debug!("contents: {contents:#?}");
@@ -278,7 +293,7 @@ impl AgentChat {
log::debug!("received message: {:#?}", response.message); 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) /// if permission was denied: Ok(None)
async fn call_tool<C: Future<Output = PermissionAnswer>>( async fn call_tool<C: Future<Output = PermissionAnswer>>(
&self, &self,
tool_call: &ollama_rs::generation::tools::ToolCall, tool_call: &ToolCall,
permission_request_callback: &mut impl FnMut(String, String) -> C, permission_request_callback: &mut impl FnMut(String, String) -> C,
) -> Result<Option<CallToolResult>, ChatError> { ) -> Result<RestrictedToolCallResult, ChatError> {
let arguments_json_object = tool_call let arguments_json_object = tool_call
.function .function
.arguments .arguments
@@ -303,7 +318,7 @@ impl AgentChat {
.inspect_err(|_| log::error!("arguments are not of type object"))?; .inspect_err(|_| log::error!("arguments are not of type object"))?;
log::trace!("arguments_json_object: {:?}", arguments_json_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())?; let (mcp_server_name, tool_name) = Self::parse_tool_name(full_unparsed_tool_name.clone())?;
@@ -321,7 +336,7 @@ impl AgentChat {
} }
ToolPermission::Denied => { ToolPermission::Denied => {
log::info!("denied restricted tool: {:?}", restricted_tool); log::info!("denied restricted tool: {:?}", restricted_tool);
return Ok(None); return Ok(RestrictedToolCallResult::Denied);
} }
ToolPermission::Ask => { ToolPermission::Ask => {
log::trace!("ask restricted tool: {:?}", restricted_tool); log::trace!("ask restricted tool: {:?}", restricted_tool);
@@ -333,7 +348,7 @@ impl AgentChat {
} }
PermissionAnswer::Denied => { PermissionAnswer::Denied => {
log::info!("denied restricted tool: {:?}", restricted_tool); log::info!("denied restricted tool: {:?}", restricted_tool);
return Ok(None); return Ok(RestrictedToolCallResult::Denied);
} }
} }
} }
@@ -348,7 +363,7 @@ impl AgentChat {
request_params request_params
); );
Ok(Some( Ok(RestrictedToolCallResult::Granted(
mcp_server_data.client.call_tool(request_params).await?, mcp_server_data.client.call_tool(request_params).await?,
)) ))
} }