add basic test
This commit is contained in:
+60
-16
@@ -1,16 +1,15 @@
|
||||
use crate::mcp::MCPClient;
|
||||
use ollama_rs::Ollama;
|
||||
use ollama_rs::error::OllamaError;
|
||||
use ollama_rs::generation::chat::ChatMessage;
|
||||
use ollama_rs::generation::chat::request::ChatMessageRequest;
|
||||
use ollama_rs::generation::chat::ChatMessage;
|
||||
use ollama_rs::generation::tools::ToolInfo;
|
||||
use rmcp::ServiceError;
|
||||
use ollama_rs::Ollama;
|
||||
use rmcp::model::{CallToolRequestParams, CallToolResult};
|
||||
use rmcp::ServiceError;
|
||||
use std::collections::HashMap;
|
||||
use std::ops::Index;
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Debug, Copy, Clone, Default)]
|
||||
#[derive(Debug, Copy, Clone, Default, PartialEq)]
|
||||
pub enum ToolPermission {
|
||||
Allowed,
|
||||
#[default]
|
||||
@@ -138,6 +137,13 @@ impl AgentChat {
|
||||
/// returns:
|
||||
/// Ok(tuple) => tuple of the server name and the tool or resource name
|
||||
/// Err(error) => ChatError::FunctionParseError(name)
|
||||
///
|
||||
/// # Examples
|
||||
///
|
||||
/// ```ignore
|
||||
/// let result = AgentChat::parse_tool_name("example::get_foo".to_string());
|
||||
/// assert_eq!(result.unwrap(), ("example".to_string(), "get_foo".to_string()))
|
||||
/// ```
|
||||
fn parse_tool_name(name: String) -> Result<(String, String), ChatError> {
|
||||
let (mcp_server_name, tool_name) = name
|
||||
.split_once("::")
|
||||
@@ -158,11 +164,10 @@ impl AgentChat {
|
||||
/// * `name`: must be in the format "mcp_server_name::tool_or_resource_name".
|
||||
///
|
||||
/// returns: Option<&RestrictedTool>
|
||||
pub fn get_tool<'a>(
|
||||
&self,
|
||||
mpc_server_data: &'a MCPServerData,
|
||||
pub fn get_tool(
|
||||
mpc_server_data: &MCPServerData,
|
||||
name: String,
|
||||
) -> Option<&'a RestrictedTool> {
|
||||
) -> Option<&RestrictedTool> {
|
||||
mpc_server_data
|
||||
.translated_tools
|
||||
.iter()
|
||||
@@ -248,11 +253,6 @@ impl AgentChat {
|
||||
///
|
||||
/// returns: Result<Option<CallToolResult>, ChatError>
|
||||
/// if permission was denied: Ok(None)
|
||||
/// # Examples
|
||||
///
|
||||
/// ```
|
||||
///
|
||||
/// ```
|
||||
async fn call_tool<C: Future<Output = PermissionAnswer>>(
|
||||
&self,
|
||||
tool_call: &ollama_rs::generation::tools::ToolCall,
|
||||
@@ -273,8 +273,7 @@ impl AgentChat {
|
||||
let mcp_server_data = self
|
||||
.get_mcp_server_by_name(mcp_server_name.clone())
|
||||
.ok_or(ChatError::ServiceNotFoundError(mcp_server_name.clone()))?;
|
||||
let restricted_tool = self
|
||||
.get_tool(&mcp_server_data, full_unparsed_tool_name.clone())
|
||||
let restricted_tool = AgentChat::get_tool(&mcp_server_data, full_unparsed_tool_name.clone())
|
||||
.ok_or(ChatError::ToolNotFoundError(
|
||||
full_unparsed_tool_name.clone(),
|
||||
))?;
|
||||
@@ -314,3 +313,48 @@ impl AgentChat {
|
||||
Ok(Some(mcp_server_data.client.call_tool(request_params).await?))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use crate::mcp::chat::RestrictedTool;
|
||||
use crate::mcp::get_client;
|
||||
use rmcp::model::Implementation;
|
||||
use std::collections::HashMap;
|
||||
|
||||
const MCP_TEST_SERVER_URL: &str = "https://mcpplaygroundonline.com/mcp-echo-server";
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_permission_setting() {
|
||||
let ollama = Ollama::default();
|
||||
|
||||
println!(
|
||||
"Test requires ollama running on {} and connection to {MCP_TEST_SERVER_URL}",
|
||||
ollama.url()
|
||||
);
|
||||
|
||||
let model_name = "lfm2.5-thinking:1.2b".into();
|
||||
|
||||
ollama.pull_model(model_name, false).await.unwrap();
|
||||
|
||||
let mcp_clients = HashMap::from([(
|
||||
"test".to_string(),
|
||||
get_client(
|
||||
MCP_TEST_SERVER_URL,
|
||||
None::<String>,
|
||||
Implementation::new("test", env!("CARGO_PKG_VERSION")),
|
||||
).await.unwrap(),
|
||||
)]);
|
||||
|
||||
let mut chat = AgentChat::new(ollama, "lfm2.5-thinking:1.2b".to_string(), mcp_clients, String::new()).await.unwrap();
|
||||
|
||||
let tools: Vec<RestrictedTool> = chat.get_all_tools().cloned().collect();
|
||||
assert_eq!(tools.len(), 1);
|
||||
|
||||
assert_eq!(tools[0].permission, ToolPermission::Ask);
|
||||
chat.set_permission("test::echo".to_string(), ToolPermission::Allowed).unwrap();
|
||||
|
||||
let tools: Vec<RestrictedTool> = chat.get_all_tools().cloned().collect();
|
||||
assert_eq!(tools[0].permission, ToolPermission::Allowed);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -3,7 +3,6 @@ use ollama_rs::generation::tools::{ToolFunctionInfo, ToolInfo, ToolType};
|
||||
use ollama_rs::re_exports::schemars::_private::serde_json::{Map, Value};
|
||||
use ollama_rs::re_exports::schemars::Schema;
|
||||
use rmcp::ServiceError;
|
||||
use crate::mcp::chat::ChatError::ServiceError as ChatServiceError;
|
||||
|
||||
pub(in crate::mcp) async fn get_server_tool_info(
|
||||
client: &MCPClient,
|
||||
@@ -44,6 +43,7 @@ fn tool_info_from_mcp_tool(mcp_tool: &rmcp::model::Tool, server_name: &String) -
|
||||
}
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
pub(in crate::mcp) async fn get_server_resource_tool_info(
|
||||
client: &MCPClient,
|
||||
server_name: &String,
|
||||
@@ -58,6 +58,7 @@ pub(in crate::mcp) async fn get_server_resource_tool_info(
|
||||
Ok(resource_tool_info)
|
||||
}
|
||||
|
||||
#[allow(dead_code)]
|
||||
fn tool_info_from_mcp_resource(
|
||||
mcp_resource: &rmcp::model::Resource,
|
||||
server_name: &String,
|
||||
|
||||
Reference in New Issue
Block a user