Skip to content

Commit 6b5fa35

Browse files
evanliu048evanliu048
committed
add model id in generateResponse API and coversationState
1 parent 072840c commit 6b5fa35

File tree

8 files changed

+25
-2
lines changed

8 files changed

+25
-2
lines changed

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

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,6 @@ impl StreamingClient {
126126
})
127127
}
128128

129-
//todo yifan generate response
130129
pub async fn send_message(
131130
&self,
132131
conversation_state: ConversationState,
@@ -292,6 +291,7 @@ mod tests {
292291
content: "Hello".into(),
293292
user_input_message_context: None,
294293
user_intent: None,
294+
model_id: Some("model".to_owned()),
295295
},
296296
history: None,
297297
})
@@ -318,13 +318,15 @@ mod tests {
318318
content: "How about rustc?".into(),
319319
user_input_message_context: None,
320320
user_intent: None,
321+
model_id: Some("model".to_owned()),
321322
},
322323
history: Some(vec![
323324
ChatMessage::UserInputMessage(UserInputMessage {
324325
images: None,
325326
content: "What language is the linux kernel written in, and who wrote it?".into(),
326327
user_input_message_context: None,
327328
user_intent: None,
329+
model_id: None,
328330
}),
329331
ChatMessage::AssistantResponseMessage(AssistantResponseMessage {
330332
content: "It is written in C by Linus Torvalds.".into(),

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

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -859,6 +859,7 @@ pub struct UserInputMessage {
859859
pub user_input_message_context: Option<UserInputMessageContext>,
860860
pub user_intent: Option<UserIntent>,
861861
pub images: Option<Vec<ImageBlock>>,
862+
pub model_id: Option<String>,
862863
}
863864

864865
impl From<UserInputMessage> for amzn_codewhisperer_streaming_client::types::UserInputMessage {
@@ -868,6 +869,7 @@ impl From<UserInputMessage> for amzn_codewhisperer_streaming_client::types::User
868869
.set_images(value.images.map(|images| images.into_iter().map(Into::into).collect()))
869870
.set_user_input_message_context(value.user_input_message_context.map(Into::into))
870871
.set_user_intent(value.user_intent.map(Into::into))
872+
.set_model_id(value.model_id)
871873
.origin(amzn_codewhisperer_streaming_client::types::Origin::Cli)
872874
.build()
873875
.expect("Failed to build UserInputMessage")
@@ -881,6 +883,7 @@ impl From<UserInputMessage> for amzn_qdeveloper_streaming_client::types::UserInp
881883
.set_images(value.images.map(|images| images.into_iter().map(Into::into).collect()))
882884
.set_user_input_message_context(value.user_input_message_context.map(Into::into))
883885
.set_user_intent(value.user_intent.map(Into::into))
886+
.set_model_id(value.model_id)
884887
.origin(amzn_qdeveloper_streaming_client::types::Origin::Cli)
885888
.build()
886889
.expect("Failed to build UserInputMessage")
@@ -976,6 +979,7 @@ mod tests {
976979
})]),
977980
}),
978981
user_intent: Some(UserIntent::ApplyCommonBestPractices),
982+
model_id: Some("model id".to_string()),
979983
};
980984

981985
let codewhisper_input =
@@ -989,6 +993,7 @@ mod tests {
989993
content: "test content".to_string(),
990994
user_input_message_context: None,
991995
user_intent: None,
996+
model_id: Some("model id".to_string()),
992997
};
993998

994999
let codewhisper_minimal =

‎crates/chat-cli/src/cli/chat/conversation_state.rs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -533,6 +533,7 @@ impl ConversationState {
533533
context_messages,
534534
dropped_context_files,
535535
tools: &self.tools,
536+
model_id: self.current_model_id.as_deref(),
536537
}
537538
}
538539

@@ -604,6 +605,7 @@ impl ConversationState {
604605
user_input_message_context: None,
605606
user_intent: None,
606607
images: None,
608+
model_id: self.current_model_id.to_owned(),
607609
};
608610

609611
// If the last message contains tool uses, then add cancelled tool results to the summary
@@ -835,6 +837,7 @@ pub struct BackendConversationStateImpl<'a, T, U> {
835837
pub context_messages: U,
836838
pub dropped_context_files: Vec<(String, String)>,
837839
pub tools: &'a HashMap<ToolOrigin, Vec<Tool>>,
840+
pub model_id: Option<&'a str>,
838841
}
839842

840843
impl
@@ -851,6 +854,7 @@ impl
851854
.cloned()
852855
.map(UserMessage::into_user_input_message)
853856
.ok_or(eyre::eyre!("next user message is not set"))?;
857+
user_input_message.model_id = self.model_id.map(str::to_string);
854858
if let Some(ctx) = user_input_message.user_input_message_context.as_mut() {
855859
ctx.tools = Some(self.tools.values().flatten().cloned().collect::<Vec<_>>());
856860
}

