agent chat now returns tool usage info
This commit is contained in:
+1
-1
@@ -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
@@ -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?,
|
||||||
))
|
))
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user