Merge branch 'main' into auto-tool-name-resolution

Author: dingfeli

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

@@ -291,6 +291,7 @@ mod tests {
             .send_message(ConversationState {
                 conversation_id: None,
                 user_input_message: UserInputMessage {
+                    images: None,
                     content: "Hello".into(),
                     user_input_message_context: None,
                     user_intent: None,
@@ -315,12 +316,14 @@ mod tests {
             .send_message(ConversationState {
                 conversation_id: None,
                 user_input_message: UserInputMessage {
+                    images: None,
                     content: "How about rustc?".into(),
                     user_input_message_context: None,
                     user_intent: None,
                 },
                 history: Some(vec![
                     ChatMessage::UserInputMessage(UserInputMessage {
+                        images: None,
                         content: "What language is the linux kernel written in, and who wrote it?".into(),
                         user_input_message_context: None,
                         user_intent: None,

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

@@ -1,4 +1,7 @@
-use aws_smithy_types::Document;
+use aws_smithy_types::{
+    Blob,
+    Document,
+};
 use serde::{
     Deserialize,
     Serialize,
@@ -565,17 +568,113 @@ impl From<GitState> for amzn_qdeveloper_streaming_client::types::GitState {
     }
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct ImageBlock {
+    pub format: ImageFormat,
+    pub source: ImageSource,
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
+pub enum ImageFormat {
+    Gif,
+    Jpeg,
+    Png,
+    Webp,
+}
+
+impl std::str::FromStr for ImageFormat {
+    type Err = String;
+
+    fn from_str(s: &str) -> Result<Self, Self::Err> {
+        match s.trim().to_lowercase().as_str() {
+            "gif" => Ok(ImageFormat::Gif),
+            "jpeg" => Ok(ImageFormat::Jpeg),
+            "jpg" => Ok(ImageFormat::Jpeg),
+            "png" => Ok(ImageFormat::Png),
+            "webp" => Ok(ImageFormat::Webp),
+            _ => Err(format!("Failed to parse '{}' as ImageFormat", s)),
+        }
+    }
+}
+
+impl From<ImageFormat> for amzn_codewhisperer_streaming_client::types::ImageFormat {
+    fn from(value: ImageFormat) -> Self {
+        match value {
+            ImageFormat::Gif => Self::Gif,
+            ImageFormat::Jpeg => Self::Jpeg,
+            ImageFormat::Png => Self::Png,
+            ImageFormat::Webp => Self::Webp,
+        }
+    }
+}
+impl From<ImageFormat> for amzn_qdeveloper_streaming_client::types::ImageFormat {
+    fn from(value: ImageFormat) -> Self {
+        match value {
+            ImageFormat::Gif => Self::Gif,
+            ImageFormat::Jpeg => Self::Jpeg,
+            ImageFormat::Png => Self::Png,
+            ImageFormat::Webp => Self::Webp,
+        }
+    }
+}
+
+#[non_exhaustive]
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub enum ImageSource {
+    Bytes(Vec<u8>),
+    #[non_exhaustive]
+    Unknown,
+}
+
+impl From<ImageSource> for amzn_codewhisperer_streaming_client::types::ImageSource {
+    fn from(value: ImageSource) -> Self {
+        match value {
+            ImageSource::Bytes(bytes) => Self::Bytes(Blob::new(bytes)),
+            ImageSource::Unknown => Self::Unknown,
+        }
+    }
+}
+impl From<ImageSource> for amzn_qdeveloper_streaming_client::types::ImageSource {
+    fn from(value: ImageSource) -> Self {
+        match value {
+            ImageSource::Bytes(bytes) => Self::Bytes(Blob::new(bytes)),
+            ImageSource::Unknown => Self::Unknown,
+        }
+    }
+}
+
+impl From<ImageBlock> for amzn_codewhisperer_streaming_client::types::ImageBlock {
+    fn from(value: ImageBlock) -> Self {
+        Self::builder()
+            .format(value.format.into())
+            .source(value.source.into())
+            .build()
+            .expect("Failed to build ImageBlock")
+    }
+}
+impl From<ImageBlock> for amzn_qdeveloper_streaming_client::types::ImageBlock {
+    fn from(value: ImageBlock) -> Self {
+        Self::builder()
+            .format(value.format.into())
+            .source(value.source.into())
+            .build()
+            .expect("Failed to build ImageBlock")
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct UserInputMessage {
     pub content: String,
     pub user_input_message_context: Option<UserInputMessageContext>,
     pub user_intent: Option<UserIntent>,
+    pub images: Option<Vec<ImageBlock>>,
 }
 
 impl From<UserInputMessage> for amzn_codewhisperer_streaming_client::types::UserInputMessage {
     fn from(value: UserInputMessage) -> Self {
         Self::builder()
             .content(value.content)
+            .set_images(value.images.map(|images| images.into_iter().map(Into::into).collect()))
             .set_user_input_message_context(value.user_input_message_context.map(Into::into))
             .set_user_intent(value.user_intent.map(Into::into))
             .origin(amzn_codewhisperer_streaming_client::types::Origin::Cli)
@@ -588,6 +687,7 @@ impl From<UserInputMessage> for amzn_qdeveloper_streaming_client::types::UserInp
     fn from(value: UserInputMessage) -> Self {
         Self::builder()
             .content(value.content)
+            .set_images(value.images.map(|images| images.into_iter().map(Into::into).collect()))
             .set_user_input_message_context(value.user_input_message_context.map(Into::into))
             .set_user_intent(value.user_intent.map(Into::into))
             .origin(amzn_qdeveloper_streaming_client::types::Origin::Cli)
@@ -654,6 +754,10 @@ mod tests {
     #[test]
     fn build_user_input_message() {
         let user_input_message = UserInputMessage {
+            images: Some(vec![ImageBlock {
+                format: ImageFormat::Png,
+                source: ImageSource::Bytes(vec![1, 2, 3]),
+            }]),
             content: "test content".to_string(),
             user_input_message_context: Some(UserInputMessageContext {
                 env_state: Some(EnvState {
@@ -690,6 +794,7 @@ mod tests {
         assert_eq!(format!("{codewhisper_input:?}"), format!("{qdeveloper_input:?}"));
 
         let minimal_message = UserInputMessage {
+            images: None,
             content: "test content".to_string(),
             user_input_message_context: None,
             user_intent: None,

crates/chat-cli/src/auth/builder_id.rs

@@ -69,7 +69,7 @@ impl std::fmt::Display for OAuthFlow {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match *self {
             OAuthFlow::DeviceCode => write!(f, "DeviceCode"),
-            OAuthFlow::Pkce => write!(f, "PKCE"),
+            OAuthFlow::Pkce => write!(f, "Pkce"),
         }
     }
 }
@@ -593,7 +593,7 @@ mod tests {
     #[test]
     fn test_oauth_flow_ser_deser() {
         test_ser_deser!(OAuthFlow, OAuthFlow::DeviceCode, "DeviceCode");
-        test_ser_deser!(OAuthFlow, OAuthFlow::Pkce, "PKCE");
+        test_ser_deser!(OAuthFlow, OAuthFlow::Pkce, "Pkce");
     }
 
     #[test]

crates/chat-cli/src/auth/secret_store/sqlite.rs

@@ -1,6 +1,5 @@
 #![allow(dead_code)]
 use super::Secret;
-use crate::Result;
 use crate::auth::AuthError;
 use crate::settings::sqlite::{
     Db,

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

@@ -16,6 +16,13 @@ pub const MAX_USER_MESSAGE_SIZE: usize = 600_000;
 /// In tokens
 pub const CONTEXT_WINDOW_SIZE: usize = 200_000;
 
+pub const CONTEXT_FILES_MAX_SIZE: usize = 150_000;
+
 pub const MAX_CHARS: usize = TokenCounter::token_to_chars(CONTEXT_WINDOW_SIZE); // Character-based warning threshold
 
 pub const DUMMY_TOOL_NAME: &str = "dummy";
+
+pub const MAX_NUMBER_OF_IMAGES_PER_REQUEST: usize = 10;
+
+/// In bytes - 10 MB
+pub const MAX_IMAGE_SIZE: usize = 10 * 1024 * 1024;