add modelin in telemetry

Author: evanliu048

crates/chat-cli/src/api_client/clients/client.rs

@@ -78,6 +78,7 @@ impl Client {
         telemetry_event: TelemetryEvent,
         user_context: UserContext,
         telemetry_enabled: bool,
+        model_id: Option<String>,
     ) -> Result<(), ApiClientError> {
         match &self.inner {
             inner::Inner::Codewhisperer(client) => {
@@ -90,6 +91,7 @@ impl Client {
                         false => OptOutPreference::OptOut,
                     })
                     .set_profile_arn(self.profile.as_ref().map(|p| p.arn.clone()))
+                    .set_model_id(model_id)
                     .send()
                     .await;
                 Ok(())
@@ -159,6 +161,7 @@ mod tests {
                     .build()
                     .unwrap(),
                 false,
+                Some("model".to_owned()),
             )
             .await
             .unwrap();

crates/chat-cli/src/cli/chat/mod.rs

@@ -3017,7 +3017,7 @@ impl ChatContext {
                 queue!(self.output, style::Print("\n"))?;
                 let labels: Vec<&str> = MODEL_OPTIONS.iter().map(|(l, _)| *l).collect();
                 let selection: Option<_> = match Select::with_theme(&crate::util::dialoguer_theme())
-                    .with_prompt("Select the model you want to use for chat")
+                    .with_prompt("Select a model for this chat session")
                     .items(&labels)
                     .default(0)
                     .interact_on_opt(&dialoguer::console::Term::stdout())
@@ -3031,13 +3031,16 @@ impl ChatContext {
                 if let Some(index) = selection {
                     let (label, model_opt) = MODEL_OPTIONS[index];
 
+                    self.conversation_state.current_model_id = Some(model_opt.to_string());
+                    telemetry.update_model_id(self.conversation_state.current_model_id.clone());
+
                     use crossterm::{
                         queue,
                         style,
                     };
                     queue!(self.output, style::Print("\n"), style::Print(format!(" Swtiched model to {}\n\n", label)))?;
                 }
-
+                
                 ChatState::PromptUser {
                     tool_uses: None,
                     pending_tool_index: None,

crates/chat-cli/src/telemetry/mod.rs

@@ -6,6 +6,7 @@ mod install_method;
 
 use core::ToolUseEventBuilder;
 use std::str::FromStr;
+use std::sync::{Arc, RwLock};
 
 use amzn_codewhisperer_client::types::{
     ChatAddMessageEvent,
@@ -131,20 +132,24 @@ impl TelemetryStage {
 pub struct TelemetryThread {
     handle: Option<JoinHandle<()>>,
     tx: mpsc::UnboundedSender<Event>,
+    current_model_id: Arc<RwLock<Option<String>>>,
 }
 
 impl Clone for TelemetryThread {
     fn clone(&self) -> Self {
         Self {
             handle: None,
             tx: self.tx.clone(),
+            current_model_id: Arc::clone(&self.current_model_id),
         }
     }
 }
 
 impl TelemetryThread {
     pub async fn new(env: &Env, database: &mut Database) -> Result<Self, TelemetryError> {
-        let telemetry_client = TelemetryClient::new(env, database).await?;
+        //todo yifan change to last used modelid
+        let current_model_id = Arc::new(RwLock::new(None));
+        let telemetry_client = TelemetryClient::new(env, database, current_model_id.clone()).await?;
         let (tx, mut rx) = mpsc::unbounded_channel();
         let handle = tokio::spawn(async move {
             while let Some(event) = rx.recv().await {
@@ -156,9 +161,14 @@ impl TelemetryThread {
         Ok(Self {
             handle: Some(handle),
             tx,
+            current_model_id,
         })
     }
 
+    pub fn update_model_id(&self, model_id: Option<String>) {
+        *self.current_model_id.write().unwrap() = model_id;
+    }
+
     pub async fn finish(self) -> Result<(), TelemetryError> {
         drop(self.tx);
         if let Some(handle) = self.handle {
@@ -275,10 +285,11 @@ struct TelemetryClient {
     telemetry_enabled: bool,
     codewhisperer_client: CodewhispererClient,
     toolkit_telemetry_client: Option<ToolkitTelemetryClient>,
+    current_model_id: Arc<RwLock<Option<String>>>,
 }
 
 impl TelemetryClient {
-    async fn new(env: &Env, database: &mut Database) -> Result<Self, TelemetryError> {
+    async fn new(env: &Env, database: &mut Database, current_model_id: Arc<RwLock<Option<String>>>) -> Result<Self, TelemetryError> {
         let telemetry_enabled = !cfg!(test)
             && env.get_os("Q_DISABLE_TELEMETRY").is_none()
             && database.settings.get_bool(Setting::TelemetryEnabled).unwrap_or(true);
@@ -338,6 +349,7 @@ impl TelemetryClient {
             telemetry_enabled,
             toolkit_telemetry_client,
             codewhisperer_client: CodewhispererClient::new(database, None).await?,
+            current_model_id,
         })
     }
 
@@ -356,6 +368,7 @@ impl TelemetryClient {
             ..
         } = &event.ty
         {
+            let model_id = self.current_model_id.read().unwrap().clone();
             let user_context = self.user_context().unwrap();
 
             let chat_add_message_event = match ChatAddMessageEvent::builder()
@@ -376,6 +389,7 @@ impl TelemetryClient {
                     TelemetryEvent::ChatAddMessageEvent(chat_add_message_event),
                     user_context,
                     self.telemetry_enabled,
+                    model_id,
                 )
                 .await
             {
@@ -451,7 +465,7 @@ mod test {
     #[tokio::test]
     async fn client_context() {
         let mut database = Database::new().await.unwrap();
-        let client = TelemetryClient::new(&Env::new(), &mut database).await.unwrap();
+        let client = TelemetryClient::new(&Env::new(), &mut database, Arc::new(RwLock::new(Some("model".to_owned())))).await.unwrap();
         let context = client.user_context().unwrap();
 
         assert_eq!(context.ide_category, IdeCategory::Cli);
@@ -511,7 +525,7 @@ mod test {
     #[ignore = "needs auth which is not in CI"]
     async fn test_without_optout() {
         let mut database = Database::new().await.unwrap();
-        let client = TelemetryClient::new(&Env::new(), &mut database).await.unwrap();
+        let client = TelemetryClient::new(&Env::new(), &mut database, Arc::new(RwLock::new(Some("model".to_owned())))).await.unwrap();
         client
             .codewhisperer_client
             .send_telemetry_event(
@@ -524,6 +538,7 @@ mod test {
                 ),
                 client.user_context().unwrap(),
                 false,
+                Some("model".to_owned()),
             )
             .await
             .unwrap();