ooodc/PicoBot
1
0
Fork 0
Code Pull Requests Activity
PicoBot/src/bus/message.rs
ooodc 3045a6b596 feat: Enhance ChatMessage with system context and background compaction 
- Added `system_context` field to `ChatMessage` for better message context handling.
- Introduced constants for system context prompts in `message.rs`.
- Updated `Session` to manage background history compaction, including methods to start and finish compaction.
- Implemented logic to schedule background compaction after message processing in `SessionManager`.
- Enhanced database schema to support new `system_context` field in messages.
- Added functionality to compact active history, preserving system messages and summaries.
- Updated tests to validate new compaction logic and ensure message integrity.
- Removed unused functions and cleaned up code in various modules for better maintainability.

Co-authored-by: Copilot <copilot@github.com>
2026-04-26 09:31:13 +08:00

616 lines
20 KiB
Rust
Raw Blame History

use std::collections::HashMap;
use serde::{Deserialize, Serialize};
use crate::providers::ToolCall;
pub const SYSTEM_CONTEXT_AGENT_PROMPT: &str = "agent_prompt";
pub const SYSTEM_CONTEXT_SCHEDULED_PROMPT: &str = "scheduled_system_prompt";
pub const SYSTEM_CONTEXT_HISTORY_COMPACTION: &str = "history_compaction";
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ToolMessageState {
Completed,
PendingUserAction,
}
// ============================================================================
// ContentBlock - Multimodal content representation (OpenAI-style)
// ============================================================================
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ContentBlock {
#[serde(rename = "text")]
Text { text: String },
#[serde(rename = "image_url")]
ImageUrl { image_url: ImageUrlBlock },
}
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ImageUrlBlock {
pub url: String,
}
impl ContentBlock {
pub fn text(content: impl Into<String>) -> Self {
Self::Text { text: content.into() }
}
pub fn image_url(url: impl Into<String>) -> Self {
Self::ImageUrl {
image_url: ImageUrlBlock { url: url.into() },
}
}
}
// ============================================================================
// MediaItem - Media metadata for messages
// ============================================================================
#[derive(Debug, Clone)]
pub struct MediaItem {
pub path: String, // Local file path
pub media_type: String, // "image", "audio", "file", "video"
pub mime_type: Option<String>,
pub original_key: Option<String>, // Feishu file_key for download
}
impl MediaItem {
pub fn new(path: impl Into<String>, media_type: impl Into<String>) -> Self {
Self {
path: path.into(),
media_type: media_type.into(),
mime_type: None,
original_key: None,
}
}
}
// ============================================================================
// ChatMessage - Used by AgentLoop for LLM conversation history
// ============================================================================
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ChatMessage {
pub id: String,
pub role: String,
pub content: String,
pub media_refs: Vec<String>, // Paths to media files for context
pub timestamp: i64,
#[serde(skip_serializing_if = "Option::is_none")]
pub system_context: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub reasoning_content: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_call_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_state: Option<ToolMessageState>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_calls: Option<Vec<ToolCall>>,
}
impl ChatMessage {
pub fn user(content: impl Into<String>) -> Self {
Self {
id: uuid::Uuid::new_v4().to_string(),
role: "user".to_string(),
content: content.into(),
media_refs: Vec::new(),
timestamp: current_timestamp(),
system_context: None,
reasoning_content: None,
tool_call_id: None,
tool_name: None,
tool_state: None,
tool_calls: None,
}
}
pub fn user_with_media(content: impl Into<String>, media_refs: Vec<String>) -> Self {
Self {
id: uuid::Uuid::new_v4().to_string(),
role: "user".to_string(),
content: content.into(),
media_refs,
timestamp: current_timestamp(),
system_context: None,
reasoning_content: None,
tool_call_id: None,
tool_name: None,
tool_state: None,
tool_calls: None,
}
}
pub fn assistant(content: impl Into<String>) -> Self {
Self {
id: uuid::Uuid::new_v4().to_string(),
role: "assistant".to_string(),
content: content.into(),
media_refs: Vec::new(),
timestamp: current_timestamp(),
system_context: None,
reasoning_content: None,
tool_call_id: None,
tool_name: None,
tool_state: None,
tool_calls: None,
}
}
pub fn assistant_with_reasoning(
content: impl Into<String>,
reasoning_content: impl Into<String>,
) -> Self {
let mut message = Self::assistant(content);
message.reasoning_content = Some(reasoning_content.into());
message
}
pub fn assistant_with_tool_calls(content: impl Into<String>, tool_calls: Vec<ToolCall>) -> Self {
Self {
id: uuid::Uuid::new_v4().to_string(),
role: "assistant".to_string(),
content: content.into(),
media_refs: Vec::new(),
timestamp: current_timestamp(),
system_context: None,
reasoning_content: None,
tool_call_id: None,
tool_name: None,
tool_state: None,
tool_calls: Some(tool_calls),
}
}
pub fn assistant_with_tool_calls_and_reasoning(
content: impl Into<String>,
tool_calls: Vec<ToolCall>,
reasoning_content: impl Into<String>,
) -> Self {
let mut message = Self::assistant_with_tool_calls(content, tool_calls);
message.reasoning_content = Some(reasoning_content.into());
message
}
pub fn system(content: impl Into<String>) -> Self {
Self::system_with_context(content, None::<String>)
}
pub fn system_with_context(
content: impl Into<String>,
system_context: impl Into<Option<String>>,
) -> Self {
Self {
id: uuid::Uuid::new_v4().to_string(),
role: "system".to_string(),
content: content.into(),
media_refs: Vec::new(),
timestamp: current_timestamp(),
system_context: system_context.into(),
reasoning_content: None,
tool_call_id: None,
tool_name: None,
tool_state: None,
tool_calls: None,
}
}
pub fn tool(tool_call_id: impl Into<String>, tool_name: impl Into<String>, content: impl Into<String>) -> Self {
Self::tool_with_state(tool_call_id, tool_name, content, ToolMessageState::Completed)
}
pub fn tool_with_state(
tool_call_id: impl Into<String>,
tool_name: impl Into<String>,
content: impl Into<String>,
tool_state: ToolMessageState,
) -> Self {
Self {
id: uuid::Uuid::new_v4().to_string(),
role: "tool".to_string(),
content: content.into(),
media_refs: Vec::new(),
timestamp: current_timestamp(),
system_context: None,
reasoning_content: None,
tool_call_id: Some(tool_call_id.into()),
tool_name: Some(tool_name.into()),
tool_state: Some(tool_state),
tool_calls: None,
}
}
pub fn has_system_context(&self, expected: &str) -> bool {
self.system_context.as_deref() == Some(expected)
}
pub fn is_assistant_tool_call_message(&self) -> bool {
self.role == "assistant"
&& self
.tool_calls
.as_ref()
.map(|calls| !calls.is_empty())
.unwrap_or(false)
}
}
// ============================================================================
// InboundMessage - Message from Channel to Bus (user input)
// ============================================================================
#[derive(Debug, Clone)]
pub struct InboundMessage {
pub channel: String,
pub sender_id: String,
pub chat_id: String,
pub content: String,
pub timestamp: i64,
pub media: Vec<MediaItem>,
/// Channel-specific data used internally by the channel (not forwarded).
pub metadata: HashMap<String, String>,
/// Data forwarded from inbound to outbound (copied to OutboundMessage.metadata by gateway).
pub forwarded_metadata: HashMap<String, String>,
}
impl InboundMessage {
pub fn session_key(&self) -> String {
format!("{}:{}", self.channel, self.chat_id)
}
}
// ============================================================================
// OutboundMessage - Message from Agent to Channel (bot response)
// ============================================================================
#[derive(Debug, Clone)]
pub struct OutboundMessage {
pub channel: String,
pub chat_id: String,
pub content: String,
pub reply_to: Option<String>,
pub media: Vec<MediaItem>,
pub metadata: HashMap<String, String>,
pub event_kind: OutboundEventKind,
pub role: String,
pub tool_call_id: Option<String>,
pub tool_name: Option<String>,
pub tool_arguments: Option<serde_json::Value>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum OutboundEventKind {
AssistantResponse,
ToolCall,
ToolResult,
ToolPending,
}
impl OutboundMessage {
pub fn is_stream_delta(&self) -> bool {
self.metadata.get("_stream_delta").is_some()
}
pub fn assistant(
channel: impl Into<String>,
chat_id: impl Into<String>,
content: impl Into<String>,
reply_to: Option<String>,
metadata: HashMap<String, String>,
) -> Self {
Self {
channel: channel.into(),
chat_id: chat_id.into(),
content: content.into(),
reply_to,
media: Vec::new(),
metadata,
event_kind: OutboundEventKind::AssistantResponse,
role: "assistant".to_string(),
tool_call_id: None,
tool_name: None,
tool_arguments: None,
}
}
pub fn tool_call(
channel: impl Into<String>,
chat_id: impl Into<String>,
message_id: impl Into<String>,
tool_name: impl Into<String>,
tool_arguments: serde_json::Value,
reply_to: Option<String>,
metadata: HashMap<String, String>,
) -> Self {
let tool_name = tool_name.into();
let content = format_tool_call_content(&tool_name, &tool_arguments);
Self {
channel: channel.into(),
chat_id: chat_id.into(),
content,
reply_to,
media: Vec::new(),
metadata,
event_kind: OutboundEventKind::ToolCall,
role: "assistant".to_string(),
tool_call_id: Some(message_id.into()),
tool_name: Some(tool_name),
tool_arguments: Some(tool_arguments),
}
}
pub fn tool_result(
channel: impl Into<String>,
chat_id: impl Into<String>,
tool_call_id: impl Into<String>,
tool_name: impl Into<String>,
content: impl Into<String>,
reply_to: Option<String>,
metadata: HashMap<String, String>,
) -> Self {
let tool_name = tool_name.into();
let raw_content = content.into();
let content = format_tool_result_content(&tool_name, &raw_content);
Self {
channel: channel.into(),
chat_id: chat_id.into(),
content,
reply_to,
media: Vec::new(),
metadata,
event_kind: OutboundEventKind::ToolResult,
role: "tool".to_string(),
tool_call_id: Some(tool_call_id.into()),
tool_name: Some(tool_name),
tool_arguments: None,
}
}
pub fn tool_pending(
channel: impl Into<String>,
chat_id: impl Into<String>,
tool_call_id: impl Into<String>,
tool_name: impl Into<String>,
content: impl Into<String>,
reply_to: Option<String>,
metadata: HashMap<String, String>,
) -> Self {
let tool_name = tool_name.into();
let raw_content = content.into();
let content = format_tool_result_content(&tool_name, &raw_content);
Self {
channel: channel.into(),
chat_id: chat_id.into(),
content,
reply_to,
media: Vec::new(),
metadata,
event_kind: OutboundEventKind::ToolPending,
role: "tool".to_string(),
tool_call_id: Some(tool_call_id.into()),
tool_name: Some(tool_name),
tool_arguments: None,
}
}
pub fn from_chat_message(
channel: &str,
chat_id: &str,
reply_to: Option<String>,
metadata: &HashMap<String, String>,
message: &ChatMessage,
) -> Vec<Self> {
match message.role.as_str() {
"assistant" => {
if let Some(tool_calls) = &message.tool_calls {
let mut outbound = Vec::new();
if !message.content.trim().is_empty() {
outbound.push(Self::assistant(
channel.to_string(),
chat_id.to_string(),
message.content.clone(),
reply_to.clone(),
metadata.clone(),
));
}
outbound.extend(tool_calls
.iter()
.map(|tool_call| {
Self::tool_call(
channel.to_string(),
chat_id.to_string(),
tool_call.id.clone(),
tool_call.name.clone(),
tool_call.arguments.clone(),
reply_to.clone(),
metadata.clone(),
)
})
);
outbound
} else {
vec![Self::assistant(
channel.to_string(),
chat_id.to_string(),
message.content.clone(),
reply_to,
metadata.clone(),
)]
}
}
"tool" => match message.tool_state.as_ref().unwrap_or(&ToolMessageState::Completed) {
ToolMessageState::Completed => vec![Self::tool_result(
channel.to_string(),
chat_id.to_string(),
message.tool_call_id.clone().unwrap_or_default(),
message.tool_name.clone().unwrap_or_default(),
message.content.clone(),
reply_to,
metadata.clone(),
)],
ToolMessageState::PendingUserAction => vec![Self::tool_pending(
channel.to_string(),
chat_id.to_string(),
message.tool_call_id.clone().unwrap_or_default(),
message.tool_name.clone().unwrap_or_default(),
message.content.clone(),
reply_to,
metadata.clone(),
)],
},
_ => Vec::new(),
}
}
}
pub(crate) fn format_tool_call_content(tool_name: &str, tool_arguments: &serde_json::Value) -> String {
let mut lines = vec![format!("### {}", tool_name)];
match tool_arguments {
serde_json::Value::Object(map) if !map.is_empty() => {
let mut entries: Vec<_> = map.iter().collect();
entries.sort_by(|(left, _), (right, _)| left.cmp(right));
for (key, value) in entries {
lines.push(format!("- {}: {}", key, format_tool_argument_value(value)));
}
}
serde_json::Value::Object(_) => {}
other => lines.push(format!("- args: {}", format_tool_argument_value(other))),
}
lines.join("\n")
}
fn format_tool_result_content(tool_name: &str, content: &str) -> String {
format!("工具结果: {}\n\n{}", tool_name, content)
}
fn format_tool_argument_value(value: &serde_json::Value) -> String {
match value {
serde_json::Value::String(text) => text.clone(),
serde_json::Value::Null => "null".to_string(),
other => serde_json::to_string(other).unwrap_or_else(|_| other.to_string()),
}
}
// ============================================================================
// Helpers
// ============================================================================
fn current_timestamp() -> i64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap()
.as_millis() as i64
}
#[cfg(test)]
mod tests {
use super::{ChatMessage, OutboundEventKind, OutboundMessage, ToolMessageState};
use crate::providers::ToolCall;
use serde_json::json;
use std::collections::HashMap;
#[test]
fn test_from_chat_message_expands_tool_calls() {
let message = ChatMessage::assistant_with_tool_calls(
"",
vec![
ToolCall {
id: "call-1".to_string(),
name: "calculator".to_string(),
arguments: json!({"expression": "1 + 1"}),
},
ToolCall {
id: "call-2".to_string(),
name: "file_read".to_string(),
arguments: json!({"path": "README.md"}),
},
],
);
let outbound = OutboundMessage::from_chat_message(
"feishu",
"chat-1",
None,
&HashMap::new(),
&message,
);
assert_eq!(outbound.len(), 2);
assert_eq!(outbound[0].event_kind, OutboundEventKind::ToolCall);
assert_eq!(outbound[0].tool_name.as_deref(), Some("calculator"));
assert_eq!(outbound[0].tool_arguments.as_ref().unwrap()["expression"], "1 + 1");
assert_eq!(outbound[0].content, "### calculator\n- expression: 1 + 1");
assert_eq!(outbound[1].tool_name.as_deref(), Some("file_read"));
assert_eq!(outbound[1].content, "### file_read\n- path: README.md");
}
#[test]
fn test_from_chat_message_keeps_assistant_content_when_tool_calls_exist() {
let message = ChatMessage::assistant_with_tool_calls(
"日报已整理完成。",
vec![ToolCall {
id: "call-1".to_string(),
name: "memory_manage".to_string(),
arguments: json!({"action": "put"}),
}],
);
let outbound = OutboundMessage::from_chat_message(
"feishu",
"chat-1",
None,
&HashMap::new(),
&message,
);
assert_eq!(outbound.len(), 2);
assert_eq!(outbound[0].event_kind, OutboundEventKind::AssistantResponse);
assert_eq!(outbound[0].content, "日报已整理完成。");
assert_eq!(outbound[1].event_kind, OutboundEventKind::ToolCall);
assert_eq!(outbound[1].tool_name.as_deref(), Some("memory_manage"));
}
#[test]
fn test_from_chat_message_includes_tool_result() {
let message = ChatMessage::tool("call-9", "calculator", "2");
let outbound = OutboundMessage::from_chat_message(
"feishu",
"chat-1",
None,
&HashMap::new(),
&message,
);
assert_eq!(outbound.len(), 1);
assert_eq!(outbound[0].event_kind, OutboundEventKind::ToolResult);
}
#[test]
fn test_from_chat_message_includes_tool_pending() {
let message = ChatMessage::tool_with_state(
"call-9",
"bash",
"等待你完成浏览器授权后再继续。",
ToolMessageState::PendingUserAction,
);
let outbound = OutboundMessage::from_chat_message(
"feishu",
"chat-1",
None,
&HashMap::new(),
&message,
);
assert_eq!(outbound.len(), 1);
assert_eq!(outbound[0].event_kind, OutboundEventKind::ToolPending);
}
}
View Git Blame Copy Permalink