Skip to content

Commit d6ee9ed

Browse files
tusharmathtusharmath
committed
feat: add raw_content field to TextMessage and update related tests
1 parent c5f72b0 commit d6ee9ed

14 files changed

Lines changed: 117 additions & 3 deletions

File tree

‎crates/forge_app/src/dto/anthropic/transforms/drop_invalid_toolcalls.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ mod tests {
4848
fn transform_tool_call(json_args: &str) -> Request {
4949
let fixture = Context::default().messages(vec![ContextMessage::Text(TextMessage {
5050
role: Role::User,
51+
raw_content: None,
5152
content: "Hello".to_string(),
5253
tool_calls: Some(vec![
5354
ToolCallFull::new("test_tool")
@@ -123,6 +124,7 @@ mod tests {
123124
fn test_preserves_text_content() {
124125
let fixture = Context::default().messages(vec![ContextMessage::Text(TextMessage {
125126
role: Role::User,
127+
raw_content: None,
126128
content: "Hello".to_string(),
127129
tool_calls: None,
128130
model: None,

‎crates/forge_app/src/dto/anthropic/transforms/set_cache.rs‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ mod tests {
6767
match c {
6868
's' => messages.push(ContextMessage::Text(TextMessage {
6969
role: Role::System,
70+
raw_content: None,
7071
content: c.to_string(),
7172
tool_calls: None,
7273
model: None,
@@ -81,13 +82,15 @@ mod tests {
8182
match c {
8283
'u' => messages.push(ContextMessage::Text(TextMessage {
8384
role: Role::User,
85+
raw_content: None,
8486
content: c.to_string(),
8587
tool_calls: None,
8688
model: ModelId::new("claude-3-5-sonnet-20241022").into(),
8789
reasoning_details: None,
8890
})),
8991
'a' => messages.push(ContextMessage::Text(TextMessage {
9092
role: Role::Assistant,
93+
raw_content: None,
9194
content: c.to_string(),
9295
tool_calls: None,
9396
model: None,
@@ -231,20 +234,23 @@ mod tests {
231234
messages: vec![
232235
ContextMessage::Text(TextMessage {
233236
role: Role::System,
237+
raw_content: None,
234238
content: "first".to_string(),
235239
tool_calls: None,
236240
model: None,
237241
reasoning_details: None,
238242
}),
239243
ContextMessage::Text(TextMessage {
240244
role: Role::System,
245+
raw_content: None,
241246
content: "second".to_string(),
242247
tool_calls: None,
243248
model: None,
244249
reasoning_details: None,
245250
}),
246251
ContextMessage::Text(TextMessage {
247252
role: Role::User,
253+
raw_content: None,
248254
content: "user".to_string(),
249255
tool_calls: None,
250256
model: ModelId::new("claude-3-5-sonnet-20241022").into(),

‎crates/forge_app/src/dto/openai/request.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -629,6 +629,7 @@ mod tests {
629629
fn test_user_message_conversion() {
630630
let user_message = ContextMessage::Text(TextMessage {
631631
role: Role::User,
632+
raw_content: None,
632633
content: "Hello".to_string(),
633634
tool_calls: None,
634635
model: ModelId::new("gpt-3.5-turbo").into(),
@@ -652,6 +653,7 @@ mod tests {
652653

653654
let message = ContextMessage::Text(TextMessage {
654655
role: Role::User,
656+
raw_content: None,
655657
content: xml_content.to_string(),
656658
tool_calls: None,
657659
model: ModelId::new("gpt-3.5-turbo").into(),
@@ -671,6 +673,7 @@ mod tests {
671673

672674
let assistant_message = ContextMessage::Text(TextMessage {
673675
role: Role::Assistant,
676+
raw_content: None,
674677
content: "Using tool".to_string(),
675678
tool_calls: Some(vec![tool_call]),
676679
model: ModelId::new("gpt-3.5-turbo").into(),

‎crates/forge_app/src/dto/openai/transformers/drop_tool_call.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ mod tests {
5757
messages: vec![
5858
ContextMessage::Text(TextMessage {
5959
role: Role::Assistant,
60+
raw_content: None,
6061
content: "Using tool".to_string(),
6162
tool_calls: Some(vec![tool_call]),
6263
model: None,

‎crates/forge_app/src/dto/openai/transformers/set_cache.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,20 +69,23 @@ mod tests {
6969
.map(|c| match c {
7070
's' => ContextMessage::Text(TextMessage {
7171
role: Role::System,
72+
raw_content: None,
7273
content: c.to_string(),
7374
tool_calls: None,
7475
model: None,
7576
reasoning_details: None,
7677
}),
7778
'u' => ContextMessage::Text(TextMessage {
7879
role: Role::User,
80+
raw_content: None,
7981
content: c.to_string(),
8082
tool_calls: None,
8183
model: ModelId::new("gpt-4").into(),
8284
reasoning_details: None,
8385
}),
8486
'a' => ContextMessage::Text(TextMessage {
8587
role: Role::Assistant,
88+
raw_content: None,
8689
content: c.to_string(),
8790
tool_calls: None,
8891
model: None,

‎crates/forge_app/src/orch_spec/orch_spec.rs‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -410,3 +410,48 @@ async fn test_multi_turn_conversation_stops_only_on_finish_reason() {
410410
"Should have exactly 3 assistant messages, confirming the orchestrator continued until FinishReason::Stop"
411411
);
412412
}
413+
414+
#[tokio::test]
415+
async fn test_raw_user_message_is_stored() {
416+
let mut ctx = TestContext::default().mock_assistant_responses(vec![
417+
ChatCompletionMessage::assistant(Content::full("Hello!")).finish_reason(FinishReason::Stop),
418+
]);
419+
420+
let raw_task = "This is a raw user message\nwith multiple lines\nfor testing";
421+
ctx.run(raw_task).await.unwrap();
422+
423+
let conversation = ctx.output.conversation_history.last().unwrap();
424+
let context = conversation.context.as_ref().unwrap();
425+
426+
// Find the user message
427+
let user_message = context
428+
.messages
429+
.iter()
430+
.find(|msg| msg.has_role(Role::User))
431+
.expect("Should have user message");
432+
433+
// Verify raw content is stored
434+
assert_eq!(
435+
user_message.raw_content(),
436+
Some(raw_task),
437+
"Raw user message should be stored in TextMessage"
438+
);
439+
440+
// Verify rendered content is different (has template wrapping)
441+
let rendered_content = user_message.content().unwrap();
442+
assert!(
443+
rendered_content.contains("<task>"),
444+
"Rendered message should contain template tags"
445+
);
446+
assert!(
447+
rendered_content.contains(raw_task),
448+
"Rendered message should contain the raw task"
449+
);
450+
451+
// Verify raw content doesn't have template tags
452+
let raw_content = user_message.raw_content().unwrap();
453+
assert!(
454+
!raw_content.contains("<task>"),
455+
"Raw content should not contain template tags"
456+
);
457+
}

‎crates/forge_app/src/user_prompt.rs‎

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,11 @@ impl<S> UserPromptBuilder<S> {
3131
where
3232
S: AgentService,
3333
{
34+
let raw_message = self.event.value.as_ref().map(|v| match v {
35+
serde_json::Value::String(s) => s.clone(),
36+
other => other.to_string(),
37+
});
38+
3439
let content = if let Some(user_prompt) = &self.agent.user_prompt
3540
&& self.event.value.is_some()
3641
{
@@ -54,11 +59,19 @@ impl<S> UserPromptBuilder<S> {
5459
)
5560
} else {
5661
// Use the raw event value as content if no user_prompt is provided
57-
self.event.value.as_ref().map(|v| v.to_string())
62+
raw_message.clone()
5863
};
5964

6065
if let Some(content) = content {
61-
context = context.add_message(ContextMessage::user(content, self.agent.model.clone()));
66+
let message = TextMessage {
67+
role: Role::User,
68+
content,
69+
raw_content: raw_message,
70+
tool_calls: None,
71+
reasoning_details: None,
72+
model: self.agent.model.clone(),
73+
};
74+
context = context.add_message(ContextMessage::Text(message));
6275
}
6376

6477
Ok(context)

‎crates/forge_domain/src/compact.rs‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,20 +209,23 @@ mod tests {
209209
'u' => ContextMessage::Text(TextMessage {
210210
role: Role::User,
211211
content,
212+
raw_content: None,
212213
tool_calls: None,
213214
model: None,
214215
reasoning_details: None,
215216
}),
216217
'a' => ContextMessage::Text(TextMessage {
217218
role: Role::Assistant,
218219
content,
220+
raw_content: None,
219221
tool_calls: None,
220222
model: None,
221223
reasoning_details: None,
222224
}),
223225
's' => ContextMessage::Text(TextMessage {
224226
role: Role::System,
225227
content,
228+
raw_content: None,
226229
tool_calls: None,
227230
model: None,
228231
reasoning_details: None,

‎crates/forge_domain/src/context.rs‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,16 @@ impl ContextMessage {
5858
}
5959
}
6060

61+
/// Returns the raw content before template rendering (only for User
62+
/// messages)
63+
pub fn raw_content(&self) -> Option<&str> {
64+
match self {
65+
ContextMessage::Text(text_message) => text_message.raw_content.as_deref(),
66+
ContextMessage::Tool(_) => None,
67+
ContextMessage::Image(_) => None,
68+
}
69+
}
70+
6171
/// Estimates the number of tokens in a message using character-based
6272
/// approximation.
6373
/// ref: https://github.com/openai/codex/blob/main/codex-cli/src/utils/approximate-tokens-used.ts
@@ -133,6 +143,7 @@ impl ContextMessage {
133143
TextMessage {
134144
role: Role::User,
135145
content: content.to_string(),
146+
raw_content: None,
136147
tool_calls: None,
137148
reasoning_details: None,
138149
model,
@@ -144,6 +155,7 @@ impl ContextMessage {
144155
TextMessage {
145156
role: Role::System,
146157
content: content.to_string(),
158+
raw_content: None,
147159
tool_calls: None,
148160
model: None,
149161
reasoning_details: None,
@@ -161,6 +173,7 @@ impl ContextMessage {
161173
TextMessage {
162174
role: Role::Assistant,
163175
content: content.to_string(),
176+
raw_content: None,
164177
tool_calls,
165178
reasoning_details,
166179
model: None,
@@ -240,6 +253,9 @@ fn reasoning_content_char_count(text_message: &TextMessage) -> usize {
240253
pub struct TextMessage {
241254
pub role: Role,
242255
pub content: String,
256+
/// The raw content before any template rendering (only for User messages)
257+
#[serde(default, skip_serializing_if = "Option::is_none")]
258+
pub raw_content: Option<String>,
243259
#[serde(default, skip_serializing_if = "Option::is_none")]
244260
pub tool_calls: Option<Vec<ToolCallFull>>,
245261
// note: this used to track model used for this message.
@@ -262,6 +278,7 @@ impl TextMessage {
262278
Self {
263279
role: Role::Assistant,
264280
content: content.to_string(),
281+
raw_content: None,
265282
tool_calls: None,
266283
reasoning_details,
267284
model,

‎crates/forge_domain/src/transformer/drop_reasoning_details.rs‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,13 +53,15 @@ mod tests {
5353
.add_message(ContextMessage::Text(TextMessage {
5454
role: Role::User,
5555
content: "User message with reasoning".to_string(),
56+
raw_content: None,
5657
tool_calls: None,
5758
model: None,
5859
reasoning_details: Some(reasoning_details.clone()),
5960
}))
6061
.add_message(ContextMessage::Text(TextMessage {
6162
role: Role::Assistant,
6263
content: "Assistant response with reasoning".to_string(),
64+
raw_content: None,
6365
tool_calls: None,
6466
model: None,
6567
reasoning_details: Some(reasoning_details),
@@ -77,6 +79,7 @@ mod tests {
7779
.add_message(ContextMessage::Text(TextMessage {
7880
role: Role::User,
7981
content: "User message with reasoning".to_string(),
82+
raw_content: None,
8083
tool_calls: None,
8184
model: None,
8285
reasoning_details: Some(reasoning_details),
@@ -110,6 +113,7 @@ mod tests {
110113
let fixture = Context::default().add_message(ContextMessage::Text(TextMessage {
111114
role: Role::Assistant,
112115
content: "Assistant message".to_string(),
116+
raw_content: None,
113117
tool_calls: None,
114118
model: Some(crate::ModelId::new("gpt-4")),
115119
reasoning_details: Some(reasoning_details),
@@ -158,6 +162,7 @@ mod tests {
158162
.add_message(ContextMessage::Text(TextMessage {
159163
role: Role::User,
160164
content: "User with reasoning".to_string(),
165+
raw_content: None,
161166
tool_calls: None,
162167
model: None,
163168
reasoning_details: Some(reasoning_details),

0 commit comments

Comments
 (0)