tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases Wiki Activity

Compaction overhaul (#5186)

Browse Source 
Co-authored-by: Douwe Osinga <douwe@squareup.com>
Co-authored-by: David Katz <dkatz@squareup.com>
Co-authored-by: Alex Hancock <alexhancock@block.xyz>
This commit is contained in:
Douwe Osinga
2025-10-15 18:39:12 -04:00
committed by GitHub
parent ef9e68f2c6
commit 7c1f5dcb2f
56 changed files with 3069 additions and 4465 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/conversation/message.rs
+15 -57
View File
@@ -112,12 +112,7 @@ pub struct FrontendToolRequest {
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
pub struct ContextLengthExceeded {
pub msg: String,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, ToSchema)]
pub struct SummarizationRequested {
pub struct ConversationCompacted {
pub msg: String,
}
@@ -133,8 +128,7 @@ pub enum MessageContent {
FrontendToolRequest(FrontendToolRequest),
Thinking(ThinkingContent),
RedactedThinking(RedactedThinkingContent),
ContextLengthExceeded(ContextLengthExceeded),
SummarizationRequested(SummarizationRequested),
ConversationCompacted(ConversationCompacted),
}
impl fmt::Display for MessageContent {
@@ -162,10 +156,7 @@ impl fmt::Display for MessageContent {
},
MessageContent::Thinking(t) => write!(f, "[Thinking: {}]", t.thinking),
MessageContent::RedactedThinking(_r) => write!(f, "[RedactedThinking]"),
MessageContent::ContextLengthExceeded(r) => {
write!(f, "[ContextLengthExceeded: {}]", r.msg)
}
MessageContent::SummarizationRequested(r) => {
MessageContent::ConversationCompacted(r) => {
write!(f, "[SummarizationRequested: {}]", r.msg)
}
}
@@ -246,17 +237,13 @@ impl MessageContent {
})
}
pub fn context_length_exceeded<S: Into<String>>(msg: S) -> Self {
MessageContent::ContextLengthExceeded(ContextLengthExceeded { msg: msg.into() })
}
pub fn summarization_requested<S: Into<String>>(msg: S) -> Self {
MessageContent::SummarizationRequested(SummarizationRequested { msg: msg.into() })
pub fn conversation_compacted<S: Into<String>>(msg: S) -> Self {
MessageContent::ConversationCompacted(ConversationCompacted { msg: msg.into() })
}
// Add this new method to check for summarization requested content
pub fn as_summarization_requested(&self) -> Option<&SummarizationRequested> {
if let MessageContent::SummarizationRequested(ref summarization_requested) = self {
pub fn as_summarization_requested(&self) -> Option<&ConversationCompacted> {
if let MessageContent::ConversationCompacted(ref summarization_requested) = self {
Some(summarization_requested)
} else {
None
@@ -390,10 +377,8 @@ impl From<PromptMessage> for Message {
#[serde(rename_all = "camelCase")]
pub struct MessageMetadata {
/// Whether the message should be visible to the user in the UI
#[serde(default = "default_true")]
pub user_visible: bool,
/// Whether the message should be included in the agent's context window
#[serde(default = "default_true")]
pub agent_visible: bool,
}
@@ -464,28 +449,18 @@ impl MessageMetadata {
}
}
fn default_true() -> bool {
true
}
#[derive(ToSchema, Clone, PartialEq, Serialize, Deserialize, Debug)]
/// A message to or from an LLM
#[serde(rename_all = "camelCase")]
pub struct Message {
pub id: Option<String>,
pub role: Role,
#[serde(default = "default_created")]
pub created: i64,
#[serde(deserialize_with = "deserialize_sanitized_content")]
pub content: Vec<MessageContent>,
#[serde(default)]
pub metadata: MessageMetadata,
}
fn default_created() -> i64 {
0 // old messages do not have timestamps.
}
impl Message {
pub fn new(role: Role, created: i64, content: Vec<MessageContent>) -> Self {
Message {
@@ -605,11 +580,6 @@ impl Message {
self.with_content(MessageContent::redacted_thinking(data))
}
/// Add context length exceeded content to the message
pub fn with_context_length_exceeded<S: Into<String>>(self, msg: S) -> Self {
self.with_content(MessageContent::context_length_exceeded(msg))
}
/// Get the concatenated text content of the message, separated by newlines
pub fn as_concat_text(&self) -> String {
self.content
@@ -680,9 +650,8 @@ impl Message {
.all(|c| matches!(c, MessageContent::Text(_)))
}
/// Add summarization requested to the message
pub fn with_summarization_requested<S: Into<String>>(self, msg: S) -> Self {
self.with_content(MessageContent::summarization_requested(msg))
pub fn with_conversation_compacted<S: Into<String>>(self, msg: S) -> Self {
self.with_content(MessageContent::conversation_compacted(msg))
}
/// Set the visibility metadata for the message
@@ -837,7 +806,8 @@ mod tests {
}
}
}
]
],
"metadata": { "agentVisible": true, "userVisible": true }
}"#;
let message: Message = serde_json::from_str(json_str).unwrap();
@@ -1045,7 +1015,8 @@ mod tests {
"data": "base64data",
"mimeType": "image/png"
}}
]
],
"metadata": {{ "agentVisible": true, "userVisible": true }}
}}"#,
malicious_text
);
@@ -1073,7 +1044,8 @@ mod tests {
"content": [{
"type": "text",
"text": "Hello world 世界 🌍"
}]
}],
"metadata": { "agentVisible": true, "userVisible": true }
}"#;
let message: Message = serde_json::from_str(clean_json).unwrap();
@@ -1142,20 +1114,6 @@ mod tests {
let message: Message = serde_json::from_str(json_with_metadata).unwrap();
assert!(!message.is_user_visible());
assert!(message.is_agent_visible());
// Test without metadata (should use defaults)
let json_without_metadata = r#"{
"role": "user",
"created": 1640995200,
"content": [{
"type": "text",
"text": "Test"
}]
}"#;
let message: Message = serde_json::from_str(json_without_metadata).unwrap();
assert!(message.is_user_visible());
assert!(message.is_agent_visible());
}
#[test]