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

tramex_tools/ai/
mod.rs
