use std::time::Instant;
use anyhow::{Result, ensure};
use crate::engine::display;
use crate::engine::message::ConversationMessage;
use crate::engine::model::{Embedder, Similarity};
use super::extract::extract_context;
use super::types::{
CompactError, CompactProfiler, SUMMARY_BRIEF_HEAD_LINES, SUMMARY_BRIEF_TAIL_LINES, Summary,
SummaryRelevance, compact_profile_enabled,
};
use super::util::{directive_references, load_embedder};
pub fn summarize_with_previous_and_budget(
messages: &[ConversationMessage],
previous_summary: Option<&str>,
max_tokens: usize,
) -> Result<Summary> {
let previous_summary = previous_summary.filter(|summary| !summary.trim().is_empty());
summarize_inner(
messages,
previous_summary,
max_tokens,
compact_profile_enabled(),
)
}
pub(super) fn summarize_inner(
messages: &[ConversationMessage],
previous_summary: Option<&str>,
max_tokens: usize,
print_profile: bool,
) -> Result<Summary> {
let started = Instant::now();
let mut profiler = CompactProfiler::new(print_profile);
let mut embedder = None;
let ctx = extract_context(messages, previous_summary, &mut embedder, &mut profiler)?;
let mut summary = Summary {
goals: ctx.goals,
read: ctx.file_activity.read.values,
modified: ctx.file_activity.modified.values,
created: ctx.file_activity.created.values,
commits: ctx.commits,
outstanding: ctx.outstanding,
preferences: ctx.preferences,
brief: ctx.brief,
tokens_before: estimate_tokens(messages)
+ previous_summary.map_or(0, |summary| (summary.chars().count() / 4) as u64),
};
bound_summary(&mut summary, max_tokens, &mut embedder, &mut profiler)?;
profiler.record("summarize total", started);
Ok(summary)
}
pub(super) fn estimate_tokens(messages: &[ConversationMessage]) -> u64 {
let chars: usize = messages
.iter()
.map(|msg| display::searchable_text(msg).chars().count())
.sum();
(chars / 4) as u64
}
pub fn summary_token_estimate(summary: &Summary) -> usize {
crate::engine::emit::pi_summary_text(summary)
.chars()
.count()
.div_ceil(4)
}
pub(super) fn bound_summary(
summary: &mut Summary,
max_tokens: usize,
embedder: &mut Option<Embedder>,
profiler: &mut CompactProfiler,
) -> Result<()> {
if max_tokens == 0 || summary_token_estimate(summary) <= max_tokens {
return Ok(());
}
let references = directive_references(&summary.goals, &summary.outstanding);
let has_candidates = summary.read.len() > 1
|| summary.modified.len() > 1
|| summary.created.len() > 1
|| summary.brief.lines().count() > SUMMARY_BRIEF_HEAD_LINES + SUMMARY_BRIEF_TAIL_LINES;
let relevance = if references.is_empty() || !has_candidates {
None
} else {
let embedder = load_embedder(embedder, profiler)?;
let started = Instant::now();
let relevance = summary_relevance(summary, &references, embedder)?;
profiler.record("bge budget relevance", started);
Some(relevance)
};
bound_summary_ranked(summary, max_tokens, relevance.as_ref())
}
pub(super) fn summary_relevance(
summary: &Summary,
references: &[&str],
embedder: &Embedder,
) -> Result<SummaryRelevance> {
let brief = summary.brief.lines().map(str::to_owned);
let path_text = |path: &str| {
path.chars()
.map(|ch| match ch {
'/' | '\\' | '.' | '_' | '-' => ' ',
_ => ch,
})
.collect::<String>()
};
let mut candidates = Vec::with_capacity(
summary.read.len()
+ summary.modified.len()
+ summary.created.len()
+ summary.brief.lines().count(),
);
candidates.extend(summary.read.iter().map(|path| path_text(path)));
candidates.extend(summary.modified.iter().map(|path| path_text(path)));
candidates.extend(summary.created.iter().map(|path| path_text(path)));
candidates.extend(brief);
let candidate_refs = candidates.iter().map(String::as_str).collect::<Vec<_>>();
let scores = embedder.relevance(references, &candidate_refs)?;
summary.split_relevance(scores)
}
pub(super) fn bound_summary_ranked(
summary: &mut Summary,
max_tokens: usize,
relevance: Option<&SummaryRelevance>,
) -> Result<()> {
trim_items_to_budget(
summary,
max_tokens,
"read-only paths",
|value| &mut value.read,
relevance.map(|value| value.read.as_slice()),
)?;
trim_items_to_budget(
summary,
max_tokens,
"completed items",
|value| &mut value.commits,
None,
)?;
trim_brief_to_budget(
summary,
max_tokens,
relevance.map(|value| value.brief.as_slice()),
)?;
trim_items_to_budget(
summary,
max_tokens,
"modified paths",
|value| &mut value.modified,
relevance.map(|value| value.modified.as_slice()),
)?;
trim_items_to_budget(
summary,
max_tokens,
"created paths",
|value| &mut value.created,
relevance.map(|value| value.created.as_slice()),
)?;
Ok(())
}
pub(super) fn trim_items_to_budget(
summary: &mut Summary,
max_tokens: usize,
label: &str,
items: fn(&mut Summary) -> &mut Vec<String>,
relevance: Option<&[Similarity]>,
) -> Result<()> {
if summary_token_estimate(summary) <= max_tokens {
return Ok(());
}
let before = summary_token_estimate(summary);
let original = items(summary).clone();
let mut scores = relevance.map(|scores| scores.iter().copied().map(Some).collect::<Vec<_>>());
if let Some(scores) = &scores {
ensure!(
scores.len() == original.len(),
"summary item and relevance counts differ"
);
}
while summary_token_estimate(summary) > max_tokens {
let trimmed = if let Some(scores) = &mut scores {
trim_least_relevant_item(items(summary), scores, 1, label)?
} else {
trim_oldest_item(items(summary), 1, label)
};
if !trimmed {
break;
}
}
if summary_token_estimate(summary) >= before {
*items(summary) = original;
}
Ok(())
}
pub(super) fn trim_brief_to_budget(
summary: &mut Summary,
max_tokens: usize,
relevance: Option<&[Similarity]>,
) -> Result<()> {
if summary_token_estimate(summary) <= max_tokens {
return Ok(());
}
let before = summary_token_estimate(summary);
let original = summary.brief.clone();
let mut scores = relevance.map(|scores| scores.iter().copied().map(Some).collect::<Vec<_>>());
while summary_token_estimate(summary) > max_tokens
&& shrink_brief_once(&mut summary.brief, scores.as_mut())?
{}
if summary_token_estimate(summary) >= before {
summary.brief = original;
}
Ok(())
}
pub(super) fn trim_oldest_item(items: &mut Vec<String>, min_kept: usize, label: &str) -> bool {
let marker_index = items.iter().position(|item| is_omission_marker(item));
let mut omitted = marker_index
.and_then(|index| items.get(index))
.and_then(|item| parse_omission_count(item))
.unwrap_or(0);
if let Some(index) = marker_index {
items.remove(index);
}
if items.len() <= min_kept {
if omitted > 0 {
items.insert(0, omission_marker(omitted, label));
}
return false;
}
items.remove(0);
omitted += 1;
items.insert(0, omission_marker(omitted, label));
true
}
pub(super) fn trim_least_relevant_item(
items: &mut Vec<String>,
relevance: &mut Vec<Option<Similarity>>,
min_kept: usize,
label: &str,
) -> Result<bool> {
ensure!(
items.len() == relevance.len(),
"summary item and relevance counts differ"
);
let marker_index = items.iter().position(|item| is_omission_marker(item));
let mut omitted = marker_index
.and_then(|index| items.get(index))
.and_then(|item| parse_omission_count(item))
.unwrap_or(0);
if let Some(index) = marker_index {
items.remove(index);
relevance.remove(index);
}
if items.len() <= min_kept {
if omitted > 0 {
items.insert(0, omission_marker(omitted, label));
relevance.insert(0, None);
}
return Ok(false);
}
let remove_index = relevance
.iter()
.enumerate()
.filter_map(|(index, score)| score.map(|score| (index, score)))
.min_by_key(|(_, score)| *score)
.map(|(index, _)| index)
.ok_or(CompactError::RelevanceScoresEmpty)?;
items.remove(remove_index);
relevance.remove(remove_index);
omitted += 1;
items.insert(0, omission_marker(omitted, label));
relevance.insert(0, None);
Ok(true)
}
pub(super) fn shrink_brief_once(
brief: &mut String,
relevance: Option<&mut Vec<Option<Similarity>>>,
) -> Result<bool> {
let source: Vec<&str> = brief.lines().collect();
if let Some(scores) = relevance.as_deref() {
ensure!(
source.len() == scores.len(),
"summary brief and relevance counts differ"
);
}
let ranked = relevance.is_some();
let mut omitted = 0;
let mut lines = Vec::with_capacity(source.len());
for (index, line) in source.into_iter().enumerate() {
if is_brief_omission_marker(line.trim()) {
omitted += parse_omission_count(line.trim()).unwrap_or(0);
continue;
}
let score = relevance
.as_deref()
.and_then(|scores| scores.get(index))
.copied()
.flatten();
lines.push((line.to_string(), score));
}
let minimum_lines = SUMMARY_BRIEF_HEAD_LINES + SUMMARY_BRIEF_TAIL_LINES;
if lines.len() <= minimum_lines {
return Ok(false);
}
let removable = lines.len() - SUMMARY_BRIEF_TAIL_LINES;
let remove_index = if ranked {
lines
.iter()
.enumerate()
.skip(SUMMARY_BRIEF_HEAD_LINES)
.take(removable - SUMMARY_BRIEF_HEAD_LINES)
.filter(|(_, (line, _))| {
!(line.is_empty() || line.starts_with('[') && line.ends_with(']'))
})
.filter_map(|(index, (_, score))| score.map(|score| (index, score)))
.min_by_key(|(_, score)| *score)
.map(|(index, _)| index)
} else {
Some(SUMMARY_BRIEF_HEAD_LINES)
};
let Some(remove_index) = remove_index else {
return Ok(false);
};
lines.remove(remove_index);
omitted += 1;
lines.insert(
SUMMARY_BRIEF_HEAD_LINES,
(
format!("...({omitted} earlier lines omitted; use goose_search for details)"),
None,
),
);
*brief = lines
.iter()
.map(|(line, _)| line.as_str())
.collect::<Vec<_>>()
.join("\n");
if let Some(scores) = relevance {
*scores = lines.into_iter().map(|(_, score)| score).collect();
}
Ok(true)
}
pub(super) fn omission_marker(count: usize, label: &str) -> String {
format!("...({count} older {label} omitted; use goose_search for details)")
}
pub(super) fn is_omission_marker(value: &str) -> bool {
value.starts_with("...(") && value.contains(" omitted; use goose_search for details)")
}
pub(super) fn is_brief_omission_marker(value: &str) -> bool {
value.starts_with("...(") && value.contains(" earlier lines omitted")
}
pub(super) fn parse_omission_count(value: &str) -> Option<usize> {
value
.strip_prefix("...(")?
.split_whitespace()
.next()?
.parse()
.ok()
}