‎crates/chat-cli/src/cli/chat/message.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -128,6 +128,7 @@ impl UserMessage {
128128
..Default::default()
129129
}),
130130
user_intent: None,
131+
model_id: None,
131132
}
132133
}
133134

@@ -158,6 +159,7 @@ impl UserMessage {
158159
..Default::default()
159160
}),
160161
user_intent: None,
162+
model_id: None,
161163
}
162164
}
163165

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -562,7 +562,7 @@ impl ChatContext {
562562
let output_clone = output.clone();
563563

564564
let mut existing_conversation = false;
565-
let model_id = Some(database.settings.get_string(Setting::UserDefaultModel).unwrap());
565+
let model_id = database.settings.get_string(Setting::UserDefaultModel);
566566
let conversation_state = if resume_conversation {
567567
let prior = std::env::current_dir()
568568
.ok()

‎crates/fig_api_client/src/clients/client.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ use amzn_codewhisperer_client::types::{
1212
TelemetryEvent,
1313
UserContext,
1414
};
15+
use crate::model::ConversationState;
1516
use amzn_consolas_client::Client as ConsolasClient;
1617
use amzn_consolas_client::error::DisplayErrorContext;
1718
use amzn_consolas_client::operation::generate_recommendations::GenerateRecommendationsError;
@@ -223,6 +224,7 @@ impl Client {
223224
Ok(customizations)
224225
}
225226

227+
//(todo) yifan : how to propgate modelID
226228
// .telemetry_event(TelemetryEvent::UserTriggerDecisionEvent(user_trigger_decision_event))
227229
// .user_context(user_context)
228230
// .opt_out_preference(opt_out_preference)

‎crates/fig_api_client/src/clients/streaming_client.rs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -286,6 +286,7 @@ mod tests {
286286
content: "Hello".into(),
287287
user_input_message_context: None,
288288
user_intent: None,
289+
model_id: Some("model".to_owned()),
289290
},
290291
history: None,
291292
})
@@ -311,13 +312,15 @@ mod tests {
311312
content: "How about rustc?".into(),
312313
user_input_message_context: None,
313314
user_intent: None,
315+
model_id: Some("model".to_owned()),
314316
},
315317
history: Some(vec![
316318
ChatMessage::UserInputMessage(UserInputMessage {
317319
images: None,
318320
content: "What language is the linux kernel written in, and who wrote it?".into(),
319321
user_input_message_context: None,
320322
user_intent: None,
323+
model_id: Some("model".to_owned()),
321324
}),
322325
ChatMessage::AssistantResponseMessage(AssistantResponseMessage {
323326
content: "It is written in C by Linus Torvalds.".into(),

‎crates/fig_api_client/src/model.rs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -646,6 +646,7 @@ pub struct UserInputMessage {
646646
pub user_input_message_context: Option<UserInputMessageContext>,
647647
pub user_intent: Option<UserIntent>,
648648
pub images: Option<Vec<ImageBlock>>,
649+
pub model_id: Option<String>
649650
}
650651

651652
impl From<UserInputMessage> for amzn_codewhisperer_streaming_client::types::UserInputMessage {
@@ -655,6 +656,7 @@ impl From<UserInputMessage> for amzn_codewhisperer_streaming_client::types::User
655656
.set_images(value.images.map(|images| images.into_iter().map(Into::into).collect()))
656657
.set_user_input_message_context(value.user_input_message_context.map(Into::into))
657658
.set_user_intent(value.user_intent.map(Into::into))
659+
.set_model_id(value.model_id)
658660
.origin(amzn_codewhisperer_streaming_client::types::Origin::Cli)
659661
.build()
660662
.expect("Failed to build UserInputMessage")
@@ -668,6 +670,7 @@ impl From<UserInputMessage> for amzn_qdeveloper_streaming_client::types::UserInp
668670
.set_images(value.images.map(|images| images.into_iter().map(Into::into).collect()))
669671
.set_user_input_message_context(value.user_input_message_context.map(Into::into))
670672
.set_user_intent(value.user_intent.map(Into::into))
673+
.set_model_id(value.model_id)
671674
.origin(amzn_qdeveloper_streaming_client::types::Origin::Cli)
672675
.build()
673676
.expect("Failed to build UserInputMessage")
@@ -853,6 +856,7 @@ mod tests {
853856
})]),
854857
}),
855858
user_intent: Some(UserIntent::ApplyCommonBestPractices),
859+
model_id: Some("model".to_owned()),
856860
};
857861

858862
let codewhisper_input =
@@ -866,6 +870,7 @@ mod tests {
866870
user_input_message_context: None,
867871
user_intent: None,
868872
images: None,
873+
model_id: Some("model".to_owned()),
869874
};
870875

871876
let codewhisper_minimal =

0 commit comments

Comments
 (0)