tramex_tools/ai/
mod.rs

1//! AI connector module for trace explanation
2//!
3//! Provides a trait-based abstraction for AI chatbot APIs.
4//! Enable with the `ai` feature flag.
5
6pub mod anthropic;
7pub mod mistral;
8pub mod openai;
9
10use crate::data::Trace;
11use crate::errors::TramexError;
12
13/// Represents an HTTP request to be sent to an AI API
14#[derive(Debug, Clone)]
15pub struct AIRequest {
16    /// The full URL endpoint
17    pub url: String,
18    /// HTTP headers as (key, value) pairs
19    pub headers: Vec<(String, String)>,
20    /// JSON body as a string
21    pub body: String,
22}
23
24/// Status of an AI explanation request
25#[derive(Debug, Clone, Default)]
26pub enum AIExplainStatus {
27    /// No request has been made
28    #[default]
29    Idle,
30    /// Request is in flight
31    Loading,
32    /// Response received successfully
33    Done(String),
34    /// Request failed
35    Error(String),
36}
37
38/// Trait abstracting an AI chatbot connector.
39///
40/// Implementations build the HTTP request and parse the response.
41/// The actual HTTP call is handled by the UI layer (using ehttp).
42pub trait AIConnector: Send + Sync {
43    /// Human-readable name of the AI provider
44    fn name(&self) -> &'static str;
45
46    /// Build the HTTP request for explaining a trace.
47    ///
48    /// # Arguments
49    /// * `trace` - The trace to explain
50    /// * `api_key` - The API key for authentication
51    ///
52    /// # Errors
53    /// Returns an error if the request cannot be built
54    fn build_request(&self, trace: &Trace, api_key: &str) -> Result<AIRequest, TramexError>;
55
56    /// Parse the API response body into a human-readable explanation.
57    ///
58    /// # Arguments
59    /// * `response_body` - The raw JSON response from the API
60    ///
61    /// # Errors
62    /// Returns an error if the response cannot be parsed
63    fn parse_response(&self, response_body: &str) -> Result<String, TramexError>;
64}
65
66/// Available AI provider types
67#[derive(Default, Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
68pub enum AIProvider {
69    /// Mistral AI
70    #[default]
71    Mistral,
72    /// OpenAI (ChatGPT)
73    OpenAI,
74    /// Anthropic (Claude)
75    Anthropic,
76}
77
78impl std::fmt::Display for AIProvider {
79    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
80        match self {
81            Self::Mistral => write!(f, "Mistral"),
82            Self::OpenAI => write!(f, "OpenAI"),
83            Self::Anthropic => write!(f, "Anthropic"),
84        }
85    }
86}
87
88impl AIProvider {
89    /// Return the list of all available providers
90    pub fn all() -> &'static [AIProvider] {
91        &[AIProvider::Mistral, AIProvider::OpenAI, AIProvider::Anthropic]
92    }
93
94    /// Return available (display_name, model_id) pairs for this provider
95    pub fn available_models(&self) -> &'static [(&'static str, &'static str)] {
96        match self {
97            AIProvider::Mistral => &[
98                ("Small", "mistral-small-latest"),
99                ("Medium", "mistral-medium-latest"),
100                ("Large", "mistral-large-latest"),
101            ],
102            AIProvider::OpenAI => &[
103                ("GPT-4o mini", "gpt-4o-mini"),
104                ("GPT-4o", "gpt-4o"),
105                ("GPT-4.1", "gpt-4.1"),
106                ("GPT-5.3", "gpt-5.3"),
107                ("GPT-5.5", "gpt-5.5"),
108            ],
109            AIProvider::Anthropic => &[
110                ("Claude 3.5 Haiku", "claude-3-5-haiku-latest"),
111                ("Claude 3.5 Sonnet", "claude-3-5-sonnet-latest"),
112                ("Claude 3 Opus", "claude-3-opus-latest"),
113                ("Claude 4 Opus", "claude-4-opus-latest"),
114                ("Claude 5 Opus", "claude-5-opus-latest"),
115            ],
116        }
117    }
118
119    /// Return the default model ID for this provider
120    pub fn default_model(&self) -> &'static str {
121        match self {
122            AIProvider::Mistral => "mistral-medium-latest",
123            AIProvider::OpenAI => "gpt-4o-mini",
124            AIProvider::Anthropic => "claude-3-5-haiku-latest",
125        }
126    }
127}
128
129/// Create an AIConnector from a provider type and model ID
130pub fn create_connector(provider: &AIProvider, model: &str) -> Box<dyn AIConnector> {
131    match provider {
132        AIProvider::Mistral => Box::new(mistral::MistralConnector::with_model(model)),
133        AIProvider::OpenAI => Box::new(openai::OpenAIConnector::with_model(model)),
134        AIProvider::Anthropic => Box::new(anthropic::AnthropicConnector::with_model(model)),
135    }
136}
tramex_tools/ai/mod.rs

tramex_tools/ai/
mod.rs
