tramex_tools/ai/
anthropic.rs

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

tramex_tools/ai/
anthropic.rs
