ooodc/PicoBot
1
0
Fork 0
Code Pull Requests Activity

feat: 添加字符计数和文本截断功能，增强文本处理能力

Browse Source
This commit is contained in:
ooodc 2026-04-24 14:34:06 +08:00
parent e6f23858b8
commit 4b74fabb98
12 changed files with 321 additions and 43 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

33
src/agent/agent_loop.rs
View File

@ -10,15 +10,13 @@ use crate::providers::{create_provider, LLMProvider, ChatCompletionRequest, Mess
use crate::skills::SkillRuntime;
use crate::storage::SessionStore;
use crate::tools::{ToolContext, ToolRegistry};
use crate::text::{char_count, take_prefix_chars, take_suffix_chars};
use std::collections::VecDeque;
use std::hash::{Hash, Hasher};
use std::io::Read;
use std::sync::Arc;
use std::time::Instant;
/// Maximum characters in a tool result before truncation.
/// Prevents context overflow from large tool outputs.
const MAX_TOOL_RESULT_CHARS: usize = 16_000;
/// Minimum characters to keep when truncating
const TRUNCATION_SUFFIX_LEN: usize = 200;
const MEMORY_TOOL_USAGE_SYSTEM_PROMPT: &str =
@ -94,27 +92,27 @@ fn encode_image_to_base64(path: &str) -> Result<(String, String), std::io::Error
Ok((mime, encoded))
}
/// Truncate tool result if it exceeds MAX_TOOL_RESULT_CHARS.
/// Truncate tool result if it exceeds the configured limit.
/// Preserves the end of the output as it often contains the conclusion/useful result.
fn truncate_tool_result(output: &str) -> String {
let char_count = output.chars().count();
if char_count <= MAX_TOOL_RESULT_CHARS {
fn truncate_tool_result(output: &str, max_tool_result_chars: usize) -> String {
let total_chars = char_count(output);
if total_chars <= max_tool_result_chars {
return output.to_string();
}
let truncated_start_len = char_count.saturating_sub(TRUNCATION_SUFFIX_LEN);
if truncated_start_len > MAX_TOOL_RESULT_CHARS {
let truncated_start_len = total_chars.saturating_sub(TRUNCATION_SUFFIX_LEN);
if truncated_start_len > max_tool_result_chars {
// Even after removing suffix, still too long - take from beginning
let head_len = MAX_TOOL_RESULT_CHARS - 100;
let head: String = output.chars().take(head_len).collect();
let head_len = max_tool_result_chars.saturating_sub(100);
let head = take_prefix_chars(output, head_len);
format!(
"{}...\n\n[Output truncated - {} characters removed]",
head,
char_count - MAX_TOOL_RESULT_CHARS + 100
total_chars - max_tool_result_chars + 100
)
} else {
// Keep most of the end which usually contains the useful result
let tail: String = output.chars().skip(truncated_start_len).collect();
let tail = take_suffix_chars(output, total_chars.saturating_sub(truncated_start_len));
format!(
"...\n\n[Output truncated - {} characters removed]\n\n{}",
truncated_start_len,
@ -486,7 +484,10 @@ impl AgentLoop {
tracing::info!(tool = %tool_call.name, args = %args_str, "Calling tool");
// Truncate tool result if too large
let truncated_output = truncate_tool_result(&result.output);
let truncated_output = truncate_tool_result(
&result.output,
self.provider_config.tool_result_max_chars,
);
// Record tool call and check for loops
let loop_result = loop_detector.record(&tool_call.name, &tool_call.arguments);
@ -894,9 +895,9 @@ mod tests {
#[test]
fn test_truncate_tool_result_handles_utf8_char_boundaries() {
let input = "范".repeat(MAX_TOOL_RESULT_CHARS + 500);
let input = "范".repeat(20_500);
let output = truncate_tool_result(&input);
let output = truncate_tool_result(&input, 20_000);
assert!(output.contains("Output truncated"));
assert!(output.is_char_boundary(output.len()));

50
src/agent/context_compressor.rs
View File

@ -1,6 +1,7 @@
use crate::bus::ChatMessage;
use crate::config::LLMProviderConfig;
use crate::providers::{create_provider, ChatCompletionRequest, Message};
use crate::text::{char_count, take_prefix_chars};
use crate::agent::AgentError;
@ -34,8 +35,8 @@ impl Default for ContextCompressionConfig {
protect_first_n: 1,
protect_last_n: 4,
max_passes: 3,
summary_max_chars: 4000,
tool_result_trim_chars: 2000,
summary_max_chars: 20_000,
tool_result_trim_chars: 20_000,
}
}
}
@ -59,6 +60,17 @@ impl ContextCompressor {
}
}
pub fn from_provider_config(provider_config: &LLMProviderConfig) -> Self {
Self::with_config(
provider_config.token_limit,
ContextCompressionConfig {
summary_max_chars: provider_config.context_summary_max_chars,
tool_result_trim_chars: provider_config.context_tool_result_trim_chars,
..ContextCompressionConfig::default()
},
)
}
/// Create with custom configuration.
pub fn with_config(context_window: usize, config: ContextCompressionConfig) -> Self {
Self {
@ -80,11 +92,12 @@ impl ContextCompressor {
let mut modified = 0;
for msg in messages.iter_mut() {
if msg.role == "tool" && msg.content.len() > limit {
let removed = msg.content.len() - limit;
let content_chars = char_count(&msg.content);
if msg.role == "tool" && content_chars > limit {
let removed = content_chars - limit;
msg.content = format!(
"{}...\n\n[Output truncated - {} characters removed]",
&msg.content[..limit.min(msg.content.len())],
take_prefix_chars(&msg.content, limit),
removed
);
modified += 1;
@ -100,6 +113,7 @@ impl ContextCompressor {
history: Vec<ChatMessage>,
provider_config: &LLMProviderConfig,
) -> Result<Vec<ChatMessage>, AgentError> {
let mut history = history;
// Check if compression is needed
let tokens = estimate_tokens(&history);
if tokens <= self.threshold() {
@ -121,7 +135,7 @@ impl ContextCompressor {
);
// Fast trim pass first
let trimmed = self.fast_trim_tool_results(&mut history.clone());
let trimmed = self.fast_trim_tool_results(&mut history);
if trimmed > 0 {
let tokens_after = estimate_tokens(&history);
#[cfg(debug_assertions)]
@ -271,11 +285,11 @@ impl ContextCompressor {
.join("\n\n");
// Truncate transcript if too long
let transcript = if transcript.len() > self.config.summary_max_chars {
let transcript = if char_count(&transcript) > self.config.summary_max_chars {
format!(
"{}...\n\n[Transcript truncated - {} characters removed]",
&transcript[..self.config.summary_max_chars],
transcript.len() - self.config.summary_max_chars
take_prefix_chars(&transcript, self.config.summary_max_chars),
char_count(&transcript).saturating_sub(self.config.summary_max_chars)
)
} else {
transcript
@ -321,7 +335,7 @@ Be concise, aim for {} characters or less.
Err(e) => {
// Fallback: just truncate the transcript
tracing::warn!(error = %e, "LLM summarization failed, using truncated transcript");
Ok(transcript[..transcript.len().min(2000)].to_string())
Ok(take_prefix_chars(&transcript, self.config.summary_max_chars))
}
}
}
@ -365,6 +379,22 @@ mod tests {
assert!(messages[1].content.len() < 100);
}
#[test]
fn test_fast_trim_handles_utf8_char_boundaries() {
let config = ContextCompressionConfig {
tool_result_trim_chars: 5,
..Default::default()
};
let compressor = ContextCompressor::with_config(100_000, config);
let mut messages = vec![ChatMessage::tool("call1", "bash", &"雅".repeat(20))];
let modified = compressor.fast_trim_tool_results(&mut messages);
assert_eq!(modified, 1);
assert!(messages[0].content.contains("Output truncated"));
assert!(messages[0].content.is_char_boundary(messages[0].content.len()));
}
#[test]
fn test_threshold() {
let compressor = ContextCompressor::new(128_000);

5
src/agent/memory_tool_usage_system_prompt.md
View File

@ -34,12 +34,15 @@
如果用户在聊持续任务、既有偏好、历史决策、项目上下文、曾经纠正过你的内容，或当前请求看起来像是延续之前的事，优先先搜记忆。
## 写入规则
# 记忆写入
## 写入规则
- 写入或修改记忆时，再使用 memory_manage。
- 遇到高价值且未来仍有用的信息时写入记忆：用户长期偏好、稳定事实、用户对你的纠正、持续任务或项目上下文、明确决策等。
- 写入时优先使用规范 namespace：preferences、profile、tasks、decisions。
- 优先调用 memory_manage(action='put')；同一 namespace/key 可直接覆盖更新。
## 注意
- 如果你决定不再调用工具，则反思一下是否使用memory_manage
### 以下场景视为高价值加分
- 用户多次交互优化输出

28
src/channels/feishu.rs
View File

@ -13,6 +13,7 @@ use tokio::sync::{broadcast, RwLock};
use crate::bus::{MessageBus, MediaItem, OutboundMessage};
use crate::channels::base::{Channel, ChannelError};
use crate::config::{FeishuChannelConfig, LLMProviderConfig};
use crate::text::{char_count, truncate_with_ellipsis};
const FEISHU_API_BASE: &str = "https://open.feishu.cn/open-apis";
const FEISHU_WS_BASE: &str = "https://open.feishu.cn";
@ -670,8 +671,6 @@ impl FeishuChannel {
Ok(())
}
const REPLY_CONTEXT_MAX_LEN: usize = 500;
/// Fetch the text content of a Feishu message by ID.
/// Returns a "[Reply to: ...]" context string, or None on failure.
async fn get_message_content(&self, message_id: &str) -> Option<String> {
@ -752,8 +751,8 @@ impl FeishuChannel {
return None;
}
let text = if text.len() > Self::REPLY_CONTEXT_MAX_LEN {
format!("{}...", &text[..Self::REPLY_CONTEXT_MAX_LEN])
let text = if char_count(&text) > self.config.reply_context_max_chars {
truncate_with_ellipsis(&text, self.config.reply_context_max_chars)
} else {
text
};
@ -765,13 +764,12 @@ impl FeishuChannel {
async fn send_message_to_feishu(&self, receive_id: &str, receive_id_type: &str, msg_type: &str, content: &str) -> Result<(), ChannelError> {
let token = self.get_tenant_access_token().await?;
// Feishu text messages have content limits (~64KB).
// Truncate if content is too long to avoid API error 230001.
const MAX_TEXT_LENGTH: usize = 60_000;
let payload_content = if msg_type == "text" {
let truncated = if content.len() > MAX_TEXT_LENGTH {
format!("{}...\n\n[Content truncated due to length limit]", &content[..MAX_TEXT_LENGTH])
let truncated = if char_count(content) > self.config.max_message_chars {
format!(
"{}\n\n[Content truncated due to length limit]",
truncate_with_ellipsis(content, self.config.max_message_chars)
)
} else {
content.to_string()
};
@ -779,9 +777,15 @@ impl FeishuChannel {
} else {
// For post messages, content is already JSON (from markdown_to_post)
// But we still need to check length
if content.len() > MAX_TEXT_LENGTH {
if char_count(content) > self.config.max_message_chars {
// Fallback to truncated text for post as well
serde_json::json!({ "text": format!("{}...\n\n[Content truncated due to length limit]", &content[..MAX_TEXT_LENGTH]) }).to_string()
serde_json::json!({
"text": format!(
"{}\n\n[Content truncated due to length limit]",
truncate_with_ellipsis(content, self.config.max_message_chars)
)
})
.to_string()
} else {
content.to_string()
}

172
src/config/mod.rs
View File

@ -106,6 +106,10 @@ pub struct FeishuChannelConfig {
/// Emoji type for message reactions (e.g. "THUMBSUP", "OK", "EYES").
#[serde(default = "default_reaction_emoji")]
pub reaction_emoji: String,
#[serde(default = "default_channel_max_message_chars")]
pub max_message_chars: usize,
#[serde(default = "default_channel_reply_context_max_chars")]
pub reply_context_max_chars: usize,
}
fn default_allow_from() -> Vec<String> {
@ -121,6 +125,14 @@ fn default_reaction_emoji() -> String {
"Typing".to_string()
}
fn default_channel_max_message_chars() -> usize {
20_000
}
fn default_channel_reply_context_max_chars() -> usize {
20_000
}
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct ProviderConfig {
#[serde(rename = "type")]
@ -152,6 +164,12 @@ pub struct AgentConfig {
pub max_tool_iterations: usize,
#[serde(default = "default_token_limit")]
pub token_limit: usize,
#[serde(default = "default_tool_result_max_chars")]
pub tool_result_max_chars: usize,
#[serde(default = "default_context_summary_max_chars")]
pub context_summary_max_chars: usize,
#[serde(default = "default_context_tool_result_trim_chars")]
pub context_tool_result_trim_chars: usize,
}
fn default_max_tool_iterations() -> usize {
@ -162,6 +180,18 @@ fn default_token_limit() -> usize {
128_000
}
fn default_tool_result_max_chars() -> usize {
20_000
}
fn default_context_summary_max_chars() -> usize {
20_000
}
fn default_context_tool_result_trim_chars() -> usize {
20_000
}
fn default_llm_timeout_secs() -> u64 {
120
}
@ -495,6 +525,9 @@ pub struct LLMProviderConfig {
pub model_extra: HashMap<String, serde_json::Value>,
pub max_tool_iterations: usize,
pub token_limit: usize,
pub tool_result_max_chars: usize,
pub context_summary_max_chars: usize,
pub context_tool_result_trim_chars: usize,
}
fn get_default_config_path() -> PathBuf {
@ -558,6 +591,9 @@ impl Config {
model_extra: model.extra.clone(),
max_tool_iterations: agent.max_tool_iterations,
token_limit: agent.token_limit,
tool_result_max_chars: agent.tool_result_max_chars,
context_summary_max_chars: agent.context_summary_max_chars,
context_tool_result_trim_chars: agent.context_tool_result_trim_chars,
})
}
}
@ -695,6 +731,9 @@ mod tests {
assert_eq!(provider_config.model_id, "qwen-plus");
assert_eq!(provider_config.temperature, Some(0.0));
assert_eq!(provider_config.llm_timeout_secs, 120);
assert_eq!(provider_config.tool_result_max_chars, 20_000);
assert_eq!(provider_config.context_summary_max_chars, 20_000);
assert_eq!(provider_config.context_tool_result_trim_chars, 20_000);
}
#[test]
@ -917,6 +956,139 @@ mod tests {
let config = Config::load(file.path().to_str().unwrap()).unwrap();
assert_eq!(config.agents["default"].max_tool_iterations, 100);
assert_eq!(config.agents["default"].tool_result_max_chars, 20_000);
assert_eq!(config.agents["default"].context_summary_max_chars, 20_000);
assert_eq!(config.agents["default"].context_tool_result_trim_chars, 20_000);
}
#[test]
fn test_agent_config_loads_custom_truncation_limits() {
let file = tempfile::NamedTempFile::new().unwrap();
std::fs::write(
file.path(),
r#"{
"providers": {
"aliyun": {
"type": "openai",
"base_url": "https://example.invalid/v1",
"api_key": "test-key",
"extra_headers": {}
}
},
"models": {
"qwen-plus": {
"model_id": "qwen-plus"
}
},
"agents": {
"default": {
"provider": "aliyun",
"model": "qwen-plus",
"tool_result_max_chars": 1234,
"context_summary_max_chars": 2345,
"context_tool_result_trim_chars": 3456
}
}
}"#,
)
.unwrap();
let config = Config::load(file.path().to_str().unwrap()).unwrap();
let agent = &config.agents["default"];
let provider_config = config.get_provider_config("default").unwrap();
assert_eq!(agent.tool_result_max_chars, 1234);
assert_eq!(agent.context_summary_max_chars, 2345);
assert_eq!(agent.context_tool_result_trim_chars, 3456);
assert_eq!(provider_config.tool_result_max_chars, 1234);
assert_eq!(provider_config.context_summary_max_chars, 2345);
assert_eq!(provider_config.context_tool_result_trim_chars, 3456);
}
#[test]
fn test_feishu_channel_config_defaults_truncation_limits() {
let file = tempfile::NamedTempFile::new().unwrap();
std::fs::write(
file.path(),
r#"{
"providers": {
"aliyun": {
"type": "openai",
"base_url": "https://example.invalid/v1",
"api_key": "test-key",
"extra_headers": {}
}
},
"models": {
"qwen-plus": {
"model_id": "qwen-plus"
}
},
"agents": {
"default": {
"provider": "aliyun",
"model": "qwen-plus"
}
},
"channels": {
"feishu": {
"enabled": true,
"app_id": "app-id",
"app_secret": "secret"
}
}
}"#,
)
.unwrap();
let config = Config::load(file.path().to_str().unwrap()).unwrap();
let feishu = &config.channels["feishu"];
assert_eq!(feishu.max_message_chars, 20_000);
assert_eq!(feishu.reply_context_max_chars, 20_000);
}
#[test]
fn test_feishu_channel_config_loads_custom_truncation_limits() {
let file = tempfile::NamedTempFile::new().unwrap();
std::fs::write(
file.path(),
r#"{
"providers": {
"aliyun": {
"type": "openai",
"base_url": "https://example.invalid/v1",
"api_key": "test-key",
"extra_headers": {}
}
},
"models": {
"qwen-plus": {
"model_id": "qwen-plus"
}
},
"agents": {
"default": {
"provider": "aliyun",
"model": "qwen-plus"
}
},
"channels": {
"feishu": {
"enabled": true,
"app_id": "app-id",
"app_secret": "secret",
"max_message_chars": 3456,
"reply_context_max_chars": 4567
}
}
}"#,
)
.unwrap();
let config = Config::load(file.path().to_str().unwrap()).unwrap();
let feishu = &config.channels["feishu"];
assert_eq!(feishu.max_message_chars, 3456);
assert_eq!(feishu.reply_context_max_chars, 4567);
}
#[test]

18
src/gateway/default_agent_prompt.md
View File

@ -13,6 +13,17 @@
- 当现有工具是完成任务的最直接方式时，优先使用工具。
- 除非用户明确要求改变方向，否则保持用户原本目标不变。
## 记忆处理
在大多数请求开始时，先检索长期记忆，再决定如何回答或是否写入。
如果当前任务还需要其它彼此独立的只读工具，可以和记忆搜索同一轮一起调用。
在完成了一个任务之后，自动将关键的信息保存到记忆中，特别是重要内容
- 用户纠正了你
- 用户的情感很强烈
- 多次出现的内容
- 客观的信息
- 用户的偏好
## 助理原则
- 优先解决问题，而不是展示过程。
@ -27,9 +38,12 @@
- 除非用户另有要求，否则使用中文回复。
- 默认短而清楚，按信息密度组织内容。
- 如果任务涉及文件、命令、配置或下一步操作，优先给出最关键的那部分。
- 如果存在限制、风险或前提条件，要直接说明。
- 在信息不足时先补关键前提，在信息充分时直接执行
## PICO配置
- 默认路径为[basedir]:~/.picobot
- Skill安装在[basedir]/skills
## 补充要求
- 回答应以帮助用户完成当前目标为中心。
- 在信息不足时先补关键前提，在信息充分时直接执行。

23
src/gateway/session.rs
View File

@ -482,7 +482,7 @@ impl Session {
provider_config: provider_config.clone(),
tools,
skills,
compressor: ContextCompressor::new(provider_config.token_limit),
compressor: ContextCompressor::from_provider_config(&provider_config),
store,
agent_prompt_reinject_every: agent_prompt_reinject_every as i64,
})
@ -1480,6 +1480,9 @@ mod tests {
model_extra: HashMap::new(),
max_tool_iterations: 1,
token_limit: 4096,
tool_result_max_chars: 20_000,
context_summary_max_chars: 20_000,
context_tool_result_trim_chars: 20_000,
}
}
@ -1497,6 +1500,9 @@ mod tests {
model_extra: HashMap::new(),
max_tool_iterations: 1,
token_limit: 4096,
tool_result_max_chars: 20_000,
context_summary_max_chars: 20_000,
context_tool_result_trim_chars: 20_000,
}
}
@ -1703,6 +1709,9 @@ mod tests {
max_tool_iterations: 1,
token_limit: 4096,
llm_timeout_secs: 30,
tool_result_max_chars: 20_000,
context_summary_max_chars: 20_000,
context_tool_result_trim_chars: 20_000,
};
let session_manager = SessionManager::new(
@ -1740,6 +1749,9 @@ mod tests {
max_tool_iterations: 1,
token_limit: 4096,
llm_timeout_secs: 30,
tool_result_max_chars: 20_000,
context_summary_max_chars: 20_000,
context_tool_result_trim_chars: 20_000,
};
let planner_provider = LLMProviderConfig {
model_id: "planner-model".to_string(),
@ -1823,6 +1835,9 @@ mod tests {
max_tool_iterations: 1,
token_limit: 4096,
llm_timeout_secs: 30,
tool_result_max_chars: 20_000,
context_summary_max_chars: 20_000,
context_tool_result_trim_chars: 20_000,
};
let session_manager = SessionManager::new(
@ -1911,6 +1926,9 @@ mod tests {
max_tool_iterations: 1,
token_limit: 4096,
llm_timeout_secs: 30,
tool_result_max_chars: 20_000,
context_summary_max_chars: 20_000,
context_tool_result_trim_chars: 20_000,
};
let session_manager = SessionManager::new(
@ -1970,6 +1988,9 @@ mod tests {
max_tool_iterations: 1,
token_limit: 4096,
llm_timeout_secs: 30,
tool_result_max_chars: 20_000,
context_summary_max_chars: 20_000,
context_tool_result_trim_chars: 20_000,
};
let session_manager = SessionManager::new(

1
src/lib.rs
View File

@ -1,4 +1,5 @@
pub mod config;
pub mod text;
pub mod providers;
pub mod bus;
pub mod cli;

6
src/scheduler/mod.rs
View File

@ -850,6 +850,9 @@ mod tests {
model_extra: HashMap::new(),
token_limit: 4096,
max_tool_iterations: 4,
tool_result_max_chars: 20_000,
context_summary_max_chars: 20_000,
context_tool_result_trim_chars: 20_000,
};
let session_manager = SessionManager::new(
4,
@ -899,6 +902,9 @@ mod tests {
model_extra: HashMap::new(),
token_limit: 4096,
max_tool_iterations: 4,
tool_result_max_chars: 20_000,
context_summary_max_chars: 20_000,
context_tool_result_trim_chars: 20_000,
};
let session_manager = SessionManager::new(
4,

20
src/text.rs Normal file
View File

@ -0,0 +1,20 @@
pub fn char_count(text: &str) -> usize {
text.chars().count()
}
pub fn take_prefix_chars(text: &str, max_chars: usize) -> String {
text.chars().take(max_chars).collect()
}
pub fn take_suffix_chars(text: &str, max_chars: usize) -> String {
let count = char_count(text);
text.chars().skip(count.saturating_sub(max_chars)).collect()
}
pub fn truncate_with_ellipsis(text: &str, max_chars: usize) -> String {
if char_count(text) <= max_chars {
return text.to_string();
}
format!("{}...", take_prefix_chars(text, max_chars))
}

3
tests/test_integration.rs
View File

@ -26,6 +26,9 @@ fn load_config() -> Option<LLMProviderConfig> {
model_extra: HashMap::new(),
max_tool_iterations: 20,
token_limit: 128_000,
tool_result_max_chars: 20_000,
context_summary_max_chars: 20_000,
context_tool_result_trim_chars: 20_000,
})
}

3
tests/test_tool_calling.rs
View File

@ -26,6 +26,9 @@ fn load_openai_config() -> Option<LLMProviderConfig> {
model_extra: HashMap::new(),
max_tool_iterations: 20,
token_limit: 128_000,
tool_result_max_chars: 20_000,
context_summary_max_chars: 20_000,
context_tool_result_trim_chars: 20_000,
})
}