restructure into multiple crates

add mcp server collection
2026-04-07 00:06:13 +02:00
parent 2dbac8a91a
commit ee17e74753
21 changed files with 2095 additions and 164 deletions
@@ -0,0 +1,213 @@
+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::tools::ToolInfo;
+use rmcp::ServiceError;
+use rmcp::model::{CallToolRequestParams, CallToolResult, };
+use std::collections::HashMap;
+
+use thiserror::Error;
+
+#[derive(Debug, Copy, Clone, Default)]
+pub enum ToolPermission {
+    Allowed,
+    #[default]
+    Ask,
+    Denied,
+}
+
+#[derive(Debug, Clone)]
+pub struct RestrictedTool {
+    permission: ToolPermission, // TODO: USE PERMISSION
+    tool_info: ToolInfo,
+}
+
+#[derive(Debug)]
+pub struct MCPServerData {
+    name: String,
+    client: MCPClient,
+    /// tools in the ollama format. mcp resources are also translated into tools
+    translated_tools: Vec<RestrictedTool>,
+}
+
+impl MCPServerData {
+    pub async fn new(client: MCPClient, name: String) -> Result<Self, ServiceError> {
+        let mut tools: Vec<RestrictedTool> = Vec::new();
+
+        for tool_info in crate::mcp::translation::get_server_tool_info(&client, &name).await? {
+            tools.push(RestrictedTool {
+                permission: Default::default(),
+                tool_info,
+            })
+        }
+
+        match crate::mcp::translation::get_server_resource_tool_info(&client, &name).await {
+            Ok(resources_as_tool_infos) => {
+                for tool_info in resources_as_tool_infos {
+                    tools.push(RestrictedTool {
+                        permission: Default::default(),
+                        tool_info,
+                    })
+                }
+            },
+            Err(_) => {
+                log::warn!("Server `{name}` does not support resources");
+            }
+        }
+
+        Ok(MCPServerData {
+            name,
+            client,
+            translated_tools: tools,
+        })
+    }
+}
+
+#[derive(Debug)]
+pub struct AgentChat {
+    ollama_client: Ollama,
+    model: String,
+    mcp_servers: HashMap<String, MCPServerData>,
+    message_history: Vec<ChatMessage>,
+}
+
+#[derive(Debug, Error)]
+pub enum ChatError {
+    #[error(transparent)]
+    OllamaError(#[from] OllamaError),
+    #[error(transparent)]
+    ServiceError(#[from] ServiceError),
+    #[error("the function name could not be parsed")]
+    FunctionParseError(String),
+    #[error("error parsing provided arguments")]
+    ArgumentParsingError,
+    #[error("service could not be found")]
+    ServiceNotFoundError(String),
+}
+
+impl AgentChat {
+    pub async fn new(
+        ollama_client: Ollama,
+        model: String,
+        mcp_clients: HashMap<String, MCPClient>,
+        system_prompt: String
+    ) -> Result<Self, ServiceError> {
+        let history = vec![ChatMessage::system(system_prompt)];
+
+        let mut servers = HashMap::new();
+
+        for (name, client) in mcp_clients.into_iter() {
+            let server = MCPServerData::new(client, name).await?;
+            servers.insert(server.name.clone(), server);
+        }
+
+        Ok(Self {
+            ollama_client,
+            model,
+            mcp_servers: servers,
+            message_history: history,
+        })
+    }
+
+    pub fn get_all_tools(&self) -> Vec<ToolInfo> {
+        let mut all_tools: Vec<ToolInfo> = Vec::new();
+
+        for server in self.mcp_servers.values() {
+            for tool in server.translated_tools.iter() {
+                all_tools.push(tool.tool_info.clone());
+            }
+        }
+
+        all_tools
+    }
+
+    pub async fn message(&mut self, user_message: String) -> Result<ChatMessage, ChatError> {
+        let all_tools = self.get_all_tools();
+
+        log::debug!("all tools: {:#?}", all_tools);
+
+        let mut response = self
+            .ollama_client
+            .send_chat_messages_with_history(
+                &mut self.message_history,
+                ChatMessageRequest::new(self.model.clone(), vec![ChatMessage::user(user_message)])
+                    .tools(all_tools.clone()),
+            )
+            .await?;
+
+        loop {
+            if response.message.tool_calls.is_empty() {
+                log::trace!("no tool was used for message: {:#?}", response.message);
+                break;
+            }
+
+            for tool_call in &response.message.tool_calls {
+                log::debug!("calling tool {}", tool_call.function.name);
+                let result = self.call_tool(tool_call).await?;
+                let contents = result
+                    .content
+                    .iter()
+                    .filter_map(|content| content.as_text())
+                    .map(|text_content| text_content.text.clone())
+                    .collect::<Vec<_>>();
+
+                log::debug!("contents: {:#?}", contents);
+
+                self.message_history
+                    .push(ChatMessage::tool(contents.join("\n")));
+            }
+
+            response = self
+                .ollama_client
+                .send_chat_messages_with_history(
+                    &mut self.message_history,
+                    ChatMessageRequest::new(self.model.clone(), Vec::new()).tools(all_tools.clone()),
+                )
+                .await?;
+            log::debug!("received message: {:#?}", response.message);
+        }
+
+        Ok(response.message)
+    }
+
+    async fn call_tool(
+        &self,
+        tool_call: &ollama_rs::generation::tools::ToolCall,
+    ) -> Result<CallToolResult, ChatError> {
+        let arguments_json_object = tool_call
+            .function
+            .arguments
+            .as_object()
+            .ok_or(ChatError::ArgumentParsingError)
+            .inspect_err(|_| log::error!("arguments are not of type object"))?;
+        log::trace!("arguments_json_object: {:?}", arguments_json_object);
+
+        let (mcp_name, function_name) =
+            tool_call
+                .function
+                .name
+                .split_once("::")
+                .ok_or(ChatError::FunctionParseError(
+                    tool_call.function.name.clone(),
+                ))?;
+        let (mcp_name, function_name) = (mcp_name.to_string(), function_name.to_string());
+
+        let request_params =
+            CallToolRequestParams::new(function_name).with_arguments(arguments_json_object.clone());
+
+        log::debug!(
+            "calling tool {} with request_params: {:#?}",
+            tool_call.function.name,
+            request_params
+        );
+
+        Ok(self
+            .mcp_servers
+            .get(&mcp_name)
+            .ok_or(ChatError::ServiceNotFoundError(mcp_name))?
+            .client.call_tool(request_params)
+            .await?)
+    }
+}