tramex_tools/ai/
mistral.rs

1//! Mistral AI connector implementation
2
3use crate::ai::{AIConnector, AIRequest};
4use crate::data::{AdditionalInfos, Trace};
5use crate::errors::{ErrorCode, TramexError};
6
7/// System prompt for Mistral AI explaining Amarisoft traces
8const SYSTEM_PROMPT: &str = r#"You are a telecom protocol expert specializing in 4G LTE and 5G NR analysis. You are helping a user understand traces captured from an Amarisoft base station (eNB/gNB).
9
10When explaining a trace, structure your response in three sections:
111. **Message Overview**: What this message is, which protocol layer it belongs to, and its role in the signaling flow.
122. **Key Fields**: Explain the important fields and parameters present in the message. Use precise telecom terminology but provide brief clarifications for non-obvious terms.
133. **Protocol Context**: Where this message fits in the typical protocol procedure (e.g. attach, handover, bearer setup). Mention what typically precedes and follows it.
14
15Be concise but technically accurate. Target approximately 200 words. Use markdown formatting for readability."#;
16
17/// Mistral AI connector
18pub struct MistralConnector {
19    /// API endpoint
20    endpoint: String,
21    /// Model to use
22    model: String,
23}
24
25impl MistralConnector {
26    /// Create a new MistralConnector with default settings
27    pub fn new() -> Self {
28        Self {
29            endpoint: "https://api.mistral.ai/v1/chat/completions".to_string(),
30            model: "mistral-medium-latest".to_string(),
31        }
32    }
33
34    /// Create a MistralConnector with a specific model
35    pub fn with_model(model: &str) -> Self {
36        Self {
37            endpoint: "https://api.mistral.ai/v1/chat/completions".to_string(),
38            model: model.to_string(),
39        }
40    }
41
42    /// Build the user message from a Trace
43    fn build_user_message(trace: &Trace) -> String {
44        let mut parts = Vec::new();
45
46        // Structured metadata
47        parts.push(format!("Layer: {:?}", trace.layer));
48        parts.push(format!("Timestamp: {}", trace.timestamp));
49
50        // Direction and message type from additional infos
51        match &trace.additional_infos {
52            AdditionalInfos::RRCInfos(info) => {
53                parts.push(format!("Direction: {:?}", info.direction));
54                parts.push(format!("Channel/Message: {}", info.canal_msg));
55            }
56            AdditionalInfos::NASInfos(info) => {
57                parts.push(format!("Direction: {:?}", info.direction));
58                parts.push(format!("Message Type: {}", info.message_type));
59            }
60            AdditionalInfos::NGAPInfos(info) => {
61                parts.push(format!("Direction: {:?}", info.direction));
62                parts.push(format!("Message Type: {}", info.message_type));
63            }
64            AdditionalInfos::GTPUInfos(info) => {
65                parts.push(format!("Direction: {:?}", info.direction));
66                parts.push(format!("Message Type: {}", info.message_type));
67            }
68            AdditionalInfos::PHYInfos(info) => {
69                parts.push(format!("Direction: {:?}", info.direction));
70                parts.push(format!("Channel Type: {:?}", info.channel_type));
71            }
72            AdditionalInfos::None => {}
73        }
74
75        // Raw text
76        if let Some(text_lines) = &trace.text {
77            parts.push(String::new());
78            parts.push("Raw trace content:".to_string());
79            parts.push("```".to_string());
80            for line in text_lines {
81                parts.push(line.clone());
82            }
83            parts.push("```".to_string());
84        }
85
86        parts.join("\n")
87    }
88}
89
90impl Default for MistralConnector {
91    fn default() -> Self {
92        Self::new()
93    }
94}
95
96impl AIConnector for MistralConnector {
97    fn name(&self) -> &'static str {
98        "Mistral"
99    }
100
101    fn build_request(&self, trace: &Trace, api_key: &str) -> Result<AIRequest, TramexError> {
102        if api_key.is_empty() {
103            return Err(TramexError::new(
104                "API key is empty. Set it in Settings > AI.".to_string(),
105                ErrorCode::RequestError,
106            ));
107        }
108
109        let user_message = Self::build_user_message(trace);
110
111        let body = serde_json::json!({
112            "model": self.model,
113            "messages": [
114                {
115                    "role": "system",
116                    "content": SYSTEM_PROMPT
117                },
118                {
119                    "role": "user",
120                    "content": user_message
121                }
122            ],
123            "temperature": 0.3,
124            "max_tokens": 1024
125        });
126
127        let body_str = serde_json::to_string(&body)
128            .map_err(|e| TramexError::new(format!("Failed to serialize request body: {e}"), ErrorCode::RequestError))?;
129
130        Ok(AIRequest {
131            url: self.endpoint.clone(),
132            headers: vec![
133                ("Authorization".to_string(), format!("Bearer {api_key}")),
134                ("Content-Type".to_string(), "application/json".to_string()),
135            ],
136            body: body_str,
137        })
138    }
139
140    fn parse_response(&self, response_body: &str) -> Result<String, TramexError> {
141        let json: serde_json::Value = serde_json::from_str(response_body)
142            .map_err(|e| TramexError::new(format!("Failed to parse AI response: {e}"), ErrorCode::RequestError))?;
143
144        // Check for API error (OpenAI-compatible nested or string error)
145        if let Some(error) = json.get("error") {
146            let msg = error
147                .get("message")
148                .and_then(|m| m.as_str())
149                .or_else(|| error.as_str())
150                .unwrap_or("Unknown API error");
151            return Err(TramexError::new(format!("Mistral API error: {msg}"), ErrorCode::RequestError));
152        }
153
154        // Some error responses expose details at the top level
155        if let Some(msg) = json
156            .get("message")
157            .and_then(|m| m.as_str())
158            .or_else(|| json.get("detail").and_then(|d| d.as_str()))
159        {
160            return Err(TramexError::new(format!("Mistral API error: {msg}"), ErrorCode::RequestError));
161        }
162
163        // Extract the assistant's message content
164        json.get("choices")
165            .and_then(|c| c.get(0))
166            .and_then(|c| c.get("message"))
167            .and_then(|m| m.get("content"))
168            .and_then(|c| c.as_str())
169            .map(|s| s.to_string())
170            .ok_or_else(|| {
171                TramexError::new(
172                    "Unexpected response format from Mistral API".to_string(),
173                    ErrorCode::RequestError,
174                )
175            })
176    }
177}
178
179#[cfg(test)]
180mod tests {
181    use super::*;
182
183    #[test]
184    fn parse_successful_response() {
185        let connector = MistralConnector::new();
186        let body = r#"{"choices":[{"message":{"content":"Hello"}}]}"#;
187        assert_eq!(connector.parse_response(body).unwrap(), "Hello");
188    }
189
190    #[test]
191    fn parse_error_with_nested_message() {
192        let connector = MistralConnector::new();
193        let body = r#"{"error":{"message":"Invalid API key","type":"unauthorized"}}"#;
194        let err = connector.parse_response(body).unwrap_err();
195        assert!(err.get_msg().contains("Invalid API key"));
196    }
197
198    #[test]
199    fn parse_error_with_top_level_message() {
200        let connector = MistralConnector::new();
201        let body = r#"{"message":"Unauthorized","request_id":"abc"}"#;
202        let err = connector.parse_response(body).unwrap_err();
203        assert!(err.get_msg().contains("Unauthorized"));
204    }
205}
tramex_tools/ai/mistral.rs

tramex_tools/ai/
mistral.rs
