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 {
|
||||
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)
|
||||
|
||||
+31
-16
@@ -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?,
|
||||
))
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user