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
145        if let Some(error) = json.get("error") {
146            let msg = error.get("message").and_then(|m| m.as_str()).unwrap_or("Unknown API error");
147            return Err(TramexError::new(format!("Mistral API error: {msg}"), ErrorCode::RequestError));
148        }
149
150        // Extract the assistant's message content
151        json.get("choices")
152            .and_then(|c| c.get(0))
153            .and_then(|c| c.get("message"))
154            .and_then(|m| m.get("content"))
155            .and_then(|c| c.as_str())
156            .map(|s| s.to_string())
157            .ok_or_else(|| {
158                TramexError::new(
159                    "Unexpected response format from Mistral API".to_string(),
160                    ErrorCode::RequestError,
161                )
162            })
163    }
164}
tramex_tools/ai/mistral.rs

tramex_tools/ai/
mistral.rs
