feat: add history to ConversationState in StreamingClient::send_message (#280)

grant0417 · web-flow · commit 8af40e37a335 · 2024-12-31T22:28:50.000-08:00
diff --git a/crates/fig_api_client/src/clients/streaming_client.rs b/crates/fig_api_client/src/clients/streaming_client.rs
@@ -86,6 +86,7 @@ impl StreamingClient {
         let ConversationState {
             conversation_id,
             user_input_message,
+            history,
         } = conversation_state;
 
         match &self.0 {
@@ -98,7 +99,12 @@ impl StreamingClient {
                                 user_input_message.into(),
                             ),
                         )
-                        .chat_trigger_type(amzn_codewhisperer_streaming_client::types::ChatTriggerType::Manual);
+                        .chat_trigger_type(amzn_codewhisperer_streaming_client::types::ChatTriggerType::Manual)
+                        .set_history(
+                            history
+                                .map(|v| v.into_iter().map(|i| i.try_into()).collect::<Result<Vec<_>, _>>())
+                                .transpose()?,
+                        );
 
                 Ok(SendMessageOutput::Codewhisperer(
                     client
@@ -114,7 +120,12 @@ impl StreamingClient {
                     .current_message(amzn_qdeveloper_streaming_client::types::ChatMessage::UserInputMessage(
                         user_input_message.into(),
                     ))
-                    .chat_trigger_type(amzn_qdeveloper_streaming_client::types::ChatTriggerType::Manual);
+                    .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
@@ -168,7 +179,11 @@ impl RequestId for SendMessageOutput {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::model::UserInputMessage;
+    use crate::model::{
+        AssistantResponseMessage,
+        ChatMessage,
+        UserInputMessage,
+    };
 
     #[tokio::test]
     async fn create_clients() {
@@ -194,6 +209,7 @@ mod tests {
                     user_input_message_context: None,
                     user_intent: None,
                 },
+                history: None,
             })
             .await
             .unwrap();
@@ -204,4 +220,36 @@ mod tests {
         }
         assert_eq!(output_content, "Hello! How can I assist you today?");
     }
+
+    #[ignore]
+    #[tokio::test]
+    async fn assistant_response() {
+        let client = StreamingClient::new().await.unwrap();
+        let mut response = client
+            .send_message(ConversationState {
+                conversation_id: None,
+                user_input_message: UserInputMessage {
+                    content: "How about rustc?".into(),
+                    user_input_message_context: None,
+                    user_intent: None,
+                },
+                history: Some(vec![
+                    ChatMessage::UserInputMessage(UserInputMessage {
+                        content: "What language is the linux kernel written in, and who wrote it?".into(),
+                        user_input_message_context: None,
+                        user_intent: None,
+                    }),
+                    ChatMessage::AssistantResponseMessage(AssistantResponseMessage {
+                        content: "It is written in C by Linus Torvalds.".into(),
+                        message_id: None,
+                    }),
+                ]),
+            })
+            .await
+            .unwrap();
+
+        while let Some(event) = response.recv().await.unwrap() {
+            println!("{:?}", event);
+        }
+    }
 }
diff --git a/crates/fig_api_client/src/error.rs b/crates/fig_api_client/src/error.rs
@@ -44,6 +44,9 @@ pub enum Error {
     #[error("{}", SdkErrorDisplay(.0))]
     QDeveloperChatResponseStream(#[from] SdkError<QDeveloperChatResponseStreamError, RawMessage>),
 
+    #[error(transparent)]
+    SmithyBuild(#[from] aws_smithy_types::error::operation::BuildError),
+
     #[error("unsupported action by consolas: {0}")]
     UnsupportedConsolas(&'static str),
 }
@@ -62,6 +65,7 @@ impl Error {
             Error::QDeveloperSendMessage(e) => e.as_service_error().map_or(false, |e| e.is_throttling_error()),
             Error::CodewhispererChatResponseStream(_)
             | Error::QDeveloperChatResponseStream(_)
+            | Error::SmithyBuild(_)
             | Error::UnsupportedConsolas(_) => false,
         }
     }
@@ -77,6 +81,7 @@ impl Error {
             Error::QDeveloperSendMessage(e) => e.as_service_error().is_some(),
             Error::CodewhispererChatResponseStream(_)
             | Error::QDeveloperChatResponseStream(_)
+            | Error::SmithyBuild(_)
             | Error::UnsupportedConsolas(_) => false,
         }
     }
@@ -135,6 +140,7 @@ mod tests {
                 QDeveloperChatResponseStreamError::unhandled("<unhandled>"),
                 raw_message(),
             )),
+            Error::SmithyBuild(aws_smithy_types::error::operation::BuildError::other("<other>")),
             Error::UnsupportedConsolas("test"),
         ]
     }
diff --git a/crates/fig_api_client/src/model.rs b/crates/fig_api_client/src/model.rs
@@ -77,6 +77,71 @@ pub struct Recommendation {
 pub struct ConversationState {
     pub conversation_id: Option<String>,
     pub user_input_message: UserInputMessage,
+    pub history: Option<Vec<ChatMessage>>,
+}
+
+#[derive(Debug, Clone)]
+pub enum ChatMessage {
+    AssistantResponseMessage(AssistantResponseMessage),
+    UserInputMessage(UserInputMessage),
+}
+
+impl TryFrom<ChatMessage> for amzn_codewhisperer_streaming_client::types::ChatMessage {
+    type Error = aws_smithy_types::error::operation::BuildError;
+
+    fn try_from(value: ChatMessage) -> Result<Self, Self::Error> {
+        Ok(match value {
+            ChatMessage::AssistantResponseMessage(message) => {
+                amzn_codewhisperer_streaming_client::types::ChatMessage::AssistantResponseMessage(message.try_into()?)
+            },
+            ChatMessage::UserInputMessage(message) => {
+                amzn_codewhisperer_streaming_client::types::ChatMessage::UserInputMessage(message.into())
+            },
+        })
+    }
+}
+
+impl TryFrom<ChatMessage> for amzn_qdeveloper_streaming_client::types::ChatMessage {
+    type Error = aws_smithy_types::error::operation::BuildError;
+
+    fn try_from(value: ChatMessage) -> Result<Self, Self::Error> {
+        Ok(match value {
+            ChatMessage::AssistantResponseMessage(message) => {
+                amzn_qdeveloper_streaming_client::types::ChatMessage::AssistantResponseMessage(message.try_into()?)
+            },
+            ChatMessage::UserInputMessage(message) => {
+                amzn_qdeveloper_streaming_client::types::ChatMessage::UserInputMessage(message.into())
+            },
+        })
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct AssistantResponseMessage {
+    pub message_id: Option<String>,
+    pub content: String,
+}
+
+impl TryFrom<AssistantResponseMessage> for amzn_codewhisperer_streaming_client::types::AssistantResponseMessage {
+    type Error = aws_smithy_types::error::operation::BuildError;
+
+    fn try_from(value: AssistantResponseMessage) -> Result<Self, Self::Error> {
+        Self::builder()
+            .content(value.content)
+            .set_message_id(value.message_id)
+            .build()
+    }
+}
+
+impl TryFrom<AssistantResponseMessage> for amzn_qdeveloper_streaming_client::types::AssistantResponseMessage {
+    type Error = aws_smithy_types::error::operation::BuildError;
+
+    fn try_from(value: AssistantResponseMessage) -> Result<Self, Self::Error> {
+        Self::builder()
+            .content(value.content)
+            .set_message_id(value.message_id)
+            .build()
+    }
 }
 
 #[non_exhaustive]
diff --git a/crates/q_cli/src/cli/chat/api.rs b/crates/q_cli/src/cli/chat/api.rs
@@ -311,6 +311,7 @@ pub(super) async fn send_message(
     let conversation_state = ConversationState {
         conversation_id,
         user_input_message,
+        history: None,
     };
 
     tokio::spawn(async move {

Original file line number	Diff line number	Diff line change
`@@ -44,6 +44,9 @@ pub enum Error {`
`44`	`44`	`#[error("{}", SdkErrorDisplay(.0))]`
`45`	`45`	`QDeveloperChatResponseStream(#[from] SdkError<QDeveloperChatResponseStreamError, RawMessage>),`
`46`	`46`
	`47`	`+ #[error(transparent)]`
	`48`	`+ SmithyBuild(#[from] aws_smithy_types::error::operation::BuildError),`
	`49`	`+`
`47`	`50`	`#[error("unsupported action by consolas: {0}")]`
`48`	`51`	`UnsupportedConsolas(&'static str),`
`49`	`52`	`}`
`@@ -62,6 +65,7 @@ impl Error {`
`62`	`65`	`Error::QDeveloperSendMessage(e) => e.as_service_error().map_or(false, \|e\| e.is_throttling_error()),`
`63`	`66`	`Error::CodewhispererChatResponseStream(_)`
`64`	`67`	`\| Error::QDeveloperChatResponseStream(_)`
	`68`	`+ \| Error::SmithyBuild(_)`
`65`	`69`	`\| Error::UnsupportedConsolas(_) => false,`
`66`	`70`	`}`
`67`	`71`	`}`
`@@ -77,6 +81,7 @@ impl Error {`
`77`	`81`	`Error::QDeveloperSendMessage(e) => e.as_service_error().is_some(),`
`78`	`82`	`Error::CodewhispererChatResponseStream(_)`
`79`	`83`	`\| Error::QDeveloperChatResponseStream(_)`
	`84`	`+ \| Error::SmithyBuild(_)`
`80`	`85`	`\| Error::UnsupportedConsolas(_) => false,`
`81`	`86`	`}`
`82`	`87`	`}`
`@@ -135,6 +140,7 @@ mod tests {`
`135`	`140`	`QDeveloperChatResponseStreamError::unhandled("<unhandled>"),`
`136`	`141`	`raw_message(),`
`137`	`142`	`)),`
	`143`	`+ Error::SmithyBuild(aws_smithy_types::error::operation::BuildError::other("<other>")),`
`138`	`144`	`Error::UnsupportedConsolas("test"),`
`139`	`145`	`]`
`140`	`146`	`}`