own_assist/own_mcp/src/mcp/llmclient/mistral/mod.rs

use crate::mcp::llmclient::{LLMClient, LLMError};
use async_trait::async_trait;
use ollama_rs::generation::chat::ChatMessage as OllamaChatMessage;
use ollama_rs::generation::chat::request::ChatMessageRequest as OllamaChatMessageRequest;
use ollama_rs::generation::chat::{ChatMessageResponse, MessageRole};
use ollama_rs::generation::tools::{ToolCall, ToolInfo};
use ollama_rs::history::ChatHistory;
use serde::{Deserialize, Serialize};

#[derive(Debug, Clone)]
pub struct Mistral {
    api_key: String,
    system_prompt: Option<String>,
}

impl Mistral {
    pub fn new(api_key: String) -> Mistral {
        Mistral {
            api_key,
            system_prompt: None,
        }
    }

    pub fn with_system_prompt(mut self, system_prompt: Option<String>) -> Mistral {
        self.system_prompt = system_prompt;
        self
    }
}

#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct Message {
    content: String,
    role: MessageRole,
    #[serde(
        skip_serializing_if = "Vec::is_empty",
        default = "Vec::new",
        deserialize_with = "parse_tool_calls"
    )]
    tool_calls: Vec<ToolCall>,
}

fn parse_tool_calls<'de, D>(d: D) -> Result<Vec<ToolCall>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    Deserialize::deserialize(d).map(|x: Option<_>| x.unwrap_or(Vec::new()))
}

impl From<OllamaChatMessage> for Message {
    fn from(message: OllamaChatMessage) -> Message {
        Message {
            content: message.content,
            role: message.role,
            tool_calls: message.tool_calls,
        }
    }
}

impl From<Message> for OllamaChatMessage {
    fn from(message: Message) -> OllamaChatMessage {
        OllamaChatMessage {
            role: message.role,
            content: message.content,
            tool_calls: message.tool_calls,
            images: None,
            thinking: None,
        }
    }
}

#[derive(Debug, Serialize)]
struct ChatCompletionRequest {
    messages: Vec<Message>,
    model: String,
    tools: Vec<ToolInfo>,
}

impl From<OllamaChatMessageRequest> for ChatCompletionRequest {
    fn from(request: OllamaChatMessageRequest) -> ChatCompletionRequest {
        ChatCompletionRequest {
            messages: request.messages.into_iter().map(|m| m.into()).collect(),
            model: request.model_name,
            tools: request.tools,
        }
    }
}

#[derive(Debug, Clone, Deserialize)]
struct ChatCompletionChoice {
    message: Message,
}

#[derive(Debug, Clone, Deserialize)]
struct ChatCompletionResponse {
    choices: Vec<ChatCompletionChoice>,
    model: String,
    #[serde(rename = "created")]
    created_unix_seconds: i64,
}

impl TryFrom<ChatCompletionResponse> for ChatMessageResponse {
    type Error = LLMError;

    fn try_from(value: ChatCompletionResponse) -> Result<ChatMessageResponse, LLMError> {
        let created = chrono::DateTime::from_timestamp(value.created_unix_seconds, 0)
            .ok_or(LLMError::Other("failed converting timestamp".to_string()))?;

        Ok(ChatMessageResponse {
            model: value.model,
            created_at: created.to_rfc3339(),
            message: value
                .choices
                .first()
                .ok_or(LLMError::Other("answer does not exist".to_string()))?
                .message
                .clone()
                .into(),
            logprobs: None,
            done: false,
            final_data: None,
        })
    }
}

#[async_trait]
impl LLMClient for Mistral {
    async fn send_chat_messages_with_history(
        &self,
        history: &mut Vec<OllamaChatMessage>,
        mut request: ollama_rs::generation::chat::request::ChatMessageRequest,
    ) -> Result<ChatMessageResponse, LLMError> {
        if let Some(first_message) = history.first() // add system prompt as first message
            && first_message.role == MessageRole::System
            && let Some(system_prompt) = self.system_prompt.clone()
        {
            history.insert(
                0,
                OllamaChatMessage::new(MessageRole::System, system_prompt),
            );
        }

        history.append(&mut request.messages);

        request.messages = history.messages().to_vec();

        let client = reqwest::Client::new();
        let chat_completion_request: ChatCompletionRequest = request.into();
        let response = client
            .post("https://api.mistral.ai/v1/chat/completions")
            .header("Authorization", format!("Bearer {}", self.api_key))
            .json(&chat_completion_request)
            .send()
            .await
            .map_err(|e| LLMError::NetworkError(e.to_string()))?;

        let response = response
            .error_for_status()
            .map_err(|e| LLMError::Other(e.to_string()))?;

        let completion_response: ChatCompletionResponse = response
            .json()
            .await
            .map_err(|e| LLMError::Other(e.to_string()))?;
        let chat_message_response: ChatMessageResponse = completion_response.try_into()?;

        Ok(chat_message_response)
    }

    async fn preload_model(&self, _model: &str) {
        // do nothing
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ollama_rs::generation::chat::request::ChatMessageRequest;
    use ollama_rs::generation::chat::{ChatMessage, MessageRole};
    use rmcp::schemars;

    #[derive(schemars::JsonSchema)]
    struct BeepToolParameters {
        location: String,
    }

    #[tokio::test]
    async fn sending_message() {
        let client = Mistral::new(include_str!(".MISTRAL_API_KEY").to_string());

        let mut history = vec![];

        dbg!(
            client
                .send_chat_messages_with_history(
                    &mut history,
                    ChatMessageRequest::new(
                        "mistral-medium-3-5".to_string(),
                        vec![ChatMessage::new(
                            MessageRole::User,
                            "Mach ein Geräusch in der Küche".to_string()
                        )],
                    )
                    .tools(vec![ToolInfo {
                        tool_type: ollama_rs::generation::tools::ToolType::Function,
                        function: ollama_rs::generation::tools::ToolFunctionInfo {
                            name: "make_sound".to_string(),
                            description: "alerts the user with a beep".to_string(),
                            parameters: schemars::schema_for!(BeepToolParameters),
                        }
                    }]),
                )
                .await
                .unwrap()
        );
    }
}