feat: Support specifying the compression model to use when compressing a session for greater cost controls
This commit is contained in:
7 files changed
+118
-4
No files matched your search
@@ -635,6 +635,10 @@ impl Agent {
|
||||
self.config.compression_keep_last
|
||||
}
|
||||
|
||||
pub fn compression_model(&self) -> Option<&str> {
|
||||
self.config.compression_model.as_deref()
|
||||
}
|
||||
|
||||
pub fn is_dynamic_instructions(&self) -> bool {
|
||||
self.config.dynamic_instructions
|
||||
}
|
||||
@@ -832,6 +836,8 @@ pub struct AgentConfig {
|
||||
pub max_concurrent_jobs: Option<usize>,
|
||||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||||
pub compression_keep_last: Option<usize>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub compression_model: Option<String>,
|
||||
#[serde(default)]
|
||||
pub description: String,
|
||||
#[serde(default)]
|
||||
|
||||
@@ -66,6 +66,7 @@ pub struct AppConfig {
|
||||
pub save_session: Option<bool>,
|
||||
pub compression_threshold: usize,
|
||||
pub compression_keep_last: usize,
|
||||
pub compression_model: Option<String>,
|
||||
pub summarization_prompt: Option<String>,
|
||||
pub summary_context_prompt: Option<String>,
|
||||
pub max_tool_result_chars: Option<usize>,
|
||||
@@ -156,6 +157,7 @@ impl Default for AppConfig {
|
||||
save_session: None,
|
||||
compression_threshold: 4000,
|
||||
compression_keep_last: 0,
|
||||
compression_model: None,
|
||||
summarization_prompt: None,
|
||||
summary_context_prompt: None,
|
||||
max_tool_result_chars: None,
|
||||
@@ -247,6 +249,7 @@ impl AppConfig {
|
||||
save_session: config.save_session,
|
||||
compression_threshold: config.compression_threshold,
|
||||
compression_keep_last: config.compression_keep_last,
|
||||
compression_model: config.compression_model,
|
||||
summarization_prompt: config.summarization_prompt,
|
||||
summary_context_prompt: config.summary_context_prompt,
|
||||
max_tool_result_chars: config.max_tool_result_chars,
|
||||
|
||||
@@ -317,6 +317,10 @@ impl Input {
|
||||
&self.role
|
||||
}
|
||||
|
||||
pub fn set_role_model(&mut self, model: Model) {
|
||||
self.role.set_model(model);
|
||||
}
|
||||
|
||||
pub fn session<'a>(&self, session: &'a Option<Session>) -> Option<&'a Session> {
|
||||
if self.with_session {
|
||||
session.as_ref()
|
||||
@@ -858,6 +862,17 @@ mod tests {
|
||||
assert!(input.session(&session).is_some());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn set_role_model_swaps_model_preserving_with_session() {
|
||||
let mut ctx = create_test_ctx();
|
||||
ctx.session = Some(Session::default());
|
||||
let mut input = Input::from_str(&ctx, "test", None).unwrap();
|
||||
assert!(input.with_session());
|
||||
input.set_role_model(Model::new("provider", "compress-model"));
|
||||
assert_eq!(input.role().model().id(), "provider:compress-model");
|
||||
assert!(input.with_session());
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn is_image_recognizes_image_extensions() {
|
||||
assert!(is_image("photo.png"));
|
||||
|
||||
@@ -276,6 +276,7 @@ pub struct Config {
|
||||
pub save_session: Option<bool>,
|
||||
pub compression_threshold: usize,
|
||||
pub compression_keep_last: usize,
|
||||
pub compression_model: Option<String>,
|
||||
pub summarization_prompt: Option<String>,
|
||||
pub summary_context_prompt: Option<String>,
|
||||
pub max_tool_result_chars: Option<usize>,
|
||||
@@ -363,6 +364,7 @@ impl Default for Config {
|
||||
save_session: None,
|
||||
compression_threshold: 4000,
|
||||
compression_keep_last: 0,
|
||||
compression_model: None,
|
||||
summarization_prompt: None,
|
||||
summary_context_prompt: None,
|
||||
max_tool_result_chars: None,
|
||||
@@ -1300,6 +1302,7 @@ clients:
|
||||
assert!(cfg.function_calling_support);
|
||||
assert!(cfg.mcp_server_support);
|
||||
assert_eq!(cfg.compression_threshold, 4000);
|
||||
assert!(cfg.compression_model.is_none());
|
||||
assert_eq!(cfg.rag_top_k, 5);
|
||||
assert!(cfg.save_shell_history);
|
||||
assert_eq!(cfg.keybindings, "emacs");
|
||||
|
||||
@@ -1006,6 +1006,13 @@ impl RequestContext {
|
||||
.unwrap_or(self.app.config.compression_keep_last)
|
||||
}
|
||||
|
||||
pub fn compression_model(&self) -> Option<String> {
|
||||
self.agent
|
||||
.as_ref()
|
||||
.and_then(|a| a.compression_model().map(|s| s.to_string()))
|
||||
.or_else(|| self.app.config.compression_model.clone())
|
||||
}
|
||||
|
||||
pub fn role_like_mut(&mut self) -> Option<&mut dyn RoleLike> {
|
||||
if let Some(session) = self.session.as_mut() {
|
||||
Some(session)
|
||||
@@ -5129,10 +5136,45 @@ impl RequestContext {
|
||||
.summarization_prompt
|
||||
.clone()
|
||||
.unwrap_or_else(|| SUMMARIZATION_PROMPT.into());
|
||||
let input = Input::from_str(self, &prompt, None)?;
|
||||
let summary = tokio::time::timeout(Duration::from_secs(120), input.fetch_chat_text())
|
||||
.await
|
||||
.map_err(|_| anyhow::anyhow!("Compression LLM call timed out after 120 s"))??;
|
||||
|
||||
async fn timed_fetch(input: &Input) -> Result<String> {
|
||||
tokio::time::timeout(Duration::from_secs(120), input.fetch_chat_text())
|
||||
.await
|
||||
.map_err(|_| anyhow::anyhow!("Compression LLM call timed out after 120 s"))?
|
||||
}
|
||||
|
||||
let mut input = Input::from_str(self, &prompt, None)?;
|
||||
let compression_model_id = self
|
||||
.compression_model()
|
||||
.filter(|id| *id != self.current_model().id());
|
||||
let summary = match compression_model_id {
|
||||
Some(model_id) => {
|
||||
match Model::retrieve_model(self.app.config.as_ref(), &model_id, ModelType::Chat) {
|
||||
Ok(model) => {
|
||||
input.set_role_model(model);
|
||||
match timed_fetch(&input).await {
|
||||
Ok(summary) => summary,
|
||||
Err(err) => {
|
||||
eprintln!(
|
||||
"Warning: compression model '{model_id}' failed: {err}; falling back to the current model '{}' for session compression",
|
||||
self.current_model().id()
|
||||
);
|
||||
let input = Input::from_str(self, &prompt, None)?;
|
||||
timed_fetch(&input).await?
|
||||
}
|
||||
}
|
||||
}
|
||||
Err(err) => {
|
||||
eprintln!(
|
||||
"Warning: compression model '{model_id}' could not be used ({err}); falling back to the current model '{}' for session compression",
|
||||
self.current_model().id()
|
||||
);
|
||||
timed_fetch(&input).await?
|
||||
}
|
||||
}
|
||||
}
|
||||
None => timed_fetch(&input).await?,
|
||||
};
|
||||
let summary_context_prompt = self
|
||||
.app
|
||||
.config
|
||||
@@ -5800,6 +5842,49 @@ mod tests {
|
||||
assert_eq!(extracted.reasoning_effort(), Some("low".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compression_model_none_when_unset() {
|
||||
let ctx = create_test_ctx();
|
||||
assert_eq!(ctx.compression_model(), None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compression_model_uses_app_config_without_agent() {
|
||||
let mut ctx = create_test_ctx();
|
||||
ctx.app = Arc::new(AppState {
|
||||
config: Arc::new(AppConfig {
|
||||
compression_model: Some("openai:app-model".to_string()),
|
||||
..(*ctx.app.config).clone()
|
||||
}),
|
||||
..(*ctx.app).clone()
|
||||
});
|
||||
assert_eq!(
|
||||
ctx.compression_model(),
|
||||
Some("openai:app-model".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn compression_model_agent_overrides_app_config() {
|
||||
let mut ctx = create_test_ctx();
|
||||
ctx.app = Arc::new(AppState {
|
||||
config: Arc::new(AppConfig {
|
||||
compression_model: Some("openai:app-model".to_string()),
|
||||
..(*ctx.app.config).clone()
|
||||
}),
|
||||
..(*ctx.app).clone()
|
||||
});
|
||||
ctx.agent = Some(Agent::test_new(AgentConfig {
|
||||
name: "test-agent".to_string(),
|
||||
compression_model: Some("openai:agent-model".to_string()),
|
||||
..AgentConfig::default()
|
||||
}));
|
||||
assert_eq!(
|
||||
ctx.compression_model(),
|
||||
Some("openai:agent-model".to_string())
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn should_inject_skill_instructions_requires_function_calling() {
|
||||
let app = AppConfig {
|
||||
|
||||
Reference in new issue
Block a user