add basic test

This commit is contained in:
milan
2026-04-07 21:09:48 +02:00
parent 94811c5a89
commit 1895e1ec66
3 changed files with 3168 additions and 17 deletions
+3106
View File
File diff suppressed because it is too large Load Diff
+60 -16
View File
@@ -1,16 +1,15 @@
use crate::mcp::MCPClient; 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::chat::request::ChatMessageRequest;
use ollama_rs::generation::chat::ChatMessage;
use ollama_rs::generation::tools::ToolInfo; use ollama_rs::generation::tools::ToolInfo;
use rmcp::ServiceError; use ollama_rs::Ollama;
use rmcp::model::{CallToolRequestParams, CallToolResult}; use rmcp::model::{CallToolRequestParams, CallToolResult};
use rmcp::ServiceError;
use std::collections::HashMap; use std::collections::HashMap;
use std::ops::Index;
use thiserror::Error; use thiserror::Error;
#[derive(Debug, Copy, Clone, Default)] #[derive(Debug, Copy, Clone, Default, PartialEq)]
pub enum ToolPermission { pub enum ToolPermission {
Allowed, Allowed,
#[default] #[default]
@@ -138,6 +137,13 @@ impl AgentChat {
/// returns: /// returns:
/// Ok(tuple) => tuple of the server name and the tool or resource name /// Ok(tuple) => tuple of the server name and the tool or resource name
/// Err(error) => ChatError::FunctionParseError(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> { fn parse_tool_name(name: String) -> Result<(String, String), ChatError> {
let (mcp_server_name, tool_name) = name let (mcp_server_name, tool_name) = name
.split_once("::") .split_once("::")
@@ -158,11 +164,10 @@ impl AgentChat {
/// * `name`: must be in the format "mcp_server_name::tool_or_resource_name". /// * `name`: must be in the format "mcp_server_name::tool_or_resource_name".
/// ///
/// returns: Option<&RestrictedTool> /// returns: Option<&RestrictedTool>
pub fn get_tool<'a>( pub fn get_tool(
&self, mpc_server_data: &MCPServerData,
mpc_server_data: &'a MCPServerData,
name: String, name: String,
) -> Option<&'a RestrictedTool> { ) -> Option<&RestrictedTool> {
mpc_server_data mpc_server_data
.translated_tools .translated_tools
.iter() .iter()
@@ -248,11 +253,6 @@ impl AgentChat {
/// ///
/// returns: Result<Option<CallToolResult>, ChatError> /// returns: Result<Option<CallToolResult>, ChatError>
/// if permission was denied: Ok(None) /// if permission was denied: Ok(None)
/// # Examples
///
/// ```
///
/// ```
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: &ollama_rs::generation::tools::ToolCall,
@@ -273,8 +273,7 @@ impl AgentChat {
let mcp_server_data = self let mcp_server_data = self
.get_mcp_server_by_name(mcp_server_name.clone()) .get_mcp_server_by_name(mcp_server_name.clone())
.ok_or(ChatError::ServiceNotFoundError(mcp_server_name.clone()))?; .ok_or(ChatError::ServiceNotFoundError(mcp_server_name.clone()))?;
let restricted_tool = self let restricted_tool = AgentChat::get_tool(&mcp_server_data, full_unparsed_tool_name.clone())
.get_tool(&mcp_server_data, full_unparsed_tool_name.clone())
.ok_or(ChatError::ToolNotFoundError( .ok_or(ChatError::ToolNotFoundError(
full_unparsed_tool_name.clone(), full_unparsed_tool_name.clone(),
))?; ))?;
@@ -314,3 +313,48 @@ impl AgentChat {
Ok(Some(mcp_server_data.client.call_tool(request_params).await?)) 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);
}
}
+2 -1
View File
@@ -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::_private::serde_json::{Map, Value};
use ollama_rs::re_exports::schemars::Schema; use ollama_rs::re_exports::schemars::Schema;
use rmcp::ServiceError; use rmcp::ServiceError;
use crate::mcp::chat::ChatError::ServiceError as ChatServiceError;
pub(in crate::mcp) async fn get_server_tool_info( pub(in crate::mcp) async fn get_server_tool_info(
client: &MCPClient, 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( pub(in crate::mcp) async fn get_server_resource_tool_info(
client: &MCPClient, client: &MCPClient,
server_name: &String, server_name: &String,
@@ -58,6 +58,7 @@ pub(in crate::mcp) async fn get_server_resource_tool_info(
Ok(resource_tool_info) Ok(resource_tool_info)
} }
#[allow(dead_code)]
fn tool_info_from_mcp_resource( fn tool_info_from_mcp_resource(
mcp_resource: &rmcp::model::Resource, mcp_resource: &rmcp::model::Resource,
server_name: &String, server_name: &String,