merge

Author: evanliu048

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

@@ -82,7 +82,7 @@ impl Client {
     ) -> Result<(), ApiClientError> {
         match &self.inner {
             inner::Inner::Codewhisperer(client) => {
-                let _ = client
+                client
                     .send_telemetry_event()
                     .telemetry_event(telemetry_event)
                     .user_context(user_context)
@@ -93,7 +93,7 @@ impl Client {
                     .set_profile_arn(self.profile.as_ref().map(|p| p.arn.clone()))
                     .set_model_id(model_id)
                     .send()
-                    .await;
+                    .await?;
                 Ok(())
             },
             inner::Inner::Mock => Ok(()),

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

@@ -8,11 +8,7 @@ use aws_credential_types::provider::ProvideCredentials;
 use aws_types::SdkConfig;
 use aws_types::sdk_config::StalledStreamProtectionConfig;
 
-use crate::api_client::credentials::CredentialsChain;
-use crate::api_client::{
-    ApiClientError,
-    Endpoint,
-};
+use crate::api_client::Endpoint;
 use crate::aws_common::behavior_version;
 use crate::database::Database;
 use crate::database::settings::Setting;
@@ -59,13 +55,3 @@ pub async fn bearer_sdk_config(database: &Database, endpoint: &Endpoint) -> SdkC
     let credentials = Credentials::new("xxx", "xxx", None, None, "xxx");
     base_sdk_config(database, endpoint.region().clone(), credentials).await
 }
-
-pub async fn sigv4_sdk_config(database: &Database, endpoint: &Endpoint) -> Result<SdkConfig, ApiClientError> {
-    let credentials_chain = CredentialsChain::new().await;
-
-    if let Err(err) = credentials_chain.provide_credentials().await {
-        return Err(ApiClientError::Credentials(err));
-    };
-
-    Ok(base_sdk_config(database, endpoint.region().clone(), credentials_chain).await)
-}

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

@@ -4,7 +4,6 @@ use std::sync::{
 };
 
 use amzn_codewhisperer_streaming_client::Client as CodewhispererStreamingClient;
-use amzn_qdeveloper_streaming_client::Client as QDeveloperStreamingClient;
 use aws_types::request_id::RequestId;
 use tracing::{
     debug,
@@ -13,7 +12,6 @@ use tracing::{
 
 use super::shared::{
     bearer_sdk_config,
-    sigv4_sdk_config,
     stalled_stream_protection_config,
 };
 use crate::api_client::interceptor::opt_out::OptOutInterceptor;
@@ -42,14 +40,12 @@ mod inner {
     };
 
     use amzn_codewhisperer_streaming_client::Client as CodewhispererStreamingClient;
-    use amzn_qdeveloper_streaming_client::Client as QDeveloperStreamingClient;
 
     use crate::api_client::model::ChatResponseStream;
 
     #[derive(Clone, Debug)]
     pub enum Inner {
         Codewhisperer(CodewhispererStreamingClient),
-        QDeveloper(QDeveloperStreamingClient),
         Mock(Arc<Mutex<std::vec::IntoIter<Vec<ChatResponseStream>>>>),
     }
 }
@@ -62,15 +58,7 @@ pub struct StreamingClient {
 
 impl StreamingClient {
     pub async fn new(database: &mut Database) -> Result<Self, ApiClientError> {
-        Ok(
-            if crate::util::system_info::in_cloudshell()
-                || std::env::var("Q_USE_SENDMESSAGE").is_ok_and(|v| !v.is_empty())
-            {
-                Self::new_qdeveloper_client(database, &Endpoint::load_q(database)).await?
-            } else {
-                Self::new_codewhisperer_client(database, &Endpoint::load_codewhisperer(database)).await?
-            },
-        )
+        Self::new_codewhisperer_client(database, &Endpoint::load_codewhisperer(database)).await
     }
 
     pub fn mock(events: Vec<Vec<ChatResponseStream>>) -> Self {
@@ -108,24 +96,6 @@ impl StreamingClient {
         Ok(Self { inner, profile })
     }
 
-    pub async fn new_qdeveloper_client(database: &Database, endpoint: &Endpoint) -> Result<Self, ApiClientError> {
-        let conf_builder: amzn_qdeveloper_streaming_client::config::Builder =
-            (&sigv4_sdk_config(database, endpoint).await?).into();
-        let conf = conf_builder
-            .http_client(crate::aws_common::http_client::client())
-            .interceptor(OptOutInterceptor::new(database))
-            .interceptor(UserAgentOverrideInterceptor::new())
-            .app_name(app_name())
-            .endpoint_url(endpoint.url())
-            .stalled_stream_protection(stalled_stream_protection_config())
-            .build();
-        let client = QDeveloperStreamingClient::from_conf(conf);
-        Ok(Self {
-            inner: inner::Inner::QDeveloper(client),
-            profile: None,
-        })
-    }
-
     pub async fn send_message(
         &self,
         conversation_state: ConversationState,
@@ -193,27 +163,6 @@ impl StreamingClient {
                     },
                 }
             },
-            inner::Inner::QDeveloper(client) => {
-                let conversation_state_builder = amzn_qdeveloper_streaming_client::types::ConversationState::builder()
-                    .set_conversation_id(conversation_id)
-                    .current_message(amzn_qdeveloper_streaming_client::types::ChatMessage::UserInputMessage(
-                        user_input_message.into(),
-                    ))
-                    .chat_trigger_type(amzn_qdeveloper_streaming_client::types::ChatTriggerType::Manual)
-                    .set_history(
-                        history
-                            .map(|v| v.into_iter().map(|i| i.try_into()).collect::<Result<Vec<_>, _>>())
-                            .transpose()?,
-                    );
-
-                Ok(SendMessageOutput::QDeveloper(
-                    client
-                        .send_message()
-                        .conversation_state(conversation_state_builder.build().expect("fix me"))
-                        .send()
-                        .await?,
-                ))
-            },
             inner::Inner::Mock(events) => {
                 let mut new_events = events.lock().unwrap().next().unwrap_or_default().clone();
                 new_events.reverse();
@@ -228,15 +177,13 @@ pub enum SendMessageOutput {
     Codewhisperer(
         amzn_codewhisperer_streaming_client::operation::generate_assistant_response::GenerateAssistantResponseOutput,
     ),
-    QDeveloper(amzn_qdeveloper_streaming_client::operation::send_message::SendMessageOutput),
     Mock(Vec<ChatResponseStream>),
 }
 
 impl SendMessageOutput {
     pub fn request_id(&self) -> Option<&str> {
         match self {
             SendMessageOutput::Codewhisperer(output) => output.request_id(),
-            SendMessageOutput::QDeveloper(output) => output.request_id(),
             SendMessageOutput::Mock(_) => None,
         }
     }
@@ -248,7 +195,6 @@ impl SendMessageOutput {
                 .recv()
                 .await?
                 .map(|s| s.into())),
-            SendMessageOutput::QDeveloper(output) => Ok(output.send_message_response.recv().await?.map(|s| s.into())),
             SendMessageOutput::Mock(vec) => Ok(vec.pop()),
         }
     }
@@ -258,7 +204,6 @@ impl RequestId for SendMessageOutput {
     fn request_id(&self) -> Option<&str> {
         match self {
             SendMessageOutput::Codewhisperer(output) => output.request_id(),
-            SendMessageOutput::QDeveloper(output) => output.request_id(),
             SendMessageOutput::Mock(_) => Some("<mock-request-id>"),
         }
     }
@@ -280,7 +225,6 @@ mod tests {
 
         let _ = StreamingClient::new(&mut database).await;
         let _ = StreamingClient::new_codewhisperer_client(&mut database, &endpoint).await;
-        let _ = StreamingClient::new_qdeveloper_client(&database, &endpoint).await;
     }
 
     #[tokio::test]

crates/chat-cli/src/api_client/consts.rs

@@ -4,9 +4,6 @@ use aws_config::Region;
 pub const PROD_CODEWHISPERER_ENDPOINT_URL: &str = "https://codewhisperer.us-east-1.amazonaws.com";
 pub const PROD_CODEWHISPERER_ENDPOINT_REGION: Region = Region::from_static("us-east-1");
 
-pub const PROD_Q_ENDPOINT_URL: &str = "https://q.us-east-1.amazonaws.com";
-pub const PROD_Q_ENDPOINT_REGION: Region = Region::from_static("us-east-1");
-
 // FRA endpoint constants
 pub const PROD_CODEWHISPERER_FRA_ENDPOINT_URL: &str = "https://q.eu-central-1.amazonaws.com/";
 pub const PROD_CODEWHISPERER_FRA_ENDPOINT_REGION: Region = Region::from_static("eu-central-1");

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

@@ -9,8 +9,6 @@ use crate::api_client::consts::{
     PROD_CODEWHISPERER_ENDPOINT_URL,
     PROD_CODEWHISPERER_FRA_ENDPOINT_REGION,
     PROD_CODEWHISPERER_FRA_ENDPOINT_URL,
-    PROD_Q_ENDPOINT_REGION,
-    PROD_Q_ENDPOINT_URL,
 };
 use crate::database::Database;
 use crate::database::settings::Setting;
@@ -30,10 +28,6 @@ impl Endpoint {
         url: Cow::Borrowed(PROD_CODEWHISPERER_ENDPOINT_URL),
         region: PROD_CODEWHISPERER_ENDPOINT_REGION,
     };
-    pub const PROD_Q: Self = Self {
-        url: Cow::Borrowed(PROD_Q_ENDPOINT_URL),
-        region: PROD_Q_ENDPOINT_REGION,
-    };
 
     pub fn load_codewhisperer(database: &Database) -> Self {
         let (endpoint, region) = if let Some(Value::Object(o)) = database.settings.get(Setting::ApiCodeWhispererService)
@@ -69,24 +63,6 @@ impl Endpoint {
         }
     }
 
-    pub fn load_q(database: &Database) -> Self {
-        match database.settings.get(Setting::ApiQService) {
-            Some(Value::Object(o)) => {
-                let endpoint = o.get("endpoint").and_then(|v| v.as_str());
-                let region = o.get("region").and_then(|v| v.as_str());
-
-                match (endpoint, region) {
-                    (Some(endpoint), Some(region)) => Self {
-                        url: endpoint.to_owned().into(),
-                        region: Region::new(region.to_owned()),
-                    },
-                    _ => Endpoint::PROD_Q,
-                }
-            },
-            _ => Endpoint::PROD_Q,
-        }
-    }
-
     pub(crate) fn url(&self) -> &str {
         &self.url
     }
@@ -106,7 +82,6 @@ mod tests {
     async fn test_endpoints() {
         let database = Database::new().await.unwrap();
         let _ = Endpoint::load_codewhisperer(&database);
-        let _ = Endpoint::load_q(&database);
 
         let prod = &Endpoint::DEFAULT_ENDPOINT;
         Url::parse(prod.url()).unwrap();

crates/chat-cli/src/api_client/endpoints.rs

@@ -1,13 +1,13 @@
 use amzn_codewhisperer_client::operation::generate_completions::GenerateCompletionsError;
 use amzn_codewhisperer_client::operation::list_available_customizations::ListAvailableCustomizationsError;
 use amzn_codewhisperer_client::operation::list_available_profiles::ListAvailableProfilesError;
+use amzn_codewhisperer_client::operation::send_telemetry_event::SendTelemetryEventError;
 pub use amzn_codewhisperer_streaming_client::operation::generate_assistant_response::GenerateAssistantResponseError;
 use amzn_codewhisperer_streaming_client::types::error::ChatResponseStreamError as CodewhispererChatResponseStreamError;
 use amzn_consolas_client::operation::generate_recommendations::GenerateRecommendationsError;
 use amzn_consolas_client::operation::list_customizations::ListCustomizationsError;
 use amzn_qdeveloper_streaming_client::operation::send_message::SendMessageError as QDeveloperSendMessageError;
 use amzn_qdeveloper_streaming_client::types::error::ChatResponseStreamError as QDeveloperChatResponseStreamError;
-use aws_credential_types::provider::error::CredentialsError;
 use aws_smithy_runtime_api::client::orchestrator::HttpResponse;
 pub use aws_smithy_runtime_api::client::result::SdkError;
 use aws_smithy_types::event_stream::RawMessage;
@@ -18,9 +18,6 @@ use crate::aws_common::SdkErrorDisplay;
 
 #[derive(Debug, Error)]
 pub enum ApiClientError {
-    #[error("failed to load credentials: {}", .0)]
-    Credentials(CredentialsError),
-
     // Generate completions errors
     #[error("{}", SdkErrorDisplay(.0))]
     GenerateCompletions(#[from] SdkError<GenerateCompletionsError, HttpResponse>),
@@ -33,6 +30,10 @@ pub enum ApiClientError {
     #[error("{}", SdkErrorDisplay(.0))]
     ListAvailableServices(#[from] SdkError<ListCustomizationsError, HttpResponse>),
 
+    // Telemetry client error
+    #[error("{}", SdkErrorDisplay(.0))]
+    SendTelemetryEvent(#[from] SdkError<SendTelemetryEventError, HttpResponse>),
+
     // Send message errors
     #[error("{}", SdkErrorDisplay(.0))]
     CodewhispererGenerateAssistantResponse(#[from] SdkError<GenerateAssistantResponseError, HttpResponse>),
@@ -92,7 +93,6 @@ mod tests {
 
     fn all_errors() -> Vec<ApiClientError> {
         vec![
-            ApiClientError::Credentials(CredentialsError::unhandled("<unhandled>")),
             ApiClientError::GenerateCompletions(SdkError::service_error(
                 GenerateCompletionsError::unhandled("<unhandled>"),
                 response(),

crates/chat-cli/src/api_client/error.rs

@@ -1,6 +1,5 @@
 pub mod clients;
 pub(crate) mod consts;
-pub(crate) mod credentials;
 pub mod customization;
 mod endpoints;
 mod error;