MythEclipse/zesdex
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

refactor: streamline token counting and message shaping logic

Browse Source
This commit is contained in:
asepharyana
2026-07-17 10:29:52 +07:00
parent 0dfde96f81
commit 5aad7e1eb1
3 changed files with 81 additions and 45 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/zesdex-backend/src/app/runtime/context/shaping.rs
+68 -27
View File
@@ -3,6 +3,8 @@
//! the LLM API. Ported from the former `runtime::shortsend` — behavior
//! is unchanged, only its token-counting now goes through
//! `context::tokens` instead of an inline heuristic.
use std::sync::atomic::{AtomicBool, Ordering};
use super::tokens::count_tokens;
use crate::dto::chat::message::ChatMessage;
@@ -37,6 +39,11 @@ pub fn should_shape(token_estimate: usize, max_wire_tokens: usize, prev_shaped:
/// `[prior conversation compacted]` (or LLM-generated summary, if
/// `client` is `Some`) system message in between.
///
/// The LLM summarization checks the abort flag before calling the LLM,
/// so a user-requested abort is respected promptly. The turn loop already
/// runs on a background thread (spawned in `spawn_turn`), so the blocking
/// summarization call does not freeze the UI.
///
/// Why: keeps context-size overhead roughly constant regardless of
/// session length.
///
@@ -48,6 +55,7 @@ pub fn shape_messages(
max_wire_tokens: usize,
force: bool,
client: Option<&crate::service::provider::LlmClient>,
abort_flag: Option<&AtomicBool>,
) -> Vec<ChatMessage> {
if !force && (token_count <= max_wire_tokens || messages.len() < 5) {
return messages.to_vec();
@@ -88,30 +96,63 @@ pub fn shape_messages(
let mut summary_text = "[prior conversation compacted]".to_string();
if let Some(llm) = client {
let prompt = format!(
"Summarize the following dropped conversation history briefly. Focus on main goals, decisions made, and files modified, so the context is preserved for future turns. Keep it concise.\n\nHistory:\n{}",
dropped_msgs.iter()
.map(|m| format!("[{}]: {}", if m.role == crate::dto::chat::message::Role::User { "User" } else { "Assistant" }, m.content.as_deref().unwrap_or("")))
.collect::<Vec<_>>()
.join("\n\n")
);
// Check abort before starting the blocking LLM summarization call.
// The turn loop already runs on a background thread (spawned in
// `spawn_turn`), so this blocking call does not freeze the UI.
let aborted = abort_flag.is_some_and(|f| f.load(Ordering::SeqCst));
if !aborted {
let prompt = format!(
"You are a context-preservation summarizer. The following conversation \
history is being dropped to free up context window space. \
Produce a structured summary that preserves the information the AI \
agent needs to continue working seamlessly.\n\n\
Structure your summary into these sections:\n\
1. **Goals & Objectives** — what the user asked for, what tasks remain\n\
2. **Key Decisions** — architectural choices, design decisions, approach changes\n\
3. **Files Modified/Created** — paths and brief description of changes\n\
4. **Findings & State** — important discoveries, test results, current state\n\
5. **Open Items** — unresolved issues, pending tasks, next steps\n\n\
Be concise but thorough. Preserve file paths, error messages, and \
specific details the agent needs to continue.\n\n\
History to summarize:\n{}",
dropped_msgs.iter()
.map(|m| {
let role_label = match m.role {
crate::dto::chat::message::Role::User => "User",
crate::dto::chat::message::Role::Assistant => "Assistant",
crate::dto::chat::message::Role::System => "System",
crate::dto::chat::message::Role::Tool => "Tool",
};
let has_tool_calls = m.tool_calls.is_some()
&& m.tool_calls.as_ref().is_some_and(|c| !c.is_empty());
let mut entry = format!("[{role_label}]: {}", m.content.as_deref().unwrap_or(""));
if has_tool_calls {
if let Some(calls) = &m.tool_calls {
let names: Vec<&str> = calls.iter().map(|c| c.function.name.as_str()).collect();
entry.push_str(&format!("\n [tool calls: {}]", names.join(", ")));
}
}
entry
})
.collect::<Vec<_>>()
.join("\n\n---\n\n")
);
let req_msgs = vec![ChatMessage::user(prompt)];
match llm.chat_with_tools_non_streaming(&req_msgs, None) {
Ok(resp) => {
if let Some(content) = resp.0.content {
summary_text =
format!("[Summary of compacted prior conversation:\n{content}\n]");
let req_msgs = vec![ChatMessage::user(prompt)];
match llm.chat_with_tools_non_streaming(&req_msgs, None) {
Ok(resp) => {
if let Some(content) = resp.0.content {
summary_text =
format!("[Summary of compacted prior conversation:\n{content}\n]");
}
}
Err(e) => {
tracing::warn!(
"[context::shaping] LLM summarization failed: {}. \
Falling back to static placeholder.",
e,
);
}
}
Err(e) => {
tracing::warn!(
"[context::shaping] LLM summarization failed: {}. \
Prior conversation history is lost — no summary available. \
This means the model will lose context about earlier parts of \
the conversation.",
e,
);
}
}
}
@@ -150,7 +191,7 @@ mod tests {
ChatMessage::user("hi"),
ChatMessage::assistant(Some("hello".to_string())),
];
let result = shape_messages(&messages, 10, 1000, false, None);
let result = shape_messages(&messages, 10, 1000, false, None, None);
assert_eq!(result.len(), messages.len());
}
@@ -178,7 +219,7 @@ mod tests {
for i in 0..20 {
messages.push(ChatMessage::user(padded_message(i)));
}
let result = shape_messages(&messages, 100_000, 1000, true, None);
let result = shape_messages(&messages, 100_000, 1000, true, None, None);
assert_eq!(result[0].content.as_deref(), Some("system prompt"));
}
@@ -188,7 +229,7 @@ mod tests {
for i in 0..20 {
messages.push(ChatMessage::user(padded_message(i)));
}
let result = shape_messages(&messages, 100_000, 1000, true, None);
let result = shape_messages(&messages, 100_000, 1000, true, None, None);
let has_placeholder = result
.iter()
.any(|m| m.content.as_deref() == Some("[prior conversation compacted]"));
@@ -201,11 +242,11 @@ mod tests {
for i in 0..20 {
messages.push(ChatMessage::user(padded_message(i)));
}
let result = shape_messages(&messages, 100_000, 1000, true, None);
let result = shape_messages(&messages, 100_000, 1000, true, None, None);
let last_content = messages.last().unwrap().content.clone();
assert!(
result.iter().any(|m| m.content == last_content),
"most recent message must survive shaping"
);
}
}
